"""Average checkpoints into a soup (the lever that carried this release).

Same-basin checkpoints from the same recipe, saved at different points of
short continuation runs, average into a better model than any member: on this
line the final soups scored +0.09 SCOREQ above their best single member, for
zero training.

    # uniform average
    python average_checkpoints.py --out decoder_soup.pt a.pt b.pt c.pt d.pt

    # weighted average (the acoustic blend used weights 2,2,1,1,1)
    python average_checkpoints.py --out acoustic_G4.pt \
        --weights 2,2,1,1,1 bs015.pt ft2.pt t1.pt t2.pt t3.pt

Rules (measured on this line):
  - Only average checkpoints with IDENTICAL key sets; the script refuses to mix
    different architectures.
  - Float tensors are averaged; non-float tensors are copied from the first
    member.
  - Members must share a basin (same recipe, diverged only by training noise
    and step count). Mixing distant runs loses.
  - 2 to 6 members is the proven range; more members dilute. The shipped
    decoder is the average of its two best soups, which itself beat larger
    blends, and the shipped acoustic is a pair-heavy blend of five
    checkpoints.
  - Averaging can be applied to soups again (nested), e.g. the shipped decoder
    is avg(e5, ss2) where e5 and ss2 are themselves averages.
"""
import argparse

import torch


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--out", required=True)
    ap.add_argument("--weights", default=None,
                    help="comma-separated per-input weights (default: uniform)")
    ap.add_argument("inputs", nargs="+", help="checkpoint files to average")
    args = ap.parse_args()

    weights = None
    if args.weights:
        weights = [float(x) for x in args.weights.split(",")]
        if len(weights) != len(args.inputs):
            raise SystemExit(f"{len(weights)} weights for {len(args.inputs)} inputs")
    else:
        weights = [1.0] * len(args.inputs)
    tot = sum(weights)

    sds = [torch.load(p, map_location="cpu", weights_only=False) for p in args.inputs]
    ks0 = set(sds[0].keys())
    for p, sd in zip(args.inputs[1:], sds[1:]):
        if set(sd.keys()) != ks0:
            raise SystemExit(f"key mismatch: {p} does not match {args.inputs[0]}")

    out = {}
    for k in sds[0].keys():
        ref = sds[0][k]
        if torch.is_tensor(ref) and ref.is_floating_point():
            acc = None
            for w, sd in zip(weights, sds):
                v = sd[k].float() * (w / tot)
                acc = v if acc is None else acc + v
            out[k] = acc
        else:
            out[k] = ref
    torch.save(out, args.out)
    n = sum(v.numel() for v in out.values() if torch.is_tensor(v))
    print(f"averaged {len(sds)} checkpoints -> {args.out} ({n:,} parameters)")


if __name__ == "__main__":
    main()