"""Accumulate TENP projection (and energy control) on routed tokens. No shared expert."""

from __future__ import annotations

from dataclasses import dataclass, field

import torch

from .score import PROJ_EPS, gated_energy, gated_signed_projection
from .temporal import DEFAULT_NUM_BINS, timestep_to_bin

VIDEO, AUDIO, TEXT = 0, 1, 2
VALID_MODALITIES = (VIDEO, AUDIO, TEXT)
DEFAULT_PERTURB_RATIO = 0.05
DEFAULT_PERTURB_REPEATS = 3
DEFAULT_TOKEN_CACHE_PER_BIN = 0


def _as_2d_table(t: torch.Tensor, d_int: int) -> torch.Tensor:
    if t.ndim == 2 and t.shape[0] == 3:
        return t
    raise ValueError(f"expected [3, d_int={d_int}] table, got {tuple(t.shape)}")


@dataclass
class ExpertStats:
    projection: torch.Tensor
    energy: torch.Tensor
    gram: torch.Tensor | None = None
    n_tokens: int = 0
    n_tokens_mod: torch.Tensor = None
    projection_bin: torch.Tensor | None = None
    energy_bin: torch.Tensor | None = None
    n_tokens_bin: torch.Tensor | None = None
    n_tokens_mod_bin: torch.Tensor | None = None
    sensitivity_sum: torch.Tensor | None = None
    sensitivity_count: torch.Tensor | None = None
    rel_abs_z: torch.Tensor | None = None
    token_cache_per_bin: int = 0
    cache_z: list | None = None
    cache_y: list | None = None
    cache_gate: list | None = None
    cache_mod: list | None = None
    _cache_seen: list = field(default_factory=list)

    @property
    def num_bins(self) -> int:
        if self.projection_bin is None:
            return 1
        return int(self.projection_bin.shape[0])

    @classmethod
    def zeros(
        cls,
        d_int: int,
        device,
        num_bins: int = DEFAULT_NUM_BINS,
        collect_gram: bool = True,
        token_cache_per_bin: int = DEFAULT_TOKEN_CACHE_PER_BIN,
    ) -> "ExpertStats":
        k = max(int(num_bins), 1)
        cap = max(int(token_cache_per_bin), 0)
        gram = (
            torch.zeros(d_int, d_int, device=device, dtype=torch.float32)
            if collect_gram
            else None
        )
        return cls(
            projection=torch.zeros(3, d_int, device=device, dtype=torch.float32),
            energy=torch.zeros(3, d_int, device=device, dtype=torch.float32),
            gram=gram,
            n_tokens_mod=torch.zeros(3, device=device, dtype=torch.long),
            projection_bin=torch.zeros(k, 3, d_int, device=device, dtype=torch.float32),
            energy_bin=torch.zeros(k, 3, d_int, device=device, dtype=torch.float32),
            n_tokens_bin=torch.zeros(k, device=device, dtype=torch.long),
            n_tokens_mod_bin=torch.zeros(k, 3, device=device, dtype=torch.long),
            sensitivity_sum=torch.zeros(k, device=device, dtype=torch.float32),
            sensitivity_count=torch.zeros(k, device=device, dtype=torch.float32),
            rel_abs_z=torch.zeros(d_int, device=device, dtype=torch.float32),
            token_cache_per_bin=cap,
            cache_z=[None] * k if cap else None,
            cache_y=[None] * k if cap else None,
            cache_gate=[None] * k if cap else None,
            _cache_seen=[0] * k if cap else [],
        )

    def _ensure_bins(self, d_int: int) -> None:
        k = self.num_bins
        device = self.projection.device
        if self.projection_bin is None:
            self.projection_bin = torch.zeros(k, 3, d_int, device=device, dtype=torch.float32)
            self.projection_bin[0] = self.projection
        if self.energy_bin is None:
            self.energy_bin = torch.zeros(k, 3, d_int, device=device, dtype=torch.float32)
            self.energy_bin[0] = self.energy
        if self.n_tokens_bin is None:
            self.n_tokens_bin = torch.zeros(k, device=device, dtype=torch.long)
            self.n_tokens_bin[0] = int(self.n_tokens)
        if self.n_tokens_mod_bin is None:
            self.n_tokens_mod_bin = torch.zeros(k, 3, device=device, dtype=torch.long)
            self.n_tokens_mod_bin[0] = self.n_tokens_mod
        if self.sensitivity_sum is None:
            self.sensitivity_sum = torch.zeros(k, device=device, dtype=torch.float32)
        if self.sensitivity_count is None:
            self.sensitivity_count = torch.zeros(k, device=device, dtype=torch.float32)
        if self.rel_abs_z is None:
            self.rel_abs_z = torch.zeros(d_int, device=device, dtype=torch.float32)

    def sensitivity(self, eps: float = 1e-8) -> torch.Tensor:
        self._ensure_bins(int(self.projection.shape[-1]))
        return self.sensitivity_sum / self.sensitivity_count.clamp_min(eps)

    def update(
        self,
        z: torch.Tensor,
        y: torch.Tensor,
        gate: torch.Tensor,
        modality: torch.Tensor,
        W_down: torch.Tensor,
        bin_id: int = 0,
        collect_sensitivity: bool = False,
        perturb_ratio: float = DEFAULT_PERTURB_RATIO,
        perturb_repeats: int = DEFAULT_PERTURB_REPEATS,
        generator: torch.Generator | None = None,
    ) -> None:
        valid = (modality >= 0) & (modality <= TEXT)
        if int(valid.sum()) == 0:
            return
        z = z[valid]
        y = y[valid]
        gate = gate[valid]
        modality = modality[valid]
        q = gated_signed_projection(z, y, gate, W_down)
        e = gated_energy(z, gate)
        d_int = int(z.shape[-1])
        self._ensure_bins(d_int)
        k = self.num_bins
        b = min(max(int(bin_id), 0), k - 1)
        if self.gram is not None:
            z_g = z.float() * gate.float().reshape(-1, 1)
            self.gram.add_(z_g.T @ z_g)
        n = int(q.shape[0])
        self.n_tokens += n
        self.n_tokens_bin[b] += n
        y_norm = y.float().norm(dim=-1, keepdim=True).clamp_min(PROJ_EPS)
        self.rel_abs_z.add_((z.float().abs() / y_norm).sum(dim=0))
        for m in VALID_MODALITIES:
            mask = modality == m
            n_m = int(mask.sum())
            if n_m == 0:
                continue
            self.n_tokens_mod[m] += n_m
            self.n_tokens_mod_bin[b, m] += n_m
            q_sum = q[mask].sum(dim=0)
            e_sum = e[mask].sum(dim=0)
            self.projection[m] += q_sum
            self.energy[m] += e_sum
            self.projection_bin[b, m] += q_sum
            self.energy_bin[b, m] += e_sum
        if collect_sensitivity:
            d = perturbation_sensitivity(
                z, y, W_down, perturb_ratio, perturb_repeats, generator
            )
            self.sensitivity_sum[b] += float(d) * n
            self.sensitivity_count[b] += n
        self._maybe_cache(z, y, gate, modality, b)

    def _maybe_cache(
        self,
        z: torch.Tensor,
        y: torch.Tensor,
        gate: torch.Tensor,
        modality: torch.Tensor,
        bin_id: int,
    ) -> None:
        """Reservoir of `token_cache_per_bin` tokens per (bin, modality).

        Per-modality reservoirs keep enough audio/text tokens (they are ~5% of routed
        tokens) to measure per-modality output error after pruning. `cache_mod[bin]`
        holds the modality id aligned with `cache_z[bin]`. Tensors are stored fp16.
        """
        cap = int(self.token_cache_per_bin)
        if cap <= 0:
            return
        k = self.num_bins
        if self.cache_z is None:
            self.cache_z = [None] * k
            self.cache_y = [None] * k
            self.cache_gate = [None] * k
            self.cache_mod = [None] * k
            self._cache_seen = [[0, 0, 0] for _ in range(k)]
        if self.cache_mod is None:
            self.cache_mod = [None] * k
        if not isinstance(self._cache_seen[bin_id], list):
            self._cache_seen[bin_id] = [int(self._cache_seen[bin_id]), 0, 0]
        z_all = z.detach().to(torch.float16).cpu()
        y_all = y.detach().to(torch.float16).cpu()
        g_all = gate.detach().float().reshape(-1).cpu()
        m_all = modality.detach().to(torch.int8).reshape(-1).cpu()
        cur_z = self.cache_z[bin_id]
        cur_mod = self.cache_mod[bin_id]
        if cur_z is not None and cur_mod is None:
            cur_mod = torch.full((int(cur_z.shape[0]),), -1, dtype=torch.int8)
        parts_z, parts_y, parts_g, parts_m = [], [], [], []
        for m in VALID_MODALITIES:
            sel = (m_all == m).nonzero(as_tuple=False).reshape(-1)
            n = int(sel.numel())
            if cur_z is not None:
                keep_m = (cur_mod == m).nonzero(as_tuple=False).reshape(-1)
                z_m, y_m, g_m = cur_z[keep_m], self.cache_y[bin_id][keep_m], self.cache_gate[bin_id][keep_m]
            else:
                z_m = torch.empty(0, int(z_all.shape[-1]), dtype=torch.float16)
                y_m = torch.empty(0, int(y_all.shape[-1]), dtype=torch.float16)
                g_m = torch.empty(0, dtype=torch.float32)
            seen = int(self._cache_seen[bin_id][m])
            if n > 0:
                filled = int(z_m.shape[0])
                if filled < cap:
                    take = min(n, cap - filled)
                    z_m = torch.cat([z_m, z_all[sel[:take]]], dim=0)
                    y_m = torch.cat([y_m, y_all[sel[:take]]], dim=0)
                    g_m = torch.cat([g_m, g_all[sel[:take]]], dim=0)
                    sel = sel[take:]
                    seen += take
                rem = int(sel.numel())
                if rem > 0:
                    # Reservoir sampling for the remainder: token i (1-based within this
                    # call) replaces slot r ~ U[0, seen+i) if r < cap.
                    counts = torch.arange(1, rem + 1, dtype=torch.float64) + float(seen)
                    r = (torch.rand(rem, dtype=torch.float64) * counts).long()
                    hit = (r < cap).nonzero(as_tuple=False).reshape(-1)
                    for i in hit.tolist():  # sequential so later hits overwrite earlier
                        j = int(sel[i])
                        slot = int(r[i])
                        z_m[slot] = z_all[j]
                        y_m[slot] = y_all[j]
                        g_m[slot] = g_all[j]
                    seen += rem
            self._cache_seen[bin_id][m] = seen
            parts_z.append(z_m)
            parts_y.append(y_m)
            parts_g.append(g_m)
            parts_m.append(torch.full((int(z_m.shape[0]),), int(m), dtype=torch.int8))
        self.cache_z[bin_id] = torch.cat(parts_z, dim=0)
        self.cache_y[bin_id] = torch.cat(parts_y, dim=0)
        self.cache_gate[bin_id] = torch.cat(parts_g, dim=0)
        self.cache_mod[bin_id] = torch.cat(parts_m, dim=0)

    def state_dict(self) -> dict:
        def _cpu_list(items):
            if items is None:
                return None
            out = []
            for t in items:
                out.append(None if t is None else t.detach().cpu())
            return out

        return {
            "projection": self.projection.detach().cpu(),
            "energy": self.energy.detach().cpu(),
            "gram": None if self.gram is None else self.gram.detach().cpu(),
            "n_tokens": int(self.n_tokens),
            "n_tokens_mod": self.n_tokens_mod.detach().cpu(),
            "projection_bin": None if self.projection_bin is None else self.projection_bin.detach().cpu(),
            "energy_bin": None if self.energy_bin is None else self.energy_bin.detach().cpu(),
            "n_tokens_bin": None if self.n_tokens_bin is None else self.n_tokens_bin.detach().cpu(),
            "n_tokens_mod_bin": None if self.n_tokens_mod_bin is None else self.n_tokens_mod_bin.detach().cpu(),
            "sensitivity_sum": None if self.sensitivity_sum is None else self.sensitivity_sum.detach().cpu(),
            "sensitivity_count": None if self.sensitivity_count is None else self.sensitivity_count.detach().cpu(),
            "rel_abs_z": None if self.rel_abs_z is None else self.rel_abs_z.detach().cpu(),
            "token_cache_per_bin": int(self.token_cache_per_bin),
            "cache_z": _cpu_list(self.cache_z),
            "cache_y": _cpu_list(self.cache_y),
            "cache_gate": _cpu_list(self.cache_gate),
            "cache_mod": _cpu_list(self.cache_mod),
        }

    @classmethod
    def from_state(cls, payload: dict, device="cpu") -> "ExpertStats":
        projection = payload["projection"].to(device=device, dtype=torch.float32)
        energy = payload["energy"].to(device=device, dtype=torch.float32)
        d_int = int(projection.shape[-1])
        stats = cls(
            projection=_as_2d_table(projection, d_int),
            energy=_as_2d_table(energy, d_int),
            gram=(
                None
                if payload.get("gram") is None
                else payload["gram"].to(device=device, dtype=torch.float32)
            ),
            n_tokens=int(payload["n_tokens"]),
            n_tokens_mod=payload["n_tokens_mod"].to(device=device, dtype=torch.long),
        )
        proj_bin = payload.get("projection_bin")
        if proj_bin is None:
            stats.projection_bin = stats.projection.unsqueeze(0).clone()
            stats.energy_bin = stats.energy.unsqueeze(0).clone()
            stats.n_tokens_bin = torch.tensor(
                [stats.n_tokens], device=device, dtype=torch.long
            )
            stats.n_tokens_mod_bin = stats.n_tokens_mod.unsqueeze(0).clone()
            stats.sensitivity_sum = torch.zeros(1, device=device, dtype=torch.float32)
            stats.sensitivity_count = torch.zeros(1, device=device, dtype=torch.float32)
        else:
            stats.projection_bin = proj_bin.to(device=device, dtype=torch.float32)
            energy_bin = payload.get("energy_bin")
            stats.energy_bin = (
                stats.energy.unsqueeze(0).clone()
                if energy_bin is None
                else energy_bin.to(device=device, dtype=torch.float32)
            )
            n_bin = payload.get("n_tokens_bin")
            k = int(stats.projection_bin.shape[0])
            stats.n_tokens_bin = (
                torch.zeros(k, device=device, dtype=torch.long)
                if n_bin is None
                else n_bin.to(device=device, dtype=torch.long)
            )
            n_mod_bin = payload.get("n_tokens_mod_bin")
            stats.n_tokens_mod_bin = (
                torch.zeros(k, 3, device=device, dtype=torch.long)
                if n_mod_bin is None
                else n_mod_bin.to(device=device, dtype=torch.long)
            )
            sens_sum = payload.get("sensitivity_sum")
            stats.sensitivity_sum = (
                torch.zeros(k, device=device, dtype=torch.float32)
                if sens_sum is None
                else sens_sum.to(device=device, dtype=torch.float32)
            )
            sens_n = payload.get("sensitivity_count")
            stats.sensitivity_count = (
                torch.zeros(k, device=device, dtype=torch.float32)
                if sens_n is None
                else sens_n.to(device=device, dtype=torch.float32)
            )
        rel = payload.get("rel_abs_z")
        stats.rel_abs_z = (
            torch.zeros(d_int, device=device, dtype=torch.float32)
            if rel is None
            else rel.to(device=device, dtype=torch.float32)
        )
        stats.token_cache_per_bin = int(payload.get("token_cache_per_bin") or 0)
        cache_z = payload.get("cache_z")
        if cache_z is None:
            stats.cache_z = None
            stats.cache_y = None
            stats.cache_gate = None
            stats.cache_mod = None
        else:
            stats.cache_z = [None if t is None else t.float().cpu() for t in cache_z]
            cache_y = payload.get("cache_y") or [None] * len(cache_z)
            cache_g = payload.get("cache_gate") or [None] * len(cache_z)
            cache_m = payload.get("cache_mod") or [None] * len(cache_z)
            stats.cache_y = [None if t is None else t.float().cpu() for t in cache_y]
            stats.cache_gate = [None if t is None else t.float().cpu() for t in cache_g]
            stats.cache_mod = [None if t is None else t.to(torch.int8).cpu() for t in cache_m]
        return stats


def stacked_token_cache(
    stats: ExpertStats,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None:
    """Concatenate per-bin token reservoirs. Returns (z, y, gate) or None."""
    if not stats.cache_z:
        return None
    zs, ys, gs = [], [], []
    for z, y, g in zip(stats.cache_z, stats.cache_y or [], stats.cache_gate or []):
        if z is None or int(z.numel()) == 0:
            continue
        zs.append(z.float())
        ys.append(y.float())
        gs.append(g.float().reshape(-1))
    if not zs:
        return None
    return torch.cat(zs, dim=0), torch.cat(ys, dim=0), torch.cat(gs, dim=0)


def stacked_token_cache_mod(
    stats: ExpertStats,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor] | None:
    """Like stacked_token_cache but also returns the per-token modality id (-1 if unknown)."""
    packed = stacked_token_cache(stats)
    if packed is None:
        return None
    z, y, g = packed
    mods = []
    for i, zb in enumerate(stats.cache_z):
        if zb is None or int(zb.numel()) == 0:
            continue
        m = None if stats.cache_mod is None or i >= len(stats.cache_mod) else stats.cache_mod[i]
        if m is None or int(m.numel()) != int(zb.shape[0]):
            m = torch.full((int(zb.shape[0]),), -1, dtype=torch.int8)
        mods.append(m.to(torch.int8))
    return z, y, g, torch.cat(mods, dim=0)


def perturbation_sensitivity(
    z: torch.Tensor,
    y: torch.Tensor,
    W_down: torch.Tensor,
    ratio: float,
    repeats: int,
    generator: torch.Generator | None,
) -> float:
    """Relative L2 of randomly masking ρ mid-neurons, averaged over repeats."""
    n, d_int = z.shape
    del n
    n_mask = max(1, int(round(float(ratio) * d_int)))
    n_mask = min(n_mask, d_int)
    w = W_down.detach().to(device=z.device, dtype=torch.float32)
    z_f = z.float()
    y_norm = y.float().norm(dim=-1).clamp_min(PROJ_EPS)
    acc = 0.0
    r = max(int(repeats), 1)
    for _ in range(r):
        idx = torch.randperm(d_int, generator=generator)[:n_mask].to(device=z.device)
        delta = z_f.index_select(1, idx) @ w.index_select(1, idx).T
        acc += float((delta.norm(dim=-1) / y_norm).mean())
    return acc / r


_perturbation_sensitivity = perturbation_sensitivity


@dataclass
class LayerCollector:
    d_int: int
    num_experts: int
    device: torch.device
    mlp: object
    stats: dict[int, ExpertStats] = field(default_factory=dict)
    num_bins: int = DEFAULT_NUM_BINS
    collect_sensitivity: bool = True
    perturb_ratio: float = DEFAULT_PERTURB_RATIO
    perturb_repeats: int = DEFAULT_PERTURB_REPEATS
    collect_gram: bool = True
    token_cache_per_bin: int = DEFAULT_TOKEN_CACHE_PER_BIN
    current_timestep: float | None = None
    current_step: int | None = None
    num_steps: int | None = None

    def __post_init__(self):
        if self.mlp is None:
            raise ValueError("TENP collector needs mlp to read down_proj.weight")
        self.num_bins = max(int(self.num_bins), 1)
        self.token_cache_per_bin = max(int(self.token_cache_per_bin), 0)
        self.stats = {
            eid: ExpertStats.zeros(
                self.d_int,
                self.device,
                num_bins=self.num_bins,
                collect_gram=bool(self.collect_gram),
                token_cache_per_bin=self.token_cache_per_bin,
            )
            for eid in range(self.num_experts)
        }
        self._sens_gen = torch.Generator(device="cpu")
        self._sens_gen.manual_seed(0)

    def _down_weight(self, expert_id: int) -> torch.Tensor:
        return self.mlp.experts[int(expert_id)].down_proj.weight

    def current_bin(self) -> int:
        mlp = self.mlp
        t = getattr(mlp, "diffusion_timestep", None)
        if t is None:
            t = self.current_timestep
        step_index = getattr(mlp, "diffusion_step_index", None)
        if step_index is None:
            step_index = self.current_step
        num_steps = getattr(mlp, "diffusion_num_steps", None)
        if num_steps is None:
            num_steps = self.num_steps
        return timestep_to_bin(
            t,
            self.num_bins,
            step_index=step_index,
            num_steps=num_steps,
        )

    def __call__(
        self,
        expert_id: int,
        modality: torch.Tensor,
        gate: torch.Tensor,
        z: torch.Tensor,
        y: torch.Tensor,
    ) -> None:
        self.stats[int(expert_id)].update(
            z,
            y,
            gate,
            modality,
            W_down=self._down_weight(expert_id),
            bin_id=self.current_bin(),
            collect_sensitivity=self.collect_sensitivity,
            perturb_ratio=self.perturb_ratio,
            perturb_repeats=self.perturb_repeats,
            generator=self._sens_gen,
        )

    def state_dict(self) -> dict:
        return {
            "d_int": int(self.d_int),
            "num_experts": int(self.num_experts),
            "num_bins": int(self.num_bins),
            "stats": {str(eid): st.state_dict() for eid, st in self.stats.items()},
        }

    @classmethod
    def from_state(cls, payload: dict, mlp, device="cpu") -> "LayerCollector":
        col = cls(
            d_int=int(payload["d_int"]),
            num_experts=int(payload["num_experts"]),
            device=torch.device(device),
            mlp=mlp,
            num_bins=int(payload.get("num_bins") or DEFAULT_NUM_BINS),
            collect_sensitivity=False,
            collect_gram=False,
            token_cache_per_bin=0,
        )
        for eid_s, st in payload["stats"].items():
            col.stats[int(eid_s)] = ExpertStats.from_state(st, device=device)
        if col.stats:
            col.num_bins = next(iter(col.stats.values())).num_bins
        return col