"""Average v3 and v4 weights, refit temperatures on synthetic + spatial + MASSIVE validation."""
import json, shutil, sys, os
import numpy as np, torch
sys.path.insert(0, "./train")
import train_v2 as T
sys.path.insert(0, ".")  # repo root, so `import thx01` works
import thx01
from safetensors.torch import load_file, save_file

alpha = float(sys.argv[1]); OUT = sys.argv[2]
A_PATH = sys.argv[3] if len(sys.argv) > 3 else "."
B_PATH = sys.argv[4] if len(sys.argv) > 4 else "./ckpt/system1_v4"
a = load_file(f"{A_PATH}/model.safetensors"); b = load_file(f"{B_PATH}/model.safetensors")
os.makedirs(OUT, exist_ok=True)
save_file({k: ((1 - alpha) * a[k].float() + alpha * b[k].float()).to(torch.float16).contiguous() for k in a}, f"{OUT}/model.safetensors")
for sub in ("encoder", "tokenizer"):
    shutil.copytree(f"{A_PATH}/{sub}", f"{OUT}/{sub}", dirs_exist_ok=True)
cfg = json.load(open(f"{B_PATH}/rl_agent_config.json"))
cfg["temperature_by_options"] = {}; cfg["temperature"] = [1.0, 1.0, 1.0]
json.dump(cfg, open(f"{OUT}/rl_agent_config.json", "w"), indent=2)
ag = thx01.load(OUT, device="cuda")
val = T.synth_items("val")
for l in open(f"{T.SYN}/spatial_val.jsonl"):
    r = json.loads(l)
    for qid, g in r["gold"].items():
        val.append((r["state"], {qid: r["questions"][qid]}, {qid: g}))
for lg in ("az", "en"):
    val += [T.reshuffle(*x) for x in T.tasks.massive_scenario(lg, "validation", 400, seed=77)]
    val += [T.reshuffle(*x, min_k=5, max_k=24) for x in T.tasks.massive_intent(lg, "validation", 400, k=60, seed=78)]
    val += list(T.tasks.intent_noul(lg, "validation", 300, seed=79))
vit = [e for e in (T.encode(ag, *x) for x in val) if e]
buckets = {}
with torch.no_grad():
    for bi in T.batches(vit):
        chunk = [vit[i] for i in bi]
        bt = T.collate_items([chunk], ag.tok.pad_token_id)
        lg_, _ = T.forward(ag.model, bt); lg_ = lg_.float().cpu()
        for j, it in enumerate(chunk):
            k = len(it["markers"]); buckets.setdefault(T.temp_bucket(it["qtype"], k), []).append((lg_[j, :k], it["label"]))
temps = {}
for bk, xs in buckets.items():
    if len(xs) < 40:
        continue
    L = torch.nn.utils.rnn.pad_sequence([x for x, _ in xs], batch_first=True, padding_value=-1e4); y = torch.tensor([t for _, t in xs])
    temps[bk] = round(float(min(np.arange(0.5, 5.01, 0.05), key=lambda t: torch.nn.functional.cross_entropy(L / t, y).item())), 3)
cfg["temperature_by_options"] = temps
for qt, name in ((0, "choice"), (1, "score"), (2, "noul")):
    v = [x for k, x in temps.items() if k.startswith(name)]
    if v:
        cfg["temperature"][qt] = round(float(np.median(v)), 3)
json.dump(cfg, open(f"{OUT}/rl_agent_config.json", "w"), indent=2)
print(alpha, temps)