"""Eksploracyjnie (poza protokołem): ile tokenów wejścia zajmuje schemat pytania, a ile samo zdanie.

Tylko tokenizer i kod kodowania z repozytorium modelu, bez uruchamiania modelu. Wynik: tokeny-schematu.json.
"""
import json
import statistics
import sys
from pathlib import Path

from huggingface_hub import snapshot_download
from transformers import AutoTokenizer

REWIZJA = "17f0b0ad64efb65d273590632833508766b2aae6"
TU = Path(__file__).parent
DANE = TU / "1.1" / "data"

path = snapshot_download("Cloudflare/clef-flash", revision=REWIZJA)
sys.path.insert(0, path)
from joint_schema_model import encode_record  # noqa: E402

tok = AutoTokenizer.from_pretrained(path)
ids = [l.strip() for l in open(TU / "ids-proba.txt") if l.strip()]
intencje = sorted({json.loads(l)["intent"] for l in open(DANE / "pl-PL.jsonl")})
scenariusze = sorted({i.split("_")[0] for i in intencje})
schematy = {
    "intencja_60_opcji": {"type": "choice", "instructions": "Which intent does the user's request express?",
                          "criteria": {i: i.replace("_", " ") for i in intencje}},
    "scenariusz_18_opcji": {"type": "choice", "instructions": "Which area does the user's request concern?",
                            "criteria": {s: s for s in scenariusze}},
    "tak_nie": {"type": "noul", "instructions": "Does the user want to set an alarm?"},
}


def dlugosc(stan, pytanie):
    return len(encode_record(tok, {"state": stan, "questions": {"q": pytanie}}).input_ids)


wynik = {"rewizja_modelu": REWIZJA, "pary": len(ids), "jezyki": {}}
for loc in ("pl-PL", "en-US"):
    rek = {r["id"]: r for r in map(json.loads, open(DANE / f"{loc}.jsonl")) if r["partition"] == "test"}
    zdania = [rek[i]["utt"] for i in ids]
    w = {
        "znaki_zdania_srednio": round(statistics.mean(len(z) for z in zdania), 2),
        "tokeny_zdania_srednio": round(statistics.mean(len(tok(z, add_special_tokens=False)["input_ids"]) for z in zdania), 2),
        "schematy": {},
    }
    for nazwa, pyt in schematy.items():
        w["schematy"][nazwa] = {
            "puste_zdanie": dlugosc("", pyt),
            "cale_zapytanie_srednio": round(statistics.mean(dlugosc(z, pyt) for z in zdania), 2),
        }
    wynik["jezyki"][loc] = w
json.dump(wynik, open(TU / "tokeny-schematu.json", "w"), ensure_ascii=False, indent=2)
print(json.dumps(wynik, ensure_ascii=False, indent=1))