#!/usr/bin/env python3
"""Compare the recorded CPU results against the same scenarios run fresh on the Space's GPU.

    python scripts/compare_devices.py [--space wang2226/steering-showcase] [--model qwen25-05b]

The recording and the Space share one implementation, but shared source is not evidence of identical
output: they run on different hardware with different kernels. This measures what "the same
procedure" is worth and writes the result to data/device-check.json.
A fresh run of every scenario costs GPU time; logged-in calls use the account's ZeroGPU quota.
"""

from __future__ import annotations

import argparse
import datetime as dt
import json
import os
import sys
from pathlib import Path

ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT))
os.environ["SD_SKIP_INIT"] = "1"

import recorded  # noqa: E402


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--space", default="wang2226/steering-showcase")
    ap.add_argument("--model", default="qwen25-05b")
    ap.add_argument("--out", default=str(recorded.DEVICE_CHECK))
    args = ap.parse_args()

    from gradio_client import Client
    from huggingface_hub import get_token

    token = os.environ.get("HF_TOKEN")
    if not token and (ROOT / ".env.local").exists():
        token = next((l.split("=", 1)[1].strip() for l in (ROOT / ".env.local").read_text().splitlines()
                      if l.startswith("HF_TOKEN=")), None)
    client = Client(args.space, token=token or get_token())
    rec = recorded.record(args.model)
    if not rec:
        print(f"no recorded results for {args.model}")
        return 1

    states = same_token = same_text = 0
    worst = 0.0
    gpu = None
    for cpu in rec["scenarios"]:
        live = client.predict(args.model, cpu["id"], "", "", True, api_name="/trace")
        gpu = (live.get("env") or {}).get("device") or gpu
        if live["candidate_ids"] != cpu["candidate_ids"]:
            print(f"{cpu['id']}: candidate rows differ: {cpu['candidates']} vs {live['candidates']}")
        # Compared by token id, so rows in a different order or a different set still compare like with like.
        cpu_ids, live_ids = cpu["candidate_ids"], live["candidate_ids"]
        for a, b in zip(cpu["states"], live["states"]):
            states += 1
            same_token += cpu_ids[a["selected_index"]] == live_ids[b["selected_index"]]
            same_text += a["continuation"] == b["continuation"]
            live_pct = dict(zip(live_ids, b["percents"]))
            pairs = [(x, live_pct[i]) for i, x in zip(cpu_ids, a["percents"]) if i in live_pct]
            if set(cpu_ids) == set(live_ids):
                pairs.append((a["other"], b["other"]))
            for x, y in pairs:
                if x > 0:
                    worst = max(worst, abs(x - y) / x)
        print(f"{cpu['id']:18} {live['env'].get('device')}: done")

    result = {"model": rec["model"], "gpu": gpu, "states": states, "same_token": same_token, "same_text": same_text,
              "worst_relative": worst, "recorded_on": rec["runtime"]["device"], "checked": dt.date.today().isoformat()}
    Path(args.out).write_text(json.dumps(result, indent=1) + "\n", encoding="utf-8")
    print(json.dumps(result, indent=1))
    return 0 if same_token == states else 1


if __name__ == "__main__":
    raise SystemExit(main())