#!/usr/bin/env python3
"""Average two of our own checkpoints (same architecture)."""
from __future__ import annotations

import json
import shutil
import sys
from pathlib import Path

import torch
from safetensors.torch import load_file, save_file


def main():
    a, b, out, wa = Path(sys.argv[1]), Path(sys.argv[2]), Path(sys.argv[3]), float(sys.argv[4])
    wb = 1.0 - wa
    out.mkdir(parents=True, exist_ok=True)
    shards = sorted(a.glob("model-*.safetensors"))
    if not shards:
        raise SystemExit(f"no shards in {a}")
    index = {"metadata": {"format": "pt"}, "weight_map": {}}
    for sa in shards:
        sb = b / sa.name
        ta = load_file(str(sa), device="cpu")
        tb = load_file(str(sb), device="cpu")
        merged = {}
        for k in ta:
            if k not in tb:
                merged[k] = ta[k]
                continue
            if ta[k].dtype in (torch.float16, torch.bfloat16, torch.float32):
                merged[k] = (ta[k].float() * wa + tb[k].float() * wb).to(ta[k].dtype)
            else:
                merged[k] = ta[k]
        dest = out / sa.name
        save_file(merged, str(dest))
        for k in merged:
            index["weight_map"][k] = sa.name
        print("wrote", dest.name, "tensors", len(merged))
    (out / "model.safetensors.index.json").write_text(json.dumps(index, indent=2))
    for name in (
        "config.json",
        "generation_config.json",
        "tokenizer.json",
        "tokenizer_config.json",
        "chat_template.jinja",
        "special_tokens_map.json",
    ):
        src = a / name
        if src.exists():
            shutil.copy2(src, out / name)
    (out / "STABLE").write_text("ok\n")
    print("done", out, "w", wa, wb)


if __name__ == "__main__":
    main()