"""Archetypal MLP transcoder, RA-SAE-aligned.

Reference: Fel et al., "Archetypal SAEs", and our `ArchetypalSAEs.pdf`.

Per-feature decoder direction is parameterized as:

    D_i  =  (W_i @ C)         +   Λ_i               # row i of W_dec, in R^D
            └─── archetypal ───┘   └── relaxation ──┘

where:
- `C ∈ R^{N_pool × D}` is the *frozen* archetype pool --> K-means centroids of
  the FFN's **output** activations (`mlp_out`), since the decoder writes into
  output space and we want each atom to be a real-data convex combination of
  the *target* distribution. The pool is built once before training and
  registered as a buffer. The trainer projects drained `mlp_in` tokens through
  the frozen FFN to get `mlp_out`, then runs MiniBatchKMeans on that.
- `W ∈ R^{H × N_pool}` is a per-feature row-stochastic weight matrix; we
  parameterize it as `topk_softmax(A_logits, k)` so each feature's row is
  exactly K-sparse and sums to 1 (a sparse simplex).
- `Λ ∈ R^{H × D}` is a *relaxation* slack term, ℓ2-norm bounded by `delta`
  per row via in-place projection (`renorm_decoder_`). δ=0 ⇒ pure archetypal
  (max stability); δ→∞ ⇒ free decoder (no archetypal constraint at all).
  RA-SAE Figure 4 shows the recon-vs-stability tradeoff in δ.

Init: by default each row of A_logits is one-hot at column `i % N_pool`,
so atom `i` starts as exactly centroid `i % N_pool` (== W = eye(H, N_pool)
in the paper's notation). This is the most-archetypal possible state at
step 0 and gives a clean training trajectory.

The encoder is unchanged from MLPTranscoder (ReLU on a linear projection).
"""

from __future__ import annotations

import math
from dataclasses import dataclass

import torch
import torch.nn as nn
import torch.nn.functional as F


@dataclass
class ArchetypalConfig:
    d_in: int = 768
    d_hidden: int = 6144
    n_pool: int = 32_000              # RA-SAE default; was 10k in earlier baby
    top_k: int = 64                   # per-archetype convex-combo sparsity (decoder side)
    a_logits_init_std: float = 0.01   # only used when init_eye=False
    init_eye: bool = True             # one-hot per atom at column i % n_pool
    delta: float = 0.0                # max ‖Λ_i‖₂; 0 = pure archetypal

    # Encoder activation choice. None (default) = ReLU + L1 (the original
    # ReLU+L1 SAE setup). int K (e.g. 64, 128, 256) = Prisma-style TopK
    # encoder: per token, keep the top-K pre-activations and zero the rest,
    # then apply ReLU. This directly enforces L0 = K and removes the L1
    # hyperparameter as the sparsity dial.
    encoder_topk: int | None = None


def topk_softmax(logits: torch.Tensor, k: int) -> torch.Tensor:
    """Differentiable sparse softmax: softmax over top-k entries per row, zeros elsewhere.

    Gradient flows through the top-k values via the standard softmax backward;
    `topk_idx` is treated as a constant. If a logit drops out of the top-k from
    one step to the next, it stops receiving gradients the standard top-K
    sparsity trade-off, fine for our setting.
    """
    topk_vals, topk_idx = logits.topk(k, dim=-1)
    sm = topk_vals.softmax(dim=-1)
    out = torch.zeros_like(logits)
    return out.scatter(-1, topk_idx, sm)


class ArchetypalTranscoder(nn.Module):
    def __init__(self, cfg: ArchetypalConfig, data_pool: torch.Tensor):
        super().__init__()
        if data_pool.shape != (cfg.n_pool, cfg.d_in):
            raise ValueError(
                f"data_pool shape {tuple(data_pool.shape)} != ({cfg.n_pool}, {cfg.d_in})"
            )
        self.cfg = cfg

        # Encoder --> identical to MLPTranscoder.
        self.W_enc = nn.Parameter(torch.empty(cfg.d_hidden, cfg.d_in))
        self.b_enc = nn.Parameter(torch.zeros(cfg.d_hidden))
        nn.init.kaiming_uniform_(self.W_enc, a=5**0.5)

        # Output bias.
        self.b_dec = nn.Parameter(torch.zeros(cfg.d_in))

        # Mixing logits A_logits ∈ R^{H × N_pool}. Two init regimes:
        #   - eye:    one-hot at column (i % n_pool) → atom i starts as that centroid.
        #             After top-K softmax that single position has weight 1.
        #   - random: small Gaussian. Atoms start as roughly uniform mixtures over
        #             top-K random pool entries.
        self.A_logits = nn.Parameter(torch.zeros(cfg.d_hidden, cfg.n_pool))
        with torch.no_grad():
            if cfg.init_eye:
                row_idx = torch.arange(cfg.d_hidden)
                col_idx = row_idx % cfg.n_pool
                # Scale init logit so the post-softmax top weight is ~99%
                # regardless of top_k. Derivation:
                #   top_val = e^L / (e^L + (k-1))  ≥ 0.99
                #   ⇒ e^L ≥ 99(k-1)  ⇒ L ≥ ln(99·(k-1))
                # We use ln(100·k) which gives ≈99% for any k ≥ 1, and stays
                # small enough that the optimizer can still redistribute mass
                # during training (~ a few hundred steps to migrate an atom).
                init_logit = math.log(100.0 * max(cfg.top_k, 1))
                self.A_logits[row_idx, col_idx] = init_logit
            else:
                self.A_logits.normal_(mean=0.0, std=cfg.a_logits_init_std)

        # Relaxation slack term Λ ∈ R^{H × D}, init zero.
        self.Lambda = nn.Parameter(torch.zeros(cfg.d_hidden, cfg.d_in))

        # Frozen archetype pool --> sampled once at training start, never updated.
        self.register_buffer("data_pool", data_pool.detach().clone().to(torch.float32))

    # ---------------------------------------------------------------- helpers

    @property
    def W_dec(self) -> torch.Tensor:
        """[D, H] decoder weight = (top-K-softmax(A_logits) @ pool).T + Λ.T.

        Returned in [d_in, d_hidden] orientation to match `MLPTranscoder.W_dec`,
        so downstream code (logging, metrics) is interface-compatible.
        """
        A = topk_softmax(self.A_logits, self.cfg.top_k)
        archetypal = (A @ self.data_pool).T            # [D, H]
        relax = self.Lambda.T                           # [D, H]
        return archetypal + relax

    @torch.no_grad()
    def renorm_decoder_(self) -> None:
        """In-place project Λ rows onto the ‖·‖₂ ≤ delta ball.

        Called after every optimizer step (parity with `MLPTranscoder.renorm_decoder_`).
        With `delta == 0` this hard-zeros Λ --> the pure-archetypal regime.
        With `delta > 0` this is the bounded-Λ projection from RA-SAE step 4.

        Note: we do NOT project the simplex on `A_logits` because our
        `topk_softmax` activation enforces row-stochasticity *inside the
        forward pass*, not as a parameter projection. A_logits is free to
        roam; the decoder is always built from a top-K simplex by construction.
        """
        if self.cfg.delta <= 0.0:
            self.Lambda.zero_()
            return
        # Per-row L2 norms of Λ.
        norms = self.Lambda.norm(dim=-1, keepdim=True)            # [H, 1]
        # Multiplicative scale: 1.0 if already inside the ball, otherwise δ / ‖·‖.
        scale = torch.clamp(self.cfg.delta / norms.clamp_min(1e-12), max=1.0)
        self.Lambda.mul_(scale)

    # ----------------------------------------------------------------- forward

    def encode(self, x: torch.Tensor) -> torch.Tensor:
        pre = F.linear(x, self.W_enc, self.b_enc)
        if self.cfg.encoder_topk is None:
            # Original ReLU+L1 path: ReLU on pre-activations, L0 emerges from L1 pressure.
            return F.relu(pre)
        # TopK encoder (Prisma-style): per token, keep the top-K pre-activations
        # (passing through ReLU on the kept ones to ensure non-negativity), zero
        # the rest. Directly enforces L0 = encoder_topk by construction. This
        # decouples sparsity from the L1 coefficient.
        k = min(self.cfg.encoder_topk, pre.shape[-1])
        vals, idx = pre.topk(k, dim=-1)
        z = torch.zeros_like(pre)
        z.scatter_(-1, idx, F.relu(vals))
        return z

    def decode(self, z: torch.Tensor) -> torch.Tensor:
        return F.linear(z, self.W_dec, self.b_dec)

    def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
        z = self.encode(x)
        x_hat = self.decode(z)
        return x_hat, z