"""Pomiar Strands Decider 2B (v1-2610) na MASSIVE 1.1 (test, pl-PL i en-US). Wynik: JSONL, wznawialny.

Użycie: venv/bin/python pomiar.py --ids ids-pilot.txt --out pilot-1.jsonl [--pytanie intent|scenario]
Schemat pytania intent identyczny jak w clef-pl-2026-10/pomiar.py.
"""
import argparse
import random
import json
import platform
import time
from pathlib import Path

import peft
import torch
import transformers
from huggingface_hub import snapshot_download
from strands_decider.infer import load_engine

REPO = "StrandsAgents/strands-decider-2B-qwen3.5-v1-2610"
REWIZJA = "216d1ccb46a9e9b036eb9187b248db98480f77d4"
DANE = Path(__file__).parent / "1.1" / "data"
INSTRUKCJE = {
    "intent": "Which intent does the user's request express?",
    "scenario": "Which area does the user's request concern?",
}


def wczytaj(lokalizacja):
    return {r["id"]: r for r in map(json.loads, open(DANE / f"{lokalizacja}.jsonl")) if r["partition"] == "test"}


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--ids", required=True)
    parser.add_argument("--out", required=True)
    parser.add_argument("--pytanie", default="intent", choices=["intent", "scenario"])
    parser.add_argument("--opcje", type=int, default=0, help="0 = wszystkie; N = poprawna + N-1 losowych dystraktorów na parę")
    args = parser.parse_args()

    path = snapshot_download(REPO, revision=REWIZJA)
    cfg = json.load(open(Path(path) / "strands_decider_config.json")) if (Path(path) / "strands_decider_config.json").exists() else {}
    engine = load_engine(path, device="mps")
    wszystkie = [json.loads(l) for l in open(DANE / "pl-PL.jsonl")]
    if args.pytanie == "intent":
        etykiety = sorted({r["intent"] for r in wszystkie})
        pole = "intent"
    else:
        etykiety = sorted({r["scenario"] for r in wszystkie})
        pole = "scenario"
    pytanie = {
        "type": "choice",
        "instructions": INSTRUKCJE[args.pytanie],
        "criteria": {e: e.replace("_", " ") for e in etykiety},
    }
    dane = {"pl-PL": wczytaj("pl-PL"), "en-US": wczytaj("en-US")}
    meta = {
        "repo": REPO, "rewizja": REWIZJA, "base": cfg.get("base_model") or cfg.get("base_model_name"),
        "base_revision": cfg.get("base_revision"), "torch": torch.__version__,
        "transformers": transformers.__version__, "peft": peft.__version__, "python": platform.python_version(),
        "pytanie": args.pytanie, "liczba_opcji": len(etykiety),
    }
    print(json.dumps(meta, ensure_ascii=False), flush=True)

    ids = [l.strip() for l in open(args.ids) if l.strip()]
    out = Path(args.out)
    zrobione = {(r["id"], r["locale"]) for r in map(json.loads, open(out))} if out.exists() else set()
    kolejka = [(rid, loc) for rid in ids for loc in ("pl-PL", "en-US") if (rid, loc) not in zrobione]

    with open(out, "a") as plik:
        for rid, loc in kolejka:
            r = dane[loc][rid]
            pyt = pytanie
            if args.opcje:
                los = random.Random(f"20261002-{rid}")
                inne = [e for e in etykiety if e != r[pole]]
                wybrane = sorted([r[pole]] + los.sample(inne, args.opcje - 1))
                pyt = {**pytanie, "criteria": {e: e.replace("_", " ") for e in wybrane}}
            t0 = time.time()
            try:
                odp = engine.ask(r["utt"], {args.pytanie: pyt})
                czas = time.time() - t0
                a = odp.answers[args.pytanie]
                probs = dict(a.probabilities)
                pred = max(probs, key=probs.get)
                rekord = {
                    "id": rid, "locale": loc, "utt": r["utt"], "gold": r[pole], "pred": pred,
                    "conf": probs[pred], "confidence_decider": a.confidence, "choice": a.choice,
                    "probs": probs, "input_tokens": odp.usage.input_tokens, "czas_s": round(czas, 4),
                }
            except Exception as e:  # rekord do kategorii „nie zmierzono”
                rekord = {"id": rid, "locale": loc, "blad": repr(e)}
            plik.write(json.dumps(rekord, ensure_ascii=False) + "\n")
            plik.flush()


if __name__ == "__main__":
    main()