#!/usr/bin/env python3
"""On-device eval for a minus model dir (AX650 NPU): ad + screen-state heads.

Reuses the model dir's own infer_vlm_fused.py helpers (patchify, prompt build, layer
order) so the run is exactly the production path, but keeps the full logits.

  python3 eval_v02.py ad     <test_dir with ads/ non_ads/>      <out.json>
  python3 eval_v02.py screen <manifest.csv path,label>  <img_root> <out.json>
"""
import csv, json, os, sys, time
import numpy as np

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import infer_vlm_fused as F

AD_Q = "Is this an advertisement? Answer Yes or No."
SC_Q = "Classify this TV screen: PLAYING, PAUSED, DIALOG, MENU, or SCREENSAVER?"
SCREEN = (("PLAYING", (9436, 60595)), ("PAUSED", (7055, 16137)), ("DIALOG", (20546, 23238)),
          ("MENU", (554, 857)), ("SCREENSAVER", (8421, 14151)))


def logits_for(models, path, prompt):
    from PIL import Image
    t0 = time.time()
    patches = F.patchify_image_512(Image.open(path).convert("RGB"))
    v = models["vision"].run(None, {"pixel_values": patches})[0].astype(np.float32, copy=False)
    ids = F.build_prompt(models["tokenizer"], prompt)
    n = min(len(ids), F.PREFILL_LEN)
    arr = np.array(ids[:n], dtype=np.int64)
    pos = np.where(arr == F.IMG_TOKEN)[0]
    d = np.zeros((1, F.PREFILL_LEN, F.HIDDEN_SIZE), np.float32)
    d[0, :n] = models["embeds"][arr].astype(np.float32, copy=False)
    d[0, pos[0]:pos[0] + min(len(pos), v.shape[1])] = v[0, :min(len(pos), v.shape[1])]
    mask = np.full((1, F.PREFILL_LEN, F.PREFILL_LEN), -65536.0, np.float32)
    mask[0, :n, :n] = np.triu(np.ones((n, n), np.float32), k=1) * -65536.0
    idx = np.arange(F.PREFILL_LEN, dtype=np.int32).reshape(1, F.PREFILL_LEN)
    for i, lt in enumerate(F.LAYER_TYPES):
        if lt == "conv":
            d = models["fused"][i].run(None, {"hidden": d, "conv_state_in": np.zeros((1, F.HIDDEN_SIZE, F.CONV_L_CACHE), np.float32)})[0]
        else:
            d = models["fused"][i].run(None, {"hidden": d, "mask": mask, "indices": idx})[0]
    lg = models["decode_post"].run(None, {"input": d[:, n - 1:n]})[0].flatten()
    return lg, (time.time() - t0) * 1000, n


def main():
    mode = sys.argv[1]
    models = F.load_models()
    res, times = [], []
    if mode == "ad":
        root, out = sys.argv[2], sys.argv[3]
        items = [(os.path.join(root, d, f), 1 if d == "ads" else 0)
                 for d in ("ads", "non_ads") for f in sorted(os.listdir(os.path.join(root, d)))
                 if f.lower().endswith((".png", ".jpg", ".jpeg"))]
        for k, (p, y) in enumerate(items):
            lg, ms, n = logits_for(models, p, AD_Q)
            yl, nl = float(lg[list(F.YES_IDS)].max()), float(lg[list(F.NO_IDS)].max())
            z = np.array([lg[t] for t in list(F.YES_IDS) + list(F.NO_IDS)]); z = np.exp(z - z.max())
            res.append({"path": p, "label": y, "pred": int(yl > nl), "p_yes": float(z[:len(F.YES_IDS)].sum() / z.sum()), "ms": ms})
            times.append(ms)
            if k % 50 == 0: print(f"{k}/{len(items)} n_tokens={n} {ms:.0f}ms", flush=True)
        acc = np.mean([r["pred"] == r["label"] for r in res])
        adr = np.mean([r["pred"] for r in res if r["label"] == 1]); nar = np.mean([1 - r["pred"] for r in res if r["label"] == 0])
        summ = {"accuracy": acc, "ad_recall": adr, "nonad_recall": nar, "n": len(res),
                "fp": sum(r["pred"] and not r["label"] for r in res), "fn": sum(r["label"] and not r["pred"] for r in res)}
    else:
        man, root, out = sys.argv[2], sys.argv[3], sys.argv[4]
        items = [(os.path.join(root, os.path.basename(r["path"])), r["label"]) for r in csv.DictReader(open(man))]
        for k, (p, y) in enumerate(items):
            lg, ms, n = logits_for(models, p, SC_Q)
            pred = max(SCREEN, key=lambda c: max(lg[t] for t in c[1]))[0]
            res.append({"path": p, "label": y, "pred": pred, "ms": ms}); times.append(ms)
            if k % 50 == 0: print(f"{k}/{len(items)} n_tokens={n} {ms:.0f}ms", flush=True)
        conf = {}
        for r in res: conf.setdefault(r["label"], {}).setdefault(r["pred"], 0); conf[r["label"]][r["pred"]] += 1
        summ = {"accuracy": np.mean([r["pred"] == r["label"] for r in res]), "n": len(res), "confusion": conf,
                "playing_misread_as_action": sum(r["label"] == "PLAYING" and r["pred"] in ("PAUSED", "SCREENSAVER", "DIALOG", "MENU") for r in res)}
    summ.update({"avg_ms": float(np.mean(times)), "median_ms": float(np.median(times)), "max_ms": float(np.max(times))})
    json.dump({"summary": summ, "results": res}, open(out, "w"), indent=1, default=float)
    print(json.dumps(summ, indent=1, default=float))


if __name__ == "__main__":
    main()