#!/usr/bin/env python3
"""Final re-scoring with fence-tolerant JSON parsing + correct package keys."""
import json, os, re
from collections import defaultdict

BENCH = os.path.expanduser("~/vision_bench")
lines = [json.loads(l) for l in open(f"{BENCH}/bench_results.jsonl")]
recs = [r for r in lines if not r.get("meta")]

IDENT_KEYS = {
    "android": ["dialog", "quickstep", "launcher", "home"],
    "com.android.chrome": ["chrome", "browser", "sign", "welcome", "terms", "wikipedia", "article"],
    "com.android.launcher3": ["home", "launcher", "wallpaper", "icon"],
    "com.android.settings": ["settings", "setup"],
    "com.android.settings.intelligence": ["settings", "search", "wi-fi", "wifi"],
    "com.android.calculator2": ["calculator", "calc"],
}

def gt_pkg(sit):
    xmls = {"s1_dialog":"ui1.xml","s3_chrome":"ui3.xml"}
    p = f"{BENCH}/{xmls.get(sit, sit+'.xml')}"
    try:
        import xml.etree.ElementTree as ET
        return ET.parse(p).getroot().iter("node").__next__().get("package")
    except Exception:
        return None

def parse_json_fuzzy(text):
    t = text.strip()
    m = re.search(r"```(?:json)?\s*(.*?)```", t, re.S)
    if m: t = m.group(1).strip()
    try: return json.loads(t)
    except Exception: pass
    m = re.search(r"\{.*\}", t, re.S)
    if m:
        try: return json.loads(m.group(0))
        except Exception: pass
    return None

def rescore():
    rows = defaultdict(lambda: defaultdict(list))
    for r in recs:
        sit, task, model = r["situation"], r["task"], r["model"]
        if r.get("error"): continue
        ans = r.get("content", "")
        if task == "identify":
            pkg = gt_pkg(sit)
            keys = IDENT_KEYS.get(pkg, [])
            sc = 1.0 if keys and any(k in ans.lower() for k in keys) else 0.0
            rows[model]["ident"].append(sc)
        elif task == "describe":
            rows[model]["desc"].append(r.get("score") or 0)
        elif task == "ground":
            rows[model]["ground"].append(r.get("score") or 0)
        elif task == "agent":
            j = parse_json_fuzzy(ans)
            sc = 0.0
            if isinstance(j, dict) and "elements" in j:
                sc = 1.0 if len(j["elements"]) > 0 else 0.5
                # bonus: coverage of first clickable label
            rows[model]["agent"].append(sc)
        rows[model]["wall"].append(r.get("wall_s") or 0)
        rows[model]["eval"].append(r.get("eval_s") or 0)
        rows[model]["tok"].append(r.get("tok_s") or 0)

    vram = {}
    for r in lines:
        if r.get("meta") and r.get("vram"):
            vram[r["model"]] = r["vram"].get("vram_gb")

    def a(xs): return round(sum(xs)/len(xs), 2) if xs else None

    print(f"{'model':24s} {'ident':>5s} {'desc':>5s} {'ground':>6s} {'agent':>5s} {'wall':>5s} {'tok/s':>6s} {'vramGB':>6s} {'fileGB':>6s}")
    sizes = {"moondream:1.8b":1.7, "granite3.2-vision:2b":1.7, "qwen2.5vl:3b":2.6, "llava-phi3:3.8b":2.9,
             "gemma3:4b":3.3, "llava:7b":4.7, "qwen2.5vl:7b":5.7, "minicpm-v4.5:latest":6.1}
    results = {}
    for m in sorted(rows, key=lambda x: sizes.get(x, 99)):
        d = rows[m]
        ident, desc, ground, agent = a(d["ident"]), a(d["desc"]), a(d["ground"]), a(d["agent"])
        wall, tok = a(d["wall"]), a(d["tok"])
        vr = vram.get(m)
        results[m] = dict(ident=ident, desc=desc, ground=ground, agent=agent, wall=wall, tok=tok, vram=vr)
        print(f"{m:24s} {ident:>5} {desc:>5} {ground:>6} {agent:>5} {wall:>5} {tok:>6} {str(vr):>6s} {sizes.get(m,0):>6}")

    print("\n=== GROUNDING on screens where target label was VISIBLE (dialog/chrome/settings/article/keyboard) ===")
    for m in sorted(rows, key=lambda x: sizes.get(x, 99)):
        vals = [r.get("score") for r in recs if r["model"]==m and r["task"]=="ground"
                and r["situation"] in ("s1_dialog","s3_chrome","s_settings","s_article","s_keyboard")
                and not r.get("error") and r.get("score") is not None]
        if vals:
            print(f"{m:24s} avg={round(sum(vals)/len(vals),2)}  " + ", ".join(f"{r['situation'][:7]}:{r['score']}" for r in recs
                if r["model"]==m and r["task"]=="ground" and r["situation"] in ("s1_dialog","s3_chrome","s_settings","s_article","s_keyboard") and not r.get("error")))
    return results

rescore()