src/agiws_neural_quant/arch_replace.py

"""arch_replace — architectural replacement via dual-path QAT distillation.

Implements the "Jamba-conversion" idea (roadmap Stage 19e): replace attention
layers with Mamba (selective SSM) layers using the SAME dual-path infrastructure
as block-wise ternarization (Stage 18b), but for architectural substitution
instead of weight quantization.

Key insight: QuantizedModule(dual_path=True) + dual_path_loss is a generic
mechanism. It can ternarize (18b), Jamba-convert (19e), or train MTP (Stage 24).
One engine — three applications.

Components:
  - MambaLayer   — selective SSM (fixed recurrent state, input-dependent A/B/C/dt)
  - ReplaceModule — generic dual-path wrapper for architectural substitution
  - jamba_replace — block-wise progressive attention->Mamba replacement
"""

from __future__ import annotations

import logging
import math
from typing import Any

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

logger = logging.getLogger(__name__)


# ---------------------------------------------------------------------------
# MambaLayer — selective State-Space Model (Mamba-style)
# ---------------------------------------------------------------------------

class MambaLayer(nn.Module):
    """Selective SSM layer with fixed recurrent state (Mamba architecture).

    Unlike attention, the recurrent state is fixed-size (d_state x d_inner) and
    does NOT grow with sequence length. This makes long context cheap in memory.

    Parameters (all learnable via standard backprop, no STE needed — these are
    continuous parameters, not quantized):
      in_proj   — input projection: d_model -> 2*d_inner (gate + value branches)
      out_proj  — output projection: d_inner -> d_model
      A_log     — log-space diagonal of state matrix [d_inner, d_state]
      D         — skip-connection scale [d_inner]
      dt_proj   — timestep projection: d_inner -> d_inner (input-dependent dt)

    Forward (selective scan, simplified):
      x [B, L, d_model]
        -> in_proj -> (x_gate, x_value)  [B, L, d_inner]
        -> dt = softplus(dt_proj(x_gate))  [B, L, d_inner]
        -> A = -exp(A_log)  [d_inner, d_state]  (constant per channel)
        -> selective_scan: for each t, state[t+1] = state[t] * exp(A*dt[t]) +
                           B(t) * x_value(t);  output[t] = C(t) . state[t]
        -> out = output * silu(x_gate)  (gated)
        -> out_proj -> [B, L, d_model]

    The selective_scan is implemented as a sequential loop over time (Python-
    level). This is NOT optimized for speed (no CUDA kernel) — it is a
    reference implementation for distillation training and correctness. For
    production inference, a fused scan kernel would be needed.
    """

    def __init__(
        self,
        d_model: int,
        d_state: int = 16,
        d_inner: int | None = None,
        dt_rank: int | None = None,
        bias: bool = False,
        device: torch.device | str | None = None,
        dtype: torch.dtype = torch.float32,
    ):
        super().__init__()
        self.d_model = d_model
        self.d_state = d_state
        # d_inner defaults to 2*d_model (standard Mamba expansion).
        self.d_inner = d_inner if d_inner is not None else 2 * d_model
        # dt_rank defaults to d_model/16 (low-rank projection for dt).
        self.dt_rank = dt_rank if dt_rank is not None else max(1, d_model // 16)

        _dt = {"device": device, "dtype": dtype}

        # Input projection: x -> (gate, value, dt_base) of total 2*d_inner + dt_rank.
        self.in_proj = nn.Linear(d_model, 2 * self.d_inner + self.dt_rank, bias=bias, **_dt)

        # Timestep projection: dt_base -> d_inner (produces input-dependent dt).
        self.dt_proj = nn.Linear(self.dt_rank, self.d_inner, bias=True, **_dt)

        # A_log: learnable log-space diagonal of the state matrix.
        # Initialized as in Mamba paper: A = -arange(1, d_state+1) repeated.
        A_init = torch.arange(1, d_state + 1, dtype=torch.float32).expand(self.d_inner, d_state).contiguous()
        self.A_log = nn.Parameter(torch.log(A_init).to(**_dt))

        # D: skip-connection scale, initialized to 1.
        self.D = nn.Parameter(torch.ones(self.d_inner, **_dt))

        # Output projection: d_inner -> d_model.
        self.out_proj = nn.Linear(self.d_inner, d_model, bias=bias, **_dt)

    def _selective_scan(
        self,
        x_value: torch.Tensor,
        dt: torch.Tensor,
        B: torch.Tensor,
        C: torch.Tensor,
    ) -> torch.Tensor:
        """Sequential selective scan over time dimension.

        Args:
            x_value: [B, L, d_inner] — input values.
            dt: [B, L, d_inner] — input-dependent timesteps.
            B: [B, L, d_state] — input-dependent B matrix (interaction weights).
            C: [B, L, d_state] — input-dependent C matrix (readout weights).

        Returns:
            y: [B, L, d_inner] — scanned output.
        """
        batch, seq_len, d_inner = x_value.shape
        d_state = self.d_state
        device = x_value.device
        dtype = x_value.dtype

        # A is constant (from A_log): [d_inner, d_state], negative.
        A = -torch.exp(self.A_log.to(dtype))  # [d_inner, d_state]

        # Discretized A: dA = exp(A * dt)  [B, L, d_inner, d_state]
        # dt: [B, L, d_inner] -> [B, L, d_inner, 1]
        dt_exp = dt.unsqueeze(-1)  # [B, L, d_inner, 1]
        dA = torch.exp(A.unsqueeze(0).unsqueeze(0) * dt_exp)  # [B, L, d_inner, d_state]

        # Discretized B: dB = dt * B  [B, L, d_inner, d_state]
        # B: [B, L, d_state] -> [B, L, 1, d_state] -> broadcast to d_inner
        dB = dt_exp * B.unsqueeze(2)  # [B, L, d_inner, d_state]

        # Sequential scan (Python loop — reference implementation).
        h = torch.zeros(batch, d_inner, d_state, device=device, dtype=dtype)
        ys = []
        for t in range(seq_len):
            # state update: h = dA * h + dB * x
            h = dA[:, t] * h + dB[:, t] * x_value[:, t].unsqueeze(-1)
            # output: y = C . h  (dot product over d_state)
            y_t = (h * C[:, t].unsqueeze(1)).sum(dim=-1)  # [B, d_inner]
            ys.append(y_t)
        y = torch.stack(ys, dim=1)  # [B, L, d_inner]
        return y

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """Forward pass.

        Args:
            x: [B, L, d_model] or [L, d_model] (auto-unsqueeze batch).

        Returns:
            [B, L, d_model] or [L, d_model] (matches input batch dim).
        """
        squeeze_batch = False
        if x.dim() == 2:
            x = x.unsqueeze(0)  # [1, L, d_model]
            squeeze_batch = True

        batch, seq_len, _ = x.shape
        dtype = x.dtype

        # Input projection: x -> (gate, value, dt_base)
        xzdr = self.in_proj(x)  # [B, L, 2*d_inner + dt_rank]
        gate, value, dt_base = torch.split(
            xzdr, [self.d_inner, self.d_inner, self.dt_rank], dim=-1
        )

        # dt: input-dependent timestep via softplus(dt_proj(dt_base)).
        dt = F.softplus(self.dt_proj(dt_base))  # [B, L, d_inner]
        # Clamp dt to avoid numerical issues (exp(A*dt) overflow).
        dt = dt.clamp(max=50.0)

        # B and C matrices: derived from gate via a linear split.
        # In Mamba, B/C are input-dependent. Here we use a simplified approach:
        # project gate into d_state dims for B and C.
        # Use first d_state and second d_state slices of value for B, C.
        # This is a lightweight variant — full Mamba uses separate projections.
        # For distillation purposes (learning to mimic attention output),
        # this simplified B/C is sufficient as the layer will be trained.
        B = gate[..., :self.d_state]  # [B, L, d_state]
        C = value[..., :self.d_state]  # [B, L, d_state]

        # Selective scan.
        y = self._selective_scan(value, dt, B, C)  # [B, L, d_inner]

        # Gated output: y * silu(gate) + D * value (skip connection).
        y = y * F.silu(gate)
        y = y + self.D.unsqueeze(0).unsqueeze(0) * value

        # Output projection.
        out = self.out_proj(y)  # [B, L, d_model]

        if squeeze_batch:
            out = out.squeeze(0)
        return out

    def reset_state(self) -> None:
        """No-op — recurrent state is per-forward (ephemeral), not stored."""
        pass


# ---------------------------------------------------------------------------
# ReplaceModule — generic dual-path wrapper for architectural substitution
# ---------------------------------------------------------------------------

class ReplaceModule(nn.Module):
    """Generic dual-path wrapper: student (new architecture) vs teacher (frozen original).

    This is the architectural analogue of QuantizedModule(dual_path=True).
    Instead of quantizing weights, it replaces the module TYPE entirely:
      - student: new architecture (e.g. MambaLayer) — learnable
      - teacher: frozen original (e.g. nn.MultiheadAttention) — reference

    forward(x, path):
      - "student"  -> student.forward(x)           (for inference after training)
      - "teacher"  -> teacher.forward(x) [no_grad]  (for validation)
      - "both"     -> (student_out, teacher_out)    (for dual_path_loss)

    distillation: dual_path_loss(student_out, teacher_out, "mse")
    uses the existing training_unified.dual_path_loss function.
    """

    def __init__(
        self,
        student: nn.Module,
        teacher: nn.Module | None = None,
        dual_path: bool = True,
        teacher_forward_fn: Any = None,
        student_forward_fn: Any = None,
    ):
        super().__init__()
        self.student = student
        self._dual_path = dual_path
        self._has_teacher = teacher is not None
        # Optional custom forward functions for non-standard call signatures.
        # teacher_forward_fn(module, x) -> output  (e.g. for MultiheadAttention:
        #   lambda m, x: m(x, x, x))
        # student_forward_fn(module, x) -> output  (e.g. for MambaLayer that
        #   needs extra args)
        # If None, defaults to module(x).
        self._teacher_fn = teacher_forward_fn
        self._student_fn = student_forward_fn
        if teacher is not None:
            # Register teacher as a non-trainable submodule.
            self.teacher = teacher
            for p in self.teacher.parameters():
                p.requires_grad = False
        else:
            # No teacher — single-path mode (student only).
            self._dual_path = False

    def _call_student(self, x: torch.Tensor) -> torch.Tensor:
        if self._student_fn is not None:
            return self._student_fn(self.student, x)
        return self.student(x)

    def _call_teacher(self, x: torch.Tensor) -> torch.Tensor:
        if self._teacher_fn is not None:
            return self._teacher_fn(self.teacher, x)
        return self.teacher(x)

    def forward(
        self, x: torch.Tensor, path: str = "student"
    ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
        if path == "teacher":
            if not self._has_teacher:
                raise RuntimeError("teacher_forward called but no teacher set")
            with torch.no_grad():
                return self._call_teacher(x)
        if path == "both":
            if not self._has_teacher:
                raise RuntimeError("dual_path='both' requires a teacher")
            student_out = self._call_student(x)
            with torch.no_grad():
                teacher_out = self._call_teacher(x)
            return student_out, teacher_out
        # Default: student only.
        return self._call_student(x)

    def distillation_loss(
        self, x: torch.Tensor, loss_type: str = "mse"
    ) -> torch.Tensor:
        """Compute dual-path distillation loss (student vs teacher)."""
        from agiws_neural_quant.training_unified import dual_path_loss
        student_out, teacher_out = self.forward(x, path="both")
        return dual_path_loss(student_out, teacher_out, loss_type)

    def freeze_student(self) -> None:
        """Freeze student parameters (after training, for inference)."""
        for p in self.student.parameters():
            p.requires_grad = False

    def unfreeze_student(self) -> None:
        """Unfreeze student parameters (for further training)."""
        for p in self.student.parameters():
            p.requires_grad = True


# ---------------------------------------------------------------------------
# jamba_replace — block-wise progressive attention->Mamba replacement
# ---------------------------------------------------------------------------

def _count_parameters(module: nn.Module) -> int:
    return sum(p.numel() for p in module.parameters())


def _train_replace_block(
    replace_module: ReplaceModule,
    inputs: list[torch.Tensor],
    epochs: int = 100,
    lr: float = 1e-4,
    loss_type: str = "mse",
    verbose: bool = True,
) -> dict[str, float]:
    """Train a single ReplaceModule's student to mimic its teacher.

    Args:
        replace_module: ReplaceModule with student + teacher.
        inputs: list of calibration input tensors (each [B, L, d_model]).
        epochs: number of training epochs.
        lr: learning rate.
        loss_type: "mse", "cosine", or "kl".
        verbose: print progress.

    Returns:
        dict with 'loss_start', 'loss_end', 'cosine_end'.
    """
    from agiws_neural_quant.training_unified import dual_path_loss

    optimizer = torch.optim.Adam(
        [p for p in replace_module.student.parameters() if p.requires_grad],
        lr=lr,
    )

    loss_start = None
    loss_end = None
    cosine_end = None

    for epoch in range(epochs):
        epoch_loss = 0.0
        n_batches = 0
        for x in inputs:
            optimizer.zero_grad()
            student_out, teacher_out = replace_module(x, path="both")
            loss = dual_path_loss(student_out, teacher_out, loss_type)
            loss.backward()
            optimizer.step()
            epoch_loss += loss.item()
            n_batches += 1
        avg_loss = epoch_loss / max(n_batches, 1)
        if loss_start is None:
            loss_start = avg_loss
        loss_end = avg_loss
        if verbose and (epoch % max(1, epochs // 10) == 0 or epoch == epochs - 1):
            logger.info(f"  epoch {epoch}/{epochs}: {loss_type}={avg_loss:.6f}")

    # Final cosine similarity.
    if inputs:
        with torch.no_grad():
            s_out, t_out = replace_module(inputs[0], path="both")
            cos = F.cosine_similarity(
                s_out.flatten().unsqueeze(0),
                t_out.flatten().unsqueeze(0),
            ).item()
            cosine_end = cos

    return {
        "loss_start": loss_start or 0.0,
        "loss_end": loss_end or 0.0,
        "cosine_end": cosine_end or 0.0,
    }


def jamba_replace(
    model: nn.Module,
    calib_inputs: list[torch.Tensor],
    mamba_ratio: float = 0.85,
    d_state: int = 16,
    epochs_per_block: int = 100,
    lr: float = 1e-4,
    loss_type: str = "mse",
    teacher_call_fn: Any = None,
    verbose: bool = True,
) -> tuple[nn.Module, list[dict[str, Any]]]:
    """Replace attention layers with Mamba via block-wise dual-path distillation.

    Walks the model tree, finds attention layers (by type name containing
    "Attention" or "MultiheadAttention"), replaces mamba_ratio fraction of them
    with MambaLayer, and trains each replacement via dual-path distillation
    against the frozen original attention layer.

    Placement: attention layers are kept periodically (every N-th) so that
    ~15% remain for recall. The rest are replaced with Mamba.

    Args:
        model: source model with attention layers.
        calib_inputs: list of calibration tensors [B, L, d_model] for distillation.
        mamba_ratio: fraction of attention layers to replace (0.85 = 85% -> Mamba).
        d_state: Mamba recurrent state size.
        epochs_per_block: distillation epochs per replaced block.
        lr: learning rate for distillation.
        loss_type: "mse", "cosine", or "kl".
        teacher_call_fn: callable(module, x) -> output for non-standard teacher
                         forward signatures. Default: None — auto-detects:
                           - nn.MultiheadAttention -> lambda m, x: m(x, x, x)[0]
                             (self-attention: q=k=v=x, extract output from tuple)
                           - everything else -> module(x)
                         For cross-attention, pass your own:
                           lambda m, x: m(decoder_h, encoder_h, encoder_h)
        verbose: print progress.

    Returns:
        (modified_model, report) where report is a list of per-block dicts.
    """
    # Collect all attention modules by path (top-level only — skip nested
    # attention submodules like nn.MultiheadAttention inside a wrapper).
    all_named = list(model.named_modules())
    attn_candidates: list[str] = []
    for name, module in all_named:
        if not name:  # skip root module
            continue
        type_name = type(module).__name__
        if "Attention" in type_name or "MultiheadAttention" in type_name:
            # Skip if it's inside an already-replaced ReplaceModule.
            if not any(p in name for p in (".student.", ".teacher.")):
                attn_candidates.append(name)
    # Keep only top-level (no parent in attn_candidates).
    attn_paths: list[str] = []
    for n in attn_candidates:
        if not any(n != other and n.startswith(other + ".") for other in attn_candidates):
            attn_paths.append(n)

    if not attn_paths:
        if verbose:
            logger.info("jamba_replace: no attention layers found, nothing to replace")
        return model, []

    n_attn = len(attn_paths)
    n_keep = max(1, int(n_attn * (1.0 - mamba_ratio)))
    n_replace = n_attn - n_keep

    if verbose:
        logger.info(
            f"jamba_replace: {n_attn} attention layers, "
            f"keep {n_keep}, replace {n_replace} with Mamba"
        )

    # Placement: keep attention layers spread evenly (every N-th).
    # Indices to keep: evenly spaced across the list.
    keep_indices = set()
    if n_keep > 0:
        step = n_attn / n_keep
        for i in range(n_keep):
            keep_indices.add(int(i * step))

    report: list[dict[str, Any]] = []

    for idx, attn_path in enumerate(attn_paths):
        # Navigate to the parent module and get the attention submodule.
        parts = attn_path.split(".")
        parent = model
        for p in parts[:-1]:
            parent = getattr(parent, p)
        attr_name = parts[-1]
        attn_module = getattr(parent, attr_name)

        if idx in keep_indices:
            report.append({
                "path": attn_path,
                "action": "keep_attention",
                "index": idx,
            })
            if verbose:
                logger.info(f"  [{idx}] KEEP attention: {attn_path}")
            continue

        # Determine d_model from the attention module.
        # Try common attributes; fallback to calib_inputs shape.
        d_model = getattr(attn_module, "embed_dim", None)
        if d_model is None:
            d_model = getattr(attn_module, "d_model", None)
        if d_model is None and calib_inputs:
            d_model = calib_inputs[0].shape[-1]
        if d_model is None:
            report.append({
                "path": attn_path,
                "action": "skip_unknown_dim",
                "index": idx,
                "reason": "cannot determine d_model",
            })
            if verbose:
                logger.info(f"  [{idx}] SKIP (unknown d_model): {attn_path}")
            continue

        # Create Mamba student.
        device = next(attn_module.parameters()).device
        dtype = next(attn_module.parameters()).dtype
        mamba_student = MambaLayer(
            d_model=d_model,
            d_state=d_state,
            device=device,
            dtype=dtype,
        ).to(device=device, dtype=dtype)

        # Auto-detect teacher call function for nn.MultiheadAttention.
        effective_fn = teacher_call_fn
        if effective_fn is None:
            if isinstance(attn_module, nn.MultiheadAttention):
                # Self-attention: q=k=v=x, MHA returns (output, weights).
                effective_fn = lambda m, x: m(x, x, x)[0]

        # Create ReplaceModule with teacher = frozen attention.
        replace_mod = ReplaceModule(
            student=mamba_student,
            teacher=attn_module,
            dual_path=True,
            teacher_forward_fn=effective_fn,
        )

        # Distillation training.
        if verbose:
            logger.info(
                f"  [{idx}] REPLACE attention->Mamba: {attn_path} "
                f"(d_model={d_model}, d_state={d_state})"
            )

        block_report = _train_replace_block(
            replace_mod,
            calib_inputs,
            epochs=epochs_per_block,
            lr=lr,
            loss_type=loss_type,
            verbose=verbose,
        )

        # Freeze student and install.
        replace_mod.freeze_student()
        setattr(parent, attr_name, replace_mod)

        report.append({
            "path": attn_path,
            "action": "replace_mamba",
            "index": idx,
            "d_model": d_model,
            "d_state": d_state,
            "teacher_params": _count_parameters(attn_module),
            "student_params": _count_parameters(mamba_student),
            **block_report,
        })
        if verbose:
            logger.info(
                f"  [{idx}] DONE: cosine={block_report['cosine_end']:.4f}, "
                f"loss {block_report['loss_start']:.6f} -> {block_report['loss_end']:.6f}"
            )

    return model, report


__all__: list[str] = []