#!/usr/bin/env python3
"""Allège un checkpoint d'entraînement pour le déploiement CPU.

Un checkpoint d'entraînement contient souvent l'état de l'optimiseur (Adam =
2 copies des poids en plus) : un modèle de 500M de paramètres peut peser
plusieurs Go alors que l'inférence n'a besoin que des poids. Ce script ne garde
que `args` + `model_state_dict` (sans le préfixe `_orig_mod.`), en fp16 par
défaut : ~2 octets/paramètre, soit ~1 Go pour 500M. Au chargement, PyTorch
reconvertit en fp32 (le calcul CPU reste en fp32).

Le fichier produit garde le même format que l'original : inference.py le
charge sans modification.

Usage :
    python scripts/slim_checkpoint.py chat_model_500M.pt chat_model_500M_slim.pt
    python scripts/slim_checkpoint.py in.pt out.pt --dtype bf16
"""
import argparse
import dataclasses
import os
import sys
from pathlib import Path

import torch

sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
import inference  # noqa: E402  (pour lire les `args` picklés depuis le script d'entraînement)

DTYPES = {"fp16": torch.float16, "bf16": torch.bfloat16, "fp32": torch.float32}


def _args_to_dict(args):
    if isinstance(args, dict):
        return dict(args)
    if dataclasses.is_dataclass(args):
        return dataclasses.asdict(args)
    return {k: v for k, v in vars(args).items() if not k.startswith("_")}


def main():
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("src")
    ap.add_argument("dst")
    ap.add_argument("--dtype", choices=DTYPES, default="fp16")
    opts = ap.parse_args()
    dtype = DTYPES[opts.dtype]

    ckpt = inference._torch_load(opts.src)
    state = {k.replace("_orig_mod.", ""): v for k, v in ckpt["model_state_dict"].items()}

    slim, bad = {}, []
    for name, tensor in state.items():
        if tensor.is_floating_point():
            converted = tensor.detach().to(dtype).contiguous()
            if dtype == torch.float16 and not torch.isfinite(converted).all():
                bad.append(name)
            slim[name] = converted
        else:
            slim[name] = tensor.detach().contiguous()
    if bad:
        sys.exit(f"❌ Débordement fp16 sur {len(bad)} tenseur(s) (ex. {bad[0]}). Relance avec --dtype bf16.")

    out = {"args": _args_to_dict(ckpt["args"]), "model_state_dict": slim, "val_loss": ckpt.get("val_loss")}
    torch.save(out, opts.dst)

    n_params = sum(t.numel() for t in slim.values())
    print(f"✅ {n_params / 1e6:.0f}M paramètres | {os.path.getsize(opts.src) / 1e9:.2f} Go -> "
          f"{os.path.getsize(opts.dst) / 1e9:.2f} Go ({opts.dtype})")


if __name__ == "__main__":
    main()