#!/usr/bin/env python3
"""decision_eval.py — Gemini Decision 1 が mechArm のお手本をどこまで正しく判定できるかを測る (読むだけ).

正解: lab/Log25・Log26 の runs_table.csv (人の判定 pick/place・外した理由) と grasp_depth.py の結果 (浅い=閉じて1秒後 24px 超・手の疑い).
入力: manifest.csv (prep.py が作る. 1本につき上カメラのコマ3枚: 閉じて1秒後 after / 閉じて3秒後 lift / 最後のコマ end)
鍵:   /Volumes/KIOXIA/gemini_api_key.txt
  python3 decision_eval.py run  [--limit N] [--resolution low|medium|high]   # API を呼んで results.csv に書く
  python3 decision_eval.py score                                             # results.csv と正解を比べて score.md に書く
"""
import argparse, base64, csv, json, os, sys, time, urllib.request, statistics as st

URL = "https://generativelanguage.googleapis.com/v1beta/decisions"
KEY = "/Volumes/KIOXIA/gemini_api_key.txt"
CTX = ("Top-down camera image of a small desktop robot arm (mechArm 270) with a two-finger gripper "
       "picking a red cube from a mat and placing it into a white box.\n")
Q = {
 "after": {
  "depth": {"type": "score", "instructions": CTX + "The gripper has just closed on the red cube. How deep is the cube inside the gripper fingers?",
            "criteria": ["Deep: the cube sits well inside the fingers, not sticking out past the finger tips.",
                         "Partly: the cube sticks out slightly past the finger tips.",
                         "Shallow: the cube is held only at the very tips of the fingers or sticks out a lot."]},
  "claw_on_block": {"type": "noul", "instructions": CTX + "Is a gripper finger resting on top of the red cube instead of closing beside it?"},
  "hand": {"type": "noul", "instructions": CTX + "Is a human hand or finger visible in the image?"},
 },
 "lift": {"picked": {"type": "noul", "instructions": CTX + "Is the red cube held between the gripper fingers and lifted off the mat?"}},
 "end":  {"placed": {"type": "noul", "instructions": CTX + "Is the red cube inside the white box?"}},
}

def call(key, img, questions, res):
    body = {"model": "gemini-decision-1",
            "input": [{"type": "image", "data": base64.b64encode(open(img, "rb").read()).decode(),
                       "mime_type": "image/jpeg", "resolution": res}],
            "questions": questions}
    req = urllib.request.Request(URL, json.dumps(body).encode(), {"Content-Type": "application/json", "x-goog-api-key": key})
    t = time.time()
    with urllib.request.urlopen(req, timeout=60) as r: out = json.load(r)
    return out, time.time() - t

def run(a):
    key = open(KEY).read().strip()
    rows = list(csv.DictReader(open("manifest.csv")))[: a.limit or None]
    w = csv.writer(open("results.csv", "w")); w.writerow(["run", "image", "question", "value", "probabilities", "confidence", "latency_s", "input_tokens", "resolution"])
    for r in rows:
        for kind, qs in Q.items():
            img = r[f"img_{kind}"]
            if not img: continue
            try: out, lat = call(key, img, qs, a.resolution)
            except Exception as e:
                print("ERR", r["run"], kind, e, file=sys.stderr); continue
            tok = out.get("usage", {}).get("total_input_tokens", "")
            for q, ans in out["answers"].items():
                v = ans.get("noul", ans.get("score", ans.get("choice")))
                w.writerow([r["run"], kind, q, v, json.dumps(ans.get("probabilities", "")), ans.get("confidence"), round(lat, 3), tok, a.resolution])
            print(r["run"], kind, "ok", round(lat, 2))
            time.sleep(0.4)

def score(a):
    man = {r["run"]: r for r in csv.DictReader(open("manifest.csv"))}
    res = {}
    for r in csv.DictReader(open("results.csv")): res[(r["run"], r["question"])] = r
    truth = {  # 正解 (True = 問いに「はい」)
      "depth": lambda m: m["shallow"] == "True",            # 浅い (24px 超)
      "claw_on_block": lambda m: m["claw_on_block"] == "True",
      "hand": lambda m: m["hand_suspect"] == "True",
      "picked": lambda m: m["human_pick"] == "yes",
      "placed": lambda m: m["human_place"] == "yes"}
    L = ["# Gemini Decision 1 の判定と正解の比べ\n", "| 問い | 本数 | 正解が「はい」 | 一致 | 見逃し | 誤り | 自信0.8以上の一致 | 自信0.8未満の本数 |", "|---|---|---|---|---|---|---|---|"]
    for q, f in truth.items():
        n = tp = fn = fp = hi_n = hi_ok = lo = 0; pos = 0
        for run, m in man.items():
            r = res.get((run, q))
            if not r: continue
            t = f(m); v = float(r["value"]); c = float(r["confidence"] or 0)
            p = (v >= 1.0) if q == "depth" else (v >= 0.5)   # depth は 0〜2 の段階の平均. 1 以上を「浅い」とみなす
            n += 1; pos += t; ok = (p == t); tp += ok and t; fn += (t and not p); fp += (p and not t)
            if c >= 0.8: hi_n += 1; hi_ok += ok
            else: lo += 1
            agree = sum(1 for _ in [0])
        if n: L.append(f"| {q} | {n} | {pos} | {sum(1 for run,m in man.items() if (run,q) in res and ((float(res[(run,q)]['value'])>=(1.0 if q=='depth' else 0.5))==f(m)))} | {fn} | {fp} | {hi_ok}/{hi_n} | {lo} |")
    lats = [float(r["latency_s"]) for r in res.values()]; toks = [int(r["input_tokens"]) for r in res.values() if r["input_tokens"]]
    if lats: L.append(f"\n1回の依頼の時間 (通信込み): 中央値 {st.median(lats):.2f} 秒 / 最大 {max(lats):.2f} 秒")
    if toks: L.append(f"入力トークン: 1回あたり中央値 {int(st.median(toks))} / 合計 {sum(toks)} (費用の目安 ${sum(toks)/1e6*0.10:.4f})")
    open("score.md", "w").write("\n".join(L) + "\n"); print("\n".join(L))

if __name__ == "__main__":
    ap = argparse.ArgumentParser(); ap.add_argument("cmd", choices=["run", "score"]); ap.add_argument("--limit", type=int); ap.add_argument("--resolution", default="medium")
    a = ap.parse_args(); run(a) if a.cmd == "run" else score(a)