"""Convert a nanochat checkpoint into a HF-compatible bundle for this repo.

Usage:
    python convert_checkpoint.py \
        --input /path/to/model_026430.pt \
        --meta /path/to/meta_026430.json \
        --output-dir .

Writes:
    - model.safetensors   (weights, renamed for CognicaPoEForCausalLM)
    - config.json         (merged from existing + meta)
    - generation_config.json

The CognicaPoEForCausalLM wrapper stores the GPT under `self.gpt`, so the state
dict keys gain a `gpt.` prefix relative to the nanochat checkpoint. This script
handles that mapping automatically and strips any DDP/compile prefixes.
"""

import argparse
import json
from pathlib import Path

import torch
from safetensors.torch import save_file


# nanochat buffers that are persistent=False are not in the checkpoint; nothing to strip.
# Any remaining meta-only keys (e.g. compile wrappers) should be cleaned here.
_PREFIXES_TO_STRIP = ("_orig_mod.", "module.", "model.")


def _normalize_key(key: str) -> str:
    for p in _PREFIXES_TO_STRIP:
        if key.startswith(p):
            key = key[len(p):]
    return key


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--input", required=True, help="Path to model_XXXXX.pt")
    parser.add_argument("--meta", required=True, help="Path to meta_XXXXX.json")
    parser.add_argument("--output-dir", default=".", help="Where to write safetensors + config")
    parser.add_argument("--alpha", type=float, default=None,
                        help="Override poe_alpha in config.json (else read from meta)")
    parser.add_argument("--keep-dtype", action="store_true",
                        help="Skip the bf16 cast and preserve the original tensor dtype")
    args = parser.parse_args()

    output_dir = Path(args.output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)

    # 1. Load the nanochat checkpoint.
    state_dict = torch.load(args.input, map_location="cpu", weights_only=True)
    if isinstance(state_dict, dict) and "model" in state_dict:
        state_dict = state_dict["model"]

    # 2. Rename keys: strip DDP/compile prefixes, then add `gpt.` prefix for
    #    the CognicaPoEForCausalLM wrapper.
    renamed = {}
    for k, v in state_dict.items():
        nk = _normalize_key(k)
        hf_key = f"gpt.{nk}"
        tensor = v
        if not args.keep_dtype:
            tensor = tensor.to(torch.bfloat16)
        renamed[hf_key] = tensor.contiguous()

    # 3. Load meta file.
    with open(args.meta) as f:
        meta = json.load(f)

    # 4. Merge into existing config.json.
    config_path = output_dir / "config.json"
    with open(config_path) as f:
        config = json.load(f)

    user_config = meta.get("user_config", {})
    config["poe_mode"] = user_config.get("poe_mode", config.get("poe_mode", "flat"))
    config["poe_every"] = user_config.get("poe_every", config.get("poe_every", 6))
    config["poe_alpha"] = args.alpha if args.alpha is not None else \
        user_config.get("poe_alpha", config.get("poe_alpha", 0.0))
    n_layer = meta.get("model_config", {}).get("n_layer", config["num_hidden_layers"])
    config["poe_head_count"] = max(1, n_layer // config["poe_every"])
    config["training_step"] = meta.get("step")
    config["training_val_bpb"] = meta.get("val_bpb")

    with open(config_path, "w") as f:
        json.dump(config, f, indent=2)

    # 5. Write safetensors.
    save_file(
        renamed,
        output_dir / "model.safetensors",
        metadata={
            "format": "pt",
            "nanochat_step": str(meta.get("step", "?")),
            "val_bpb": f"{meta.get('val_bpb', 'NaN')}",
        },
    )

    # 6. Minimal generation_config.json.
    gen_config = {
        "transformers_version": config.get("transformers_version", "4.45.0"),
        "max_length": config["max_position_embeddings"],
        "do_sample": True,
        "temperature": 0.9,
        "top_p": 0.95,
    }
    with open(output_dir / "generation_config.json", "w") as f:
        json.dump(gen_config, f, indent=2)

    print(f"[ok] Wrote {len(renamed)} tensors -> {output_dir / 'model.safetensors'}")
    print(f"[ok] poe_alpha = {config['poe_alpha']}, step = {config.get('training_step')}, "
          f"val_bpb = {config.get('training_val_bpb')}")


if __name__ == "__main__":
    main()