"""Blend a specialist's decision head into a champion checkpoint (same encoder), v32b-style.

    python finetune/blend_heads.py <champion_dir> <specialist_dir> <w> <out_dir>

head = (1 - w) * champion_head + w * specialist_head for every non-encoder tensor; encoder, tokenizer and config come from
the champion.  Meant for specialists trained with HEAD_ONLY=1 from the champion (their encoder is the champion's), so
the blend moves only the head toward the new skill -- small w keeps the champion's other skills intact.
"""
import json, os, shutil, sys
import torch
from safetensors.torch import load_file, save_file

champ, spec, w, out = sys.argv[1], sys.argv[2], float(sys.argv[3]), sys.argv[4]
a, b = load_file(os.path.join(champ, "model.safetensors")), load_file(os.path.join(spec, "model.safetensors"))
assert a.keys() == b.keys(), "different architectures"
enc_diff = max(((a[k].float() - b[k].float()).abs().max().item() for k in a if k.startswith("encoder.")), default=0.0)
if enc_diff > 1e-2:
    print(f"warning: encoders differ (max |d| = {enc_diff:.3g}); the champion's encoder is kept")
blend = {k: (a[k] if (k.startswith("encoder.") or k == "temperature") else ((1 - w) * a[k].float() + w * b[k].float()).to(a[k].dtype))
         for k in a}
os.makedirs(out, exist_ok=True)
save_file({k: v.contiguous() for k, v in blend.items()}, os.path.join(out, "model.safetensors"))
for sub in ("encoder", "tokenizer"):
    shutil.copytree(os.path.join(champ, sub), os.path.join(out, sub), dirs_exist_ok=True)
cfg = json.load(open(os.path.join(champ, "rl_agent_config.json")))
cfg["blend"] = {"champion": champ, "specialist": spec, "w": w}
json.dump(cfg, open(os.path.join(out, "rl_agent_config.json"), "w"), indent=2)
print(f"blended head (w={w}) -> {out}; encoder max diff {enc_diff:.3g}")