"""Unstructured Wanda / SparseGPT on SparseMoE expert Linears.

Shapes are unchanged. Only expert `up_gate_proj` and `down_proj` are pruned.
Layer-sequential: prune block i, then re-forward before scoring block i+1.
"""

from __future__ import annotations

import json
import os
import types
from typing import Any

import torch
import torch.nn as nn

from inference.model.dit.dit_module import SparseMoEMLP

from .sparsegpt_full import SparseGPT
from .wanda_unstructured import WrappedGPT, prune_linear_wanda

FORMAT = "davinci_unstructured_v1"
CAPTURE_ENV = "MAGI_UNSTRUCTURED_CAPTURE"
OUT_ROOT_ENV = "MAGI_UNSTRUCTURED_OUT_ROOT"
METHODS_ENV = "MAGI_UNSTRUCTURED_METHODS"
SPARSITIES_ENV = "MAGI_UNSTRUCTURED_SPARSITIES"
STRIDE_ENV = "MAGI_UNSTRUCTURED_STRIDE"
MAX_ENV = "MAGI_UNSTRUCTURED_MAX"
SOURCE_ENV = "MAGI_UNSTRUCTURED_SOURCE_CKPT"

__all__ = [
    "FORMAT",
    "moe_expert_linears",
    "prune_linear_wanda",
    "sequential_unstructured",
    "global_mlp_sparsity",
    "maybe_attach_capture",
    "maybe_prune_captured",
]


def moe_expert_linears(mlp: SparseMoEMLP) -> dict[str, nn.Linear]:
    out: dict[str, nn.Linear] = {}
    for eid, expert in enumerate(mlp.experts):
        out[f"expert{eid}.up_gate_proj"] = expert.up_gate_proj
        out[f"expert{eid}.down_proj"] = expert.down_proj
    return out


def global_mlp_sparsity(model) -> float:
    total = 0
    zeros = 0
    for layer in model.block.layers:
        mlp = getattr(layer, "mlp", None)
        if not isinstance(mlp, SparseMoEMLP):
            continue
        for lin in moe_expert_linears(mlp).values():
            w = lin.weight.data
            total += int(w.numel())
            zeros += int((w == 0).sum())
    return zeros / max(total, 1)


def _run_layer(layer, sample: dict[str, Any]):
    with torch.amp.autocast("cuda", dtype=torch.bfloat16):
        out = layer(
            sample["h"],
            sample["rope"],
            sample["permute_mapping"],
            sample["inv_permute_mapping"],
            sample["varlen_handler"],
            sample["local_attn_handler"],
            sample["modality_dispatcher"],
            sample["cp_split_sizes"],
        )
    return out.detach() if torch.is_tensor(out) else out


def _clone_samples(samples: list[dict[str, Any]]) -> list[dict[str, Any]]:
    cloned = []
    for sample in samples:
        item = dict(sample)
        item["h"] = sample["h"].clone()
        cloned.append(item)
    return cloned


@torch.no_grad()
def sequential_unstructured(
    model,
    samples: list[dict[str, Any]],
    sparsity: float,
    method: str,
    verbose: bool = True,
) -> dict:
    method = method.strip().lower()
    if method not in ("wanda", "sparsegpt"):
        raise ValueError(f"method must be wanda or sparsegpt, got {method}")
    # Grouped GEMM calls F.linear on packed weights and skips nn.Linear.forward,
    # so Wanda/SparseGPT hooks would see nsamples=0 and zero the expert.
    os.environ["MAGI_MOE_GROUPED"] = "off"
    work = _clone_samples(samples)
    layer_meta: dict[str, Any] = {}
    captured_any = False
    for layer_idx, layer in enumerate(model.block.layers):
        mlp = getattr(layer, "mlp", None)
        is_moe = isinstance(mlp, SparseMoEMLP)
        linears = moe_expert_linears(mlp) if is_moe else {}
        wrappers: dict[str, Any] = {}
        handles = []
        if is_moe:
            if method == "wanda":
                wrappers = {name: WrappedGPT(lin) for name, lin in linears.items()}
            else:
                wrappers = {name: SparseGPT(lin) for name, lin in linears.items()}

            def _make_hook(wrapper):
                def _hook(_, inp, out):
                    wrapper.add_batch(inp[0].data, out.data)
                return _hook

            for name, lin in linears.items():
                handles.append(lin.register_forward_hook(_make_hook(wrappers[name])))
            for sample in work:
                _run_layer(layer, sample)
            for handle in handles:
                handle.remove()
            names_meta = {}
            layer_ns = [int(wrappers[name].nsamples) for name in linears]
            if sum(layer_ns) == 0:
                raise RuntimeError(
                    f"unstructured {method} layer {getattr(mlp, 'layer_idx', layer_idx)} "
                    "captured 0 expert activations; MAGI_MOE_GROUPED must be off"
                )
            captured_any = True
            for name, lin in linears.items():
                nsamples = int(wrappers[name].nsamples)
                if nsamples == 0:
                    names_meta[name] = {
                        "sparsity": 0.0,
                        "shape": list(lin.weight.shape),
                        "nsamples": 0,
                        "skipped": True,
                    }
                    if hasattr(wrappers[name], "free"):
                        wrappers[name].free()
                    continue
                if method == "wanda":
                    sp = prune_linear_wanda(lin, wrappers[name].scaler_row, sparsity)
                else:
                    sp = wrappers[name].fasterprune(sparsity, verbose=verbose)
                    wrappers[name].free()
                names_meta[name] = {
                    "sparsity": float(sp),
                    "shape": list(lin.weight.shape),
                    "nsamples": nsamples,
                }
            layer_meta[str(int(getattr(mlp, "layer_idx", layer_idx)))] = names_meta
        for sample in work:
            sample["h"] = _run_layer(layer, sample)
        if verbose and is_moe:
            print(
                f"unstructured {method} layer {getattr(mlp, 'layer_idx', layer_idx)} "
                f"sparsity={sparsity:.3f}",
                flush=True,
            )
        torch.cuda.empty_cache()
    if not captured_any:
        raise RuntimeError(f"unstructured {method} captured 0 expert activations")
    return {
        "method": method,
        "sparsity": float(sparsity),
        "keep_ratio": float(1.0 - sparsity),
        "global_mlp_sparsity": global_mlp_sparsity(model),
        "n_samples": len(samples),
        "layers": layer_meta,
    }


def maybe_attach_capture(model) -> dict | None:
    flag = os.getenv(CAPTURE_ENV, "").strip()
    if flag not in ("1", "true", "True", "yes"):
        return None
    stride = int(os.getenv(STRIDE_ENV, "8") or "8")
    max_samples = int(os.getenv(MAX_ENV, "40") or "40")
    block = model.block
    orig = block.forward
    bag: dict[str, Any] = {"samples": [], "calls": 0, "orig": orig}

    def wrapped(
        self,
        x,
        rope,
        permute_mapping,
        inv_permute_mapping,
        varlen_handler,
        local_attn_handler,
        modality_dispatcher,
        cp_split_sizes,
    ):
        bag["calls"] += 1
        if bag["calls"] % stride == 0 and len(bag["samples"]) < max_samples:
            bag["samples"].append(
                {
                    "h": x.detach().clone(),
                    "rope": rope,
                    "permute_mapping": permute_mapping,
                    "inv_permute_mapping": inv_permute_mapping,
                    "varlen_handler": varlen_handler,
                    "local_attn_handler": local_attn_handler,
                    "modality_dispatcher": modality_dispatcher,
                    "cp_split_sizes": cp_split_sizes,
                }
            )
        return orig(
            x,
            rope,
            permute_mapping,
            inv_permute_mapping,
            varlen_handler,
            local_attn_handler,
            modality_dispatcher,
            cp_split_sizes,
        )

    block.forward = types.MethodType(wrapped, block)
    model._unstr_capture = bag
    print(
        f"unstructured capture stride={stride} max={max_samples}",
        flush=True,
    )
    return bag


def save_unstructured_checkpoint(path: str, model, meta: dict) -> None:
    payload = {
        "format": FORMAT,
        "model": model.state_dict(),
        "unstructured_meta": meta,
    }
    os.makedirs(os.path.dirname(os.path.abspath(path)) or ".", exist_ok=True)
    torch.save(payload, path)


def maybe_prune_captured(model) -> list[str]:
    out_root = os.getenv(OUT_ROOT_ENV, "").strip()
    if not out_root:
        return []
    bag = getattr(model, "_unstr_capture", None)
    if not bag or not bag.get("samples"):
        raise RuntimeError(
            f"{OUT_ROOT_ENV} is set but no TransformerBlock inputs were captured"
        )
    block = model.block
    orig = bag.get("orig")
    if orig is not None:
        block.forward = orig
    methods = [m.strip() for m in os.getenv(METHODS_ENV, "wanda,sparsegpt").split(",") if m.strip()]
    sparsities = [
        float(x) for x in os.getenv(SPARSITIES_ENV, "0.25,0.50,0.75").split(",") if x.strip()
    ]
    source = os.getenv(SOURCE_ENV, "").strip()
    samples = bag["samples"]
    print(
        f"unstructured prune captured={len(samples)} calls={bag.get('calls')} "
        f"methods={methods} sparsities={sparsities}",
        flush=True,
    )
    orig_state = {k: v.detach().cpu().clone() for k, v in model.state_dict().items()}
    written: list[str] = []
    os.makedirs(out_root, exist_ok=True)
    for method in methods:
        for sparsity in sparsities:
            keep = 1.0 - float(sparsity)
            slug = f"{method}_unstr_keep{int(round(keep * 100)):03d}"
            out_path = os.path.join(out_root, f"{slug}.pt")
            if os.path.isfile(out_path):
                print(f"[skip prune] {out_path} exists", flush=True)
                written.append(out_path)
                continue
            print(f"unstructured apply {slug}", flush=True)
            model.load_state_dict(orig_state, strict=True)
            meta = sequential_unstructured(
                model, samples, sparsity=float(sparsity), method=method, verbose=True,
            )
            meta["source"] = source
            meta["captured_calls"] = int(bag.get("calls") or 0)
            save_unstructured_checkpoint(out_path, model, meta)
            meta_json = os.path.splitext(out_path)[0] + "_meta.json"
            with open(meta_json, "w", encoding="utf-8") as f:
                json.dump(meta, f, indent=2)
                f.write("\n")
            print(f"wrote {out_path} global_mlp_sparsity={meta['global_mlp_sparsity']:.4f}", flush=True)
            written.append(out_path)
            torch.cuda.empty_cache()
    return written