"""Save / reload physically narrowed SparseMoE experts. Rebuild shapes before load."""

from __future__ import annotations

from typing import Any

import torch

from inference.model.dit.dit_module import SparseMoEMLP

from .surgery import materialize_expert_mid, resize_expert_mid

FORMAT = "davinci_tenp_v1"


def sparse_mlp_at(model, layer_idx: int) -> SparseMoEMLP:
    mlp = model.block.layers[int(layer_idx)].mlp
    if not isinstance(mlp, SparseMoEMLP):
        raise TypeError(f"layer {layer_idx} mlp is {type(mlp).__name__}, expected SparseMoEMLP")
    return mlp


def apply_keep_map(model, keep_by_layer_expert: dict[tuple[int, int], torch.Tensor]) -> dict:
    """Physically slice experts. keys are (layer_idx, expert_id)."""
    layers: dict[str, dict] = {}
    for (layer_idx, expert_id), keep in keep_by_layer_expert.items():
        mlp = sparse_mlp_at(model, layer_idx)
        expert = mlp.experts[int(expert_id)]
        original = int(expert.down_proj.in_features)
        materialize_expert_mid(expert, keep)
        layers.setdefault(str(layer_idx), {"experts": []})
        layers[str(layer_idx)]["experts"].append(
            {
                "expert_id": int(expert_id),
                "original_width": original,
                "kept_width": int(keep.numel()),
                "keep_idx": keep.detach().cpu().long().tolist(),
            }
        )
    return layers


def rebuild_mid_shapes(model, meta: dict) -> None:
    layers = meta.get("layers") or {}
    for layer_s, layer_meta in layers.items():
        mlp = sparse_mlp_at(model, int(layer_s))
        for spec in layer_meta["experts"]:
            eid = int(spec["expert_id"])
            kept = int(spec["kept_width"])
            keep_idx = spec.get("keep_idx") or []
            if keep_idx and len(keep_idx) != kept:
                raise ValueError(
                    f"layer {layer_s} expert {eid}: keep_idx len {len(keep_idx)} != kept_width {kept}"
                )
            resize_expert_mid(mlp.experts[eid], kept)


def save_pruned_checkpoint(path: str, model, meta: dict) -> None:
    payload = {
        "format": FORMAT,
        "model": model.state_dict(),
        "tenp_meta": meta,
    }
    torch.save(payload, path)


def load_pruned_checkpoint(model, path: str, map_location="cpu") -> dict:
    payload = torch.load(path, map_location=map_location, weights_only=False)
    if payload.get("format") != FORMAT:
        raise ValueError(f"unsupported checkpoint format {payload.get('format')!r}")
    meta = payload.get("tenp_meta") or {}
    rebuild_mid_shapes(model, meta)
    model.load_state_dict(payload["model"], strict=True)
    return meta


def layer_meta_from_prune(layer_idx: int, prune_meta: dict, extra: dict | None = None) -> dict:
    return meta_from_prunes({int(layer_idx): prune_meta}, extra)


def meta_from_prunes(by_layer: dict[int, dict], extra: dict | None = None) -> dict:
    layers: dict[str, Any] = {}
    for layer_idx, prune_meta in sorted(by_layer.items()):
        experts = []
        for eid, spec in sorted(prune_meta.items()):
            experts.append(
                {
                    "expert_id": int(eid),
                    "original_width": int(spec["original_width"]),
                    "kept_width": int(spec["kept_width"]),
                    "keep_idx": list(spec["keep"]),
                    "n_tokens_mod": list(spec["n_tokens_mod"]),
                    "ols": False,
                }
            )
        layers[str(int(layer_idx))] = {"experts": experts}
    body: dict[str, Any] = {
        "version": 1,
        "score_type": (extra or {}).get("score_type", "projection"),
        "modality_fuse": (extra or {}).get("modality_fuse", "fixed311"),
        "p_mid": (extra or {}).get("p_mid"),
        "ols": False,
        "audio_primary_discount": False,
        "layers": layers,
    }
    extra = extra or {}
    for key in ("method", "variant", "allocation", "temporal", "redundancy"):
        if key in extra:
            body[key] = extra[key]
    return body