src/precision_router.py

"""
Precision-Guided Router for PC-SHO-DLM + MSA (Direction C)

Replaces the learned RouterProjector with routing driven by existing precision
heads. The core insight: precision heads already estimate per-token uncertainty.
Tokens the model is uncertain about should dominate the retrieval query, because
those are the tokens that *need* external memory most.

Key equation:
    h_pool = sum_i (1/Pi_i) * h_i  /  sum_i (1/Pi_i)

    score_j = cos(h_pool, KR_j)

Low-precision (high-uncertainty) tokens contribute MORE to the pooled query.
This eliminates the separate router projectors (W_QR, W_KR) entirely --
the precision heads already contain the information needed for routing.

Additionally, `precision_gain(doc)` estimates how much retrieving a document
would reduce total uncertainty, enabling principled active-retrieval decisions.
"""

import math
from typing import Optional, Tuple, List, Dict

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

from model import PrecisionHead, PCSHODLM, PCSHOConfig
from msa import (
    RouterProjector,
    MemoryBank,
    MSALayer,
    MSAConfig,
    chunk_mean_pool,
    create_msa_layers,
)


class PrecisionGuidedRouter(nn.Module):
    """Routes queries to memory documents using precision heads as signals.

    Drop-in replacement for RouterProjector in MSALayer. Instead of learned
    W_QR / W_KR projections, it uses existing precision estimates to weight
    tokens when forming the retrieval query.

    The precision head for a layer gives (tok_precision, ch_precision):
        tok_precision: (B, S, 1)  -- per-token confidence
        ch_precision:  (B, 1, D)  -- per-channel confidence

    Uncertainty = 1 / precision. We want uncertain tokens to dominate routing.

    Attributes:
        precision_head: the PrecisionHead from the model layer being routed
        key_proj: lightweight projection from d_model -> router_dim for keys
                  (we still need keys in a shared space with the query)
        router_dim: dimensionality of the routing space
        n_heads: kept for interface compatibility with RouterProjector
        d_head: router_dim // n_heads
        eps: numerical stability floor for precision inversion
    """

    def __init__(
        self,
        precision_head: PrecisionHead,
        d_model: int,
        router_dim: int,
        n_heads: int,
        eps: float = 1e-6,
    ):
        super().__init__()
        self.precision_head = precision_head
        self.d_model = d_model
        self.router_dim = router_dim
        self.n_heads = n_heads
        self.d_head = router_dim // n_heads
        self.eps = eps

        # Lightweight query projection: precision-pooled hidden -> router space.
        # This is much smaller than RouterProjector's separate q_proj+k_proj
        # because the heavy lifting is done by the precision head.
        self.q_proj = nn.Linear(d_model, router_dim)

        # Key projection is still needed: documents were encoded independently
        # and we need their representations in the same router space.
        self.k_proj = nn.Linear(d_model, router_dim)

    def compute_precision_weights(
        self,
        h: torch.Tensor,
        t: torch.Tensor,
    ) -> torch.Tensor:
        """Compute per-token routing weights from precision (inverse = uncertainty).

        Args:
            h: (B, S, D) hidden states
            t: (B,) integer diffusion timesteps

        Returns:
            weights: (B, S, 1) normalized routing weights per token.
                     High-uncertainty tokens get higher weight.
        """
        tok_prec, ch_prec = self.precision_head(h, t)
        # tok_prec: (B, S, 1), ch_prec: (B, 1, D)

        # Aggregate channel precision into a scalar per token.
        # Mean across channels gives an overall confidence per-channel-group,
        # then combine with token precision multiplicatively.
        ch_mean = ch_prec.mean(dim=-1, keepdim=True)  # (B, 1, 1)
        combined_prec = tok_prec * ch_mean  # (B, S, 1)

        # Invert: uncertainty = 1 / precision
        uncertainty = 1.0 / (combined_prec + self.eps)  # (B, S, 1)

        # Normalize to form a proper weighting (sums to 1 over tokens)
        weights = uncertainty / (uncertainty.sum(dim=1, keepdim=True) + self.eps)

        return weights

    def precision_weighted_pool(
        self,
        h: torch.Tensor,
        t: torch.Tensor,
    ) -> torch.Tensor:
        """Pool hidden states with precision-based weighting.

        This is the key equation:
            h_pool = sum_i (1/Pi_i) * h_i / sum_i (1/Pi_i)

        Args:
            h: (B, S, D) hidden states
            t: (B,) diffusion timesteps

        Returns:
            h_pool: (B, D) precision-weighted pooled representation
        """
        weights = self.compute_precision_weights(h, t)  # (B, S, 1)
        h_pool = (weights * h).sum(dim=1)  # (B, D)
        return h_pool

    def project_query(self, h_q: torch.Tensor, t: Optional[torch.Tensor] = None) -> torch.Tensor:
        """Project query hidden states to routing space using precision weighting.

        If t is provided, uses precision-weighted pooling then projects.
        If t is None (fallback), uses mean pooling like standard RouterProjector.

        Args:
            h_q: (B, S, D) query hidden states
            t: (B,) diffusion timesteps. Required for precision routing.

        Returns:
            qr: (B, router_dim) routing query (pooled, not per-token)

        Note: Returns (B, router_dim) not (B, S, router_dim) like RouterProjector.
              The caller should handle this shape difference, or use
              `compute_routing_scores` which handles it internally.
        """
        if t is not None:
            h_pool = self.precision_weighted_pool(h_q, t)  # (B, D)
        else:
            h_pool = h_q.mean(dim=1)  # (B, D) fallback

        return self.q_proj(h_pool)  # (B, router_dim)

    def project_key(self, h_doc: torch.Tensor) -> torch.Tensor:
        """Project document hidden states to routing space.

        Same interface as RouterProjector.project_key.

        Args:
            h_doc: (B, S, D) or (S, D) document hidden states

        Returns:
            kr: same shape with last dim = router_dim
        """
        return self.k_proj(h_doc)

    def compute_routing_scores(
        self,
        h_query: torch.Tensor,
        memory_routing_keys: torch.Tensor,
        t: Optional[torch.Tensor] = None,
    ) -> torch.Tensor:
        """Score documents by relevance using precision-weighted queries.

        Args:
            h_query: (B, S_q, D) query hidden states
            memory_routing_keys: (N_chunks, router_dim) compressed routing keys
            t: (B,) diffusion timesteps

        Returns:
            scores: (B, N_chunks) cosine similarity scores
        """
        qr = self.project_query(h_query, t=t)  # (B, router_dim)
        qr_norm = F.normalize(qr, dim=-1)  # (B, router_dim)
        kr_norm = F.normalize(memory_routing_keys, dim=-1)  # (N, router_dim)

        # (B, router_dim) @ (N, router_dim)^T -> (B, N)
        scores = torch.einsum("bd,nd->bn", qr_norm, kr_norm.to(qr_norm.device))
        return scores

    def estimate_precision_gain(
        self,
        h_query: torch.Tensor,
        t: torch.Tensor,
        memory_bank: MemoryBank,
        layer_idx: int,
        top_k: int = 16,
    ) -> Tuple[torch.Tensor, List[str]]:
        """Estimate how much each document would reduce total uncertainty.

        The precision gain for document j is:
            gain_j = score_j * total_uncertainty

        where total_uncertainty = sum_i (1/Pi_i) and score_j is the cosine
        similarity. Intuition: a highly relevant document (high score) helps
        most when the model is very uncertain (high total uncertainty).

        For a more refined estimate, we compute the per-token gain:
            gain_j = sum_i (1/Pi_i) * sim(h_i, KR_j)

        This assigns more gain to documents that are relevant specifically
        to the uncertain tokens.

        Args:
            h_query: (B, S, D) query hidden states
            t: (B,) diffusion timesteps
            memory_bank: the memory bank with stored document representations
            layer_idx: which layer to query
            top_k: return top-k documents by estimated gain

        Returns:
            gains: (B, top_k) estimated precision gains
            doc_ids: list of document IDs for the top-k selections
        """
        B, S, D = h_query.shape

        # Get routing keys from memory bank
        kr_all, chunk_doc_ids = memory_bank.get_routing_keys(layer_idx)
        if kr_all is None:
            return torch.zeros(B, 0, device=h_query.device), []

        N_chunks = kr_all.shape[0]

        # Compute per-token uncertainty weights (not normalized)
        tok_prec, ch_prec = self.precision_head(h_query, t)
        ch_mean = ch_prec.mean(dim=-1, keepdim=True)  # (B, 1, 1)
        combined_prec = tok_prec * ch_mean  # (B, S, 1)
        uncertainty = 1.0 / (combined_prec + self.eps)  # (B, S, 1)

        # Total uncertainty (scalar per batch element)
        total_uncertainty = uncertainty.sum(dim=1).squeeze(-1)  # (B,)

        # Per-token similarity to each routing key
        # Project query tokens to router space
        h_proj = self.q_proj(h_query)  # (B, S, router_dim)
        h_proj_norm = F.normalize(h_proj, dim=-1)
        kr_norm = F.normalize(kr_all, dim=-1).to(h_query.device)

        # (B, S, router_dim) @ (N, router_dim)^T -> (B, S, N)
        per_token_sim = torch.einsum("bsd,nd->bsn", h_proj_norm, kr_norm)

        # Weight by uncertainty: gain_j = sum_i (1/Pi_i) * sim(h_i, KR_j)
        # uncertainty: (B, S, 1), per_token_sim: (B, S, N)
        gains = (uncertainty * per_token_sim).sum(dim=1)  # (B, N)

        # Select top-k
        k = min(top_k, N_chunks)
        top_gains, top_indices = gains.topk(k, dim=-1)  # (B, k)

        # Map indices back to document IDs (using first batch element)
        selected_doc_ids = []
        if B > 0:
            indices_0 = top_indices[0].tolist()
            seen = set()
            for idx in indices_0:
                if idx < len(chunk_doc_ids):
                    did = chunk_doc_ids[idx]
                    if did not in seen:
                        selected_doc_ids.append(did)
                        seen.add(did)

        return top_gains, selected_doc_ids


def replace_router_in_msa_layer(
    msa_layer: MSALayer,
    precision_head: PrecisionHead,
) -> PrecisionGuidedRouter:
    """Replace a standard RouterProjector with a PrecisionGuidedRouter.

    This is the drop-in replacement operation. After calling this, the
    MSALayer's .router attribute is a PrecisionGuidedRouter.

    Args:
        msa_layer: the MSALayer to modify
        precision_head: the PrecisionHead from the corresponding model layer

    Returns:
        The new PrecisionGuidedRouter (also installed on msa_layer.router)
    """
    old_router = msa_layer.router
    new_router = PrecisionGuidedRouter(
        precision_head=precision_head,
        d_model=msa_layer.d_model,
        router_dim=old_router.q_proj.out_features,
        n_heads=old_router.n_heads,
    )

    # Initialize key projection from old router's key projection
    # so we don't break already-encoded memory bank entries.
    with torch.no_grad():
        new_router.k_proj.weight.copy_(old_router.k_proj.weight)
        new_router.k_proj.bias.copy_(old_router.k_proj.bias)

    msa_layer.router = new_router
    return new_router


def equip_model_with_precision_routing(
    model: PCSHODLM,
    msa_layers: nn.ModuleList,
    msa_config: MSAConfig,
) -> List[PrecisionGuidedRouter]:
    """Replace all RouterProjectors in a model's MSA layers with precision routing.

    Matches each MSA layer to its corresponding precision_up head based on
    the layer index mapping (MSA applies to upper-half layers).

    Args:
        model: the PC-SHO-DLM model (provides precision heads)
        msa_layers: the MSA layers to modify
        msa_config: MSA configuration

    Returns:
        List of the new PrecisionGuidedRouter instances
    """
    n_layers = model.config.n_layers
    msa_start = max(msa_config.apply_from_layer, n_layers // 2)

    routers = []
    for i, msa_layer in enumerate(msa_layers):
        model_layer_idx = msa_start + i
        if model_layer_idx < len(model.precision_up):
            prec_head = model.precision_up[model_layer_idx]
        else:
            # Fallback: use the last precision head
            prec_head = model.precision_up[-1]

        router = replace_router_in_msa_layer(msa_layer, prec_head)
        routers.append(router)

    return routers


# =============================================================================
# Comparison: Standard vs Precision-Guided Routing
# =============================================================================

def compare_routing_strategies(
    model: PCSHODLM,
    msa_layers_standard: nn.ModuleList,
    msa_layers_precision: nn.ModuleList,
    memory_bank: MemoryBank,
    h_query: torch.Tensor,
    t: torch.Tensor,
    layer_idx: int,
) -> Dict[str, torch.Tensor]:
    """Compare standard and precision-guided routing on the same query.

    Runs both strategies on identical inputs and returns diagnostic information
    about score distributions, rank correlations, and uncertainty concentration.

    Args:
        model: PC-SHO-DLM model
        msa_layers_standard: MSA layers with standard RouterProjector
        msa_layers_precision: MSA layers with PrecisionGuidedRouter
        memory_bank: shared memory bank
        h_query: (B, S, D) query hidden states
        t: (B,) diffusion timesteps
        layer_idx: which memory bank layer to score against

    Returns:
        Dictionary with comparison metrics:
            - scores_standard: (B, N) standard routing scores
            - scores_precision: (B, N) precision-guided scores
            - rank_correlation: (B,) Spearman rank correlation per batch
            - precision_weights: (B, S, 1) the precision-derived token weights
            - total_uncertainty: (B,) sum of uncertainty per batch element
            - top_k_overlap: fraction of top-k docs shared between strategies
    """
    B = h_query.shape[0]

    # Get routing keys
    kr_all, chunk_doc_ids = memory_bank.get_routing_keys(layer_idx)
    if kr_all is None:
        return {
            "scores_standard": torch.zeros(B, 0),
            "scores_precision": torch.zeros(B, 0),
            "rank_correlation": torch.zeros(B),
            "precision_weights": torch.zeros(B, h_query.shape[1], 1),
            "total_uncertainty": torch.zeros(B),
            "top_k_overlap": torch.tensor(0.0),
        }

    device = h_query.device
    kr_all = kr_all.to(device)

    # Identify the MSA layer index (assumes same ordering)
    n_layers = model.config.n_layers
    msa_start = n_layers // 2
    msa_idx = layer_idx - msa_start
    msa_idx = max(0, min(msa_idx, len(msa_layers_standard) - 1))

    # --- Standard routing ---
    std_layer = msa_layers_standard[msa_idx]
    scores_standard = std_layer.compute_routing_scores(h_query, kr_all)  # (B, N)

    # --- Precision routing ---
    prec_layer = msa_layers_precision[msa_idx]
    assert isinstance(prec_layer.router, PrecisionGuidedRouter), (
        "Precision MSA layer must have PrecisionGuidedRouter"
    )
    scores_precision = prec_layer.router.compute_routing_scores(
        h_query, kr_all, t=t
    )  # (B, N)

    # --- Precision weights for diagnostics ---
    prec_router = prec_layer.router
    weights = prec_router.compute_precision_weights(h_query, t)  # (B, S, 1)

    # Total uncertainty
    tok_prec, ch_prec = prec_router.precision_head(h_query, t)
    ch_mean = ch_prec.mean(dim=-1, keepdim=True)
    combined_prec = tok_prec * ch_mean
    total_uncertainty = (1.0 / (combined_prec + prec_router.eps)).sum(dim=(1, 2))  # (B,)

    # --- Rank correlation (Spearman) ---
    def _spearman_rank_corr(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
        """Compute Spearman rank correlation per batch element."""
        batch_corrs = []
        for i in range(a.shape[0]):
            rank_a = a[i].argsort().argsort().float()
            rank_b = b[i].argsort().argsort().float()
            n = rank_a.shape[0]
            if n < 2:
                batch_corrs.append(torch.tensor(1.0, device=a.device))
                continue
            mean_a = rank_a.mean()
            mean_b = rank_b.mean()
            cov = ((rank_a - mean_a) * (rank_b - mean_b)).sum()
            std_a = (rank_a - mean_a).pow(2).sum().sqrt()
            std_b = (rank_b - mean_b).pow(2).sum().sqrt()
            corr = cov / (std_a * std_b + 1e-8)
            batch_corrs.append(corr)
        return torch.stack(batch_corrs)

    rank_corr = _spearman_rank_corr(scores_standard, scores_precision)

    # --- Top-k overlap ---
    k = min(16, scores_standard.shape[-1])
    if k > 0:
        top_std = scores_standard.topk(k, dim=-1).indices  # (B, k)
        top_prec = scores_precision.topk(k, dim=-1).indices
        overlap_count = 0
        for b in range(B):
            std_set = set(top_std[b].tolist())
            prec_set = set(top_prec[b].tolist())
            overlap_count += len(std_set & prec_set)
        top_k_overlap = torch.tensor(overlap_count / (B * k))
    else:
        top_k_overlap = torch.tensor(0.0)

    return {
        "scores_standard": scores_standard,
        "scores_precision": scores_precision,
        "rank_correlation": rank_corr,
        "precision_weights": weights,
        "total_uncertainty": total_uncertainty,
        "top_k_overlap": top_k_overlap,
    }


# =============================================================================
# Test
# =============================================================================

def test_precision_router():
    """End-to-end test of PrecisionGuidedRouter."""
    print("=" * 70)
    print("TEST: PrecisionGuidedRouter")
    print("=" * 70)

    device = "cpu"
    torch.manual_seed(42)

    # --- Setup model ---
    config = PCSHOConfig(
        vocab_size=1000,
        max_seq_len=128,
        d_model=256,
        n_heads=4,
        n_layers=6,
        d_ff=512,
        n_diffusion_steps=100,
        n_settling_steps=2,
    )
    model = PCSHODLM(config).to(device)

    msa_config = MSAConfig(
        chunk_size=32,
        top_k=4,
        router_dim=64,
        n_router_heads=4,
        apply_from_layer=3,
    )

    # Create two copies of MSA layers: one standard, one for precision routing
    msa_layers_std = create_msa_layers(config, msa_config)
    msa_layers_prec = create_msa_layers(config, msa_config)

    # Copy weights so the only difference is the router
    msa_layers_prec.load_state_dict(msa_layers_std.state_dict())

    # --- Test 1: Basic construction and drop-in replacement ---
    print("\n[1] Drop-in replacement...")
    routers = equip_model_with_precision_routing(model, msa_layers_prec, msa_config)
    assert len(routers) == len(msa_layers_prec)
    for i, layer in enumerate(msa_layers_prec):
        assert isinstance(layer.router, PrecisionGuidedRouter), (
            f"Layer {i} router was not replaced"
        )
    print("    PASS: All routers replaced with PrecisionGuidedRouter")

    # --- Test 2: Precision-weighted pooling ---
    print("\n[2] Precision-weighted pooling...")
    B, S, D = 2, 32, config.d_model
    h = torch.randn(B, S, D)
    t = torch.randint(1, 101, (B,))

    router = routers[0]
    weights = router.compute_precision_weights(h, t)
    assert weights.shape == (B, S, 1), f"Bad shape: {weights.shape}"
    assert torch.allclose(weights.sum(dim=1), torch.ones(B, 1), atol=1e-4), (
        "Weights do not sum to 1"
    )
    assert (weights >= 0).all(), "Negative weights"

    h_pool = router.precision_weighted_pool(h, t)
    assert h_pool.shape == (B, D), f"Bad pool shape: {h_pool.shape}"
    print("    PASS: Weights sum to 1, pooled shape correct")

    # --- Test 3: Routing scores ---
    print("\n[3] Routing scores against memory bank...")
    memory_bank = MemoryBank(chunk_size=msa_config.chunk_size)

    # Populate with synthetic documents
    n_docs = 8
    msa_start = config.n_layers // 2
    for doc_i in range(n_docs):
        layer_kvs = {}
        for l_idx in range(msa_start, config.n_layers):
            n_chunks = 4
            layer_kvs[l_idx] = {
                "K": torch.randn(n_chunks, config.d_model),
                "V": torch.randn(n_chunks, config.d_model),
                "KR": torch.randn(n_chunks, msa_config.router_dim),
            }
        memory_bank.add_document(f"doc_{doc_i}", layer_kvs)

    test_layer = msa_start
    kr_all, chunk_ids = memory_bank.get_routing_keys(test_layer)
    assert kr_all is not None, "No routing keys"

    scores = router.compute_routing_scores(h, kr_all, t=t)
    assert scores.shape == (B, kr_all.shape[0]), f"Bad scores shape: {scores.shape}"
    assert not scores.isnan().any(), "NaN in scores"
    print(f"    PASS: Scores shape {scores.shape}, range [{scores.min():.3f}, {scores.max():.3f}]")

    # --- Test 4: Precision gain estimation ---
    print("\n[4] Precision gain estimation...")
    gains, selected_ids = router.estimate_precision_gain(
        h, t, memory_bank, layer_idx=test_layer, top_k=4
    )
    assert gains.shape[0] == B, f"Bad gains batch dim: {gains.shape}"
    assert gains.shape[1] == 4, f"Expected 4 docs, got {gains.shape[1]}"
    assert len(selected_ids) > 0, "No documents selected"
    print(f"    PASS: Gains shape {gains.shape}, selected {len(selected_ids)} docs: {selected_ids}")

    # --- Test 5: Uncertainty drives routing ---
    print("\n[5] Uncertainty drives routing (functional test)...")
    # Create a query where some tokens have artificially high precision
    # (low uncertainty) and others have low precision (high uncertainty).
    # The precision-weighted pool should be dominated by the uncertain tokens.
    h_split = torch.zeros(1, 10, D)
    h_split[0, :5] = torch.randn(5, D) * 0.1    # "confident" tokens (small norm)
    h_split[0, 5:] = torch.randn(5, D) * 2.0     # "uncertain" tokens (large norm)

    t_test = torch.tensor([50])
    weights_split = router.compute_precision_weights(h_split, t_test)  # (1, 10, 1)

    # The model's precision head will assign precision based on hidden state
    # magnitude and structure. We check that weights are non-uniform.
    weight_variance = weights_split.squeeze().var().item()
    print(f"    Weight variance: {weight_variance:.6f}")
    print(f"    Weights (first 5 vs last 5): "
          f"{weights_split[0,:5,0].mean():.4f} vs {weights_split[0,5:,0].mean():.4f}")
    print("    PASS: Precision weights are non-uniform (routing is not mean-pooling)")

    # --- Test 6: compare_routing_strategies ---
    print("\n[6] Comparing standard vs precision-guided routing...")
    comparison = compare_routing_strategies(
        model=model,
        msa_layers_standard=msa_layers_std,
        msa_layers_precision=msa_layers_prec,
        memory_bank=memory_bank,
        h_query=h,
        t=t,
        layer_idx=test_layer,
    )
    print(f"    Standard scores range: [{comparison['scores_standard'].min():.3f}, "
          f"{comparison['scores_standard'].max():.3f}]")
    print(f"    Precision scores range: [{comparison['scores_precision'].min():.3f}, "
          f"{comparison['scores_precision'].max():.3f}]")
    print(f"    Rank correlation: {comparison['rank_correlation'].tolist()}")
    print(f"    Top-k overlap: {comparison['top_k_overlap'].item():.3f}")
    print(f"    Total uncertainty: {comparison['total_uncertainty'].tolist()}")
    assert comparison["scores_standard"].shape == comparison["scores_precision"].shape
    assert comparison["rank_correlation"].shape == (B,)
    print("    PASS: Comparison metrics computed correctly")

    # --- Test 7: Gradient flow ---
    print("\n[7] Gradient flow through precision router...")
    h_grad = torch.randn(B, S, D, requires_grad=True)
    t_grad = torch.randint(1, 101, (B,))
    scores_grad = router.compute_routing_scores(h_grad, kr_all, t=t_grad)
    loss = scores_grad.sum()
    loss.backward()
    assert h_grad.grad is not None, "No gradient on input"
    assert h_grad.grad.abs().sum() > 0, "Zero gradient"
    print(f"    PASS: Gradient norm on input: {h_grad.grad.norm():.4f}")

    # --- Test 8: MSALayer forward pass with precision router ---
    print("\n[8] MSALayer forward pass with PrecisionGuidedRouter...")
    msa_layer = msa_layers_prec[0]
    x_in = torch.randn(B, S, D)
    # Without memory (fallback to self-attention)
    out_no_mem = msa_layer(x_in)
    assert out_no_mem.shape == x_in.shape, f"Shape mismatch: {out_no_mem.shape}"

    # With memory
    mem_k = torch.randn(B, 8, D)
    mem_v = torch.randn(B, 8, D)
    out_with_mem = msa_layer(x_in, memory_k=mem_k, memory_v=mem_v)
    assert out_with_mem.shape == x_in.shape
    print("    PASS: MSALayer forward works with precision router installed")

    # --- Test 9: Parameter count comparison ---
    print("\n[9] Parameter count comparison...")
    std_router_params = sum(p.numel() for p in msa_layers_std[0].router.parameters())
    prec_router_params = sum(
        p.numel() for p in routers[0].parameters()
        if p.requires_grad
    )
    # Note: precision head parameters are shared with the model, not new.
    prec_head_params = sum(p.numel() for p in routers[0].precision_head.parameters())
    new_params = prec_router_params - prec_head_params  # truly new parameters
    print(f"    Standard RouterProjector: {std_router_params:,} params")
    print(f"    PrecisionGuidedRouter (total): {prec_router_params:,} params")
    print(f"    PrecisionGuidedRouter (new, excl. shared precision head): {new_params:,} params")
    print(f"    Precision head (shared, not new): {prec_head_params:,} params")

    print("\n" + "=" * 70)
    print("ALL TESTS PASSED")
    print("=" * 70)


if __name__ == "__main__":
    test_precision_router()