"""Interpolate a fine-tuned Dory checkpoint with the base model (WiSE-FT style), no training.

  .venv/bin/python -I scripts/merge_dory.py --base models/dory-4b --tuned models/rA/final --t 0.5 --out models/rA-t050

w = (1 - t) * base + t * tuned for every tensor (in fp32), then matrices are projected back onto the
unit hypersphere along the hidden dimension, as DoryForCausalLM.normalize_matrices_ does after each
optimizer step. For unit rows this is a normalized lerp, i.e. close to slerp. Scalars (alphas, scales,
logit_scale) are interpolated linearly.
"""

import argparse
import json
import shutil
import sys
from pathlib import Path

import torch
from safetensors import safe_open
from safetensors.torch import save_file
from torch import nn

sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
from dory.modeling import DoryConfig, DoryForCausalLM, justnorm  # noqa: E402

BASE_FILES = ("config.json", "generation_config.json", "tokenizer.json", "tokenizer_config.json",
              "special_tokens_map.json", "chat_template.jinja")


def norm_dims(cfg):
    """Parameter name -> dim to normalize along, mirroring normalize_matrices_."""
    with torch.device("meta"):
        model = DoryForCausalLM(cfg)
    dims = {"backbone.embeddings.weight": 1, "lm_head.weight": 1}
    for name, mod in model.backbone.named_modules():
        if isinstance(mod, nn.Linear):
            dims[f"backbone.{name}.weight"] = 1 if mod.in_features == cfg.hidden_size else 0
    return dims


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--base", required=True)
    ap.add_argument("--tuned", required=True)
    ap.add_argument("--t", type=float, required=True, help="0 = base, 1 = tuned")
    ap.add_argument("--out", required=True)
    args = ap.parse_args()
    base, tuned, out = Path(args.base), Path(args.tuned), Path(args.out)

    cfg = DoryConfig.from_json(tuned / "config.json")
    dims = norm_dims(cfg)
    merged, n_norm = {}, 0
    with safe_open(base / "model.safetensors", "pt") as fb, safe_open(tuned / "model.safetensors", "pt") as ft:
        assert set(fb.keys()) == set(ft.keys()), "checkpoints have different tensors"
        assert set(dims) <= set(fb.keys()), sorted(set(dims) - set(fb.keys()))[:5]
        for k in fb.keys():
            a, b = fb.get_tensor(k).float(), ft.get_tensor(k).float()
            w = torch.lerp(a, b, args.t)
            if k in dims:
                w = justnorm(w, dim=dims[k])
                n_norm += 1
            merged[k] = w.to(torch.bfloat16).contiguous()

    out.mkdir(parents=True, exist_ok=True)
    save_file(merged, out / "model.safetensors", metadata={"format": "pt"})
    for f in BASE_FILES:
        if (tuned / f).exists():
            shutil.copy(tuned / f, out / f)
    (out / "merge_args.json").write_text(json.dumps({**vars(args), "normalized_matrices": n_norm}, indent=1))
    print(f"merged {len(merged)} tensors ({n_norm} renormalized), t={args.t} -> {out}")


if __name__ == "__main__":
    main()