"""Load a published checkpoint.

Each published run is a directory holding

    model.safetensors   the parameters
    config.json         {"kind", "cfg", "train_args", "params", ...}

`kind` selects the architecture; `cfg` is the exact dataclass field set the run
was built with, so nothing is inferred from tensor shapes.
"""
from __future__ import annotations

import json

import torch
from safetensors.torch import load_file

from .common import ckpt_root


def load_config(run: str) -> dict:
    return json.loads((ckpt_root() / run / "config.json").read_text())


def build(kind: str, cfg: dict):
    if kind == "rmt":
        from .model.rmt import RMT, RMTConfig
        return RMT(RMTConfig(**{k: v for k, v in cfg.items()
                                if k in RMTConfig.__dataclass_fields__}))
    if kind in ("pnsr", "pnsr_k1"):
        from .model.pnsr import PNSR, PNSRConfig
        return PNSR(PNSRConfig(**{k: v for k, v in cfg.items()
                                  if k in PNSRConfig.__dataclass_fields__}))
    if kind in ("bind", "unbound"):
        from .model.bind import BindConfig, PNSBind
        return PNSBind(BindConfig(**{k: v for k, v in cfg.items()
                                     if k in BindConfig.__dataclass_fields__}))
    from .model.tx import TX, TXConfig
    return TX(TXConfig(**{k: v for k, v in cfg.items()
                          if k in TXConfig.__dataclass_fields__}))


def load_model(run: str, device: str = "cuda"):
    """Returns (model in eval mode on `device`, kind, config dict)."""
    conf = load_config(run)
    model = build(conf["kind"], conf["cfg"])
    state = load_file(str(ckpt_root() / run / "model.safetensors"), device="cpu")
    model.load_state_dict(state, strict=True)
    model.eval().to(device)
    return model, conf["kind"], conf


def param_count(model) -> int:
    """Number of distinct trainable parameters — the paper's budget figure.

    The event encoder and the record encoder *share* one 8192x384 token
    embedding table (the parameter is tied, not duplicated), so `state_dict`
    lists it under two keys while the model holds it once. `parameters()`
    already de-duplicates shared tensors, which is exactly the quantity the
    matched-parameter comparison uses.
    """
    return sum(p.numel() for p in model.parameters())


@torch.no_grad()
def verify(run: str, device: str = "cpu") -> dict:
    """Load a checkpoint and check it against its recorded provenance."""
    conf = load_config(run)
    model, _, _ = load_model(run, device)
    n = param_count(model)
    tensors = sum(v.numel() for v in model.state_dict().values())
    return dict(run=run, kind=conf["kind"], loaded=True,
                params=n, params_match=n == conf["params"],
                state_dict_tensor_total=tensors,
                tensor_total_match=tensors == conf["state_dict_tensor_total"])