"""Convert a sharded training checkpoint into one plain model file for eval.

Training saves model + optimizer split across all GPUs. Eval only needs the
model weights, in a single file that one GPU can load.
"""
import argparse

import torch
import torch.distributed.checkpoint as dcp

from ar.model.config import ARConfig
from ar.model.model import ARModel


def main(ckpt_dir, out_path):
    # build on meta then allocate on CPU: skips a slow 4B-param random init
    with torch.device("meta"):
        model = ARModel(ARConfig())
    model.to_empty(device="cpu")
    sd = model.state_dict()

    # see which weights the checkpoint actually holds (its keys look like "model.<name>")
    meta = dcp.FileSystemReader(ckpt_dir).read_metadata()
    saved = {k[len("model."):] for k in meta.state_dict_metadata if k.startswith("model.")}
    wanted = {k: v for k, v in sd.items() if k in saved}
    missing = [k for k in sd if k not in saved]
    print(f"loading {len(wanted)}/{len(sd)} tensors; not in checkpoint: {missing}")

    # load ONLY the model part, filling the tensors in place (optimizer is skipped)
    dcp.load(state_dict={"model": wanted}, checkpoint_id=ckpt_dir)

    # bf16 halves the file (~8 GB); eval runs in bf16 anyway
    out = {k: v.to(torch.bfloat16) for k, v in model.state_dict().items()}
    torch.save(out, out_path)
    print(f"saved {len(out)} tensors -> {out_path}")


if __name__ == "__main__":
    p = argparse.ArgumentParser()
    p.add_argument("--ckpt", required=True)
    p.add_argument("--out", required=True)
    args = p.parse_args()
    main(args.ckpt, args.out)