code/training/model.py

"""
Encoder-only Transformer for predicting circuit summary statistics.

Architecture:
  1. Each scalar input feature → learned linear projection to d_model
  2. Add learned positional + type embeddings (ACh token gets special type)
  3. Prepend a [CLS] aggregation token
  4. Transformer encoder (N layers, multi-head self-attention)
  5. [CLS] output → MLP head → 11 predicted statistics

Two configurations:
  - Model A: 10 input tokens (no ACh), predicts 11 stats
  - Model B: 11 input tokens (with ACh), predicts 11 stats
"""

from __future__ import annotations

import math

import torch
import torch.nn as nn


class FeatureTokenizer(nn.Module):
    """Project each scalar feature to d_model via per-feature linear layers.

    Input:  (B, n_features)  — raw normalized scalars
    Output: (B, n_features, d_model) — token embeddings
    """

    def __init__(self, n_features: int, d_model: int):
        super().__init__()
        self.projections = nn.ModuleList([
            nn.Linear(1, d_model) for _ in range(n_features)
        ])

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # x: (B, n_features)
        tokens = []
        for i, proj in enumerate(self.projections):
            tokens.append(proj(x[:, i : i + 1]))  # (B, 1) → (B, d_model)
        return torch.stack(tokens, dim=1)  # (B, n_features, d_model)


class CircuitTransformer(nn.Module):
    """Encoder-only transformer for circuit statistics prediction.

    Args:
        n_features: Number of input features (10 for Model A, 11 for Model B)
        n_outputs: Number of output statistics (11)
        d_model: Embedding dimension
        n_heads: Number of attention heads
        n_layers: Number of transformer encoder layers
        d_ff: Feed-forward hidden dimension
        dropout: Dropout rate
        has_ach: Whether ACh is included (for type embedding)
    """

    def __init__(
        self,
        n_features: int,
        n_outputs: int = 11,
        d_model: int = 64,
        n_heads: int = 4,
        n_layers: int = 4,
        d_ff: int = 256,
        dropout: float = 0.1,
        has_ach: bool = True,
    ):
        super().__init__()
        self.n_features = n_features
        self.n_outputs = n_outputs
        self.d_model = d_model
        self.has_ach = has_ach

        # ── Input embedding ──────────────────────────────────────────────
        self.tokenizer = FeatureTokenizer(n_features, d_model)

        # Learned [CLS] token
        self.cls_token = nn.Parameter(torch.randn(1, 1, d_model) * 0.02)

        # Positional embedding: n_features + 1 (for [CLS])
        self.pos_embed = nn.Parameter(
            torch.randn(1, n_features + 1, d_model) * 0.02
        )

        # Type embedding: 0 = structural param, 1 = ACh token, 2 = CLS
        self.type_embed = nn.Embedding(3, d_model)

        self.embed_dropout = nn.Dropout(dropout)
        self.embed_norm = nn.LayerNorm(d_model)

        # ── Transformer encoder ──────────────────────────────────────────
        encoder_layer = nn.TransformerEncoderLayer(
            d_model=d_model,
            nhead=n_heads,
            dim_feedforward=d_ff,
            dropout=dropout,
            activation="gelu",
            batch_first=True,
            norm_first=True,  # Pre-norm for training stability
        )
        self.encoder = nn.TransformerEncoder(
            encoder_layer, num_layers=n_layers
        )

        # ── Output head ──────────────────────────────────────────────────
        self.output_norm = nn.LayerNorm(d_model)
        self.output_head = nn.Sequential(
            nn.Linear(d_model, d_ff),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(d_ff, n_outputs),
        )

        # Initialize weights
        self.apply(self._init_weights)

    def _init_weights(self, module: nn.Module):
        if isinstance(module, nn.Linear):
            nn.init.trunc_normal_(module.weight, std=0.02)
            if module.bias is not None:
                nn.init.zeros_(module.bias)
        elif isinstance(module, nn.Embedding):
            nn.init.trunc_normal_(module.weight, std=0.02)
        elif isinstance(module, nn.LayerNorm):
            nn.init.ones_(module.weight)
            nn.init.zeros_(module.bias)

    def _build_type_ids(self, batch_size: int, device: torch.device) -> torch.Tensor:
        """Build type IDs: [CLS]=2, structural=0, ACh=1."""
        # Sequence: [CLS, feat_0, feat_1, ..., feat_{n-1}]
        type_ids = torch.zeros(
            batch_size, self.n_features + 1, dtype=torch.long, device=device
        )
        type_ids[:, 0] = 2  # CLS token

        if self.has_ach:
            # Last feature position is ACh
            type_ids[:, -1] = 1  # ACh token type

        return type_ids

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Args:
            x: (B, n_features) — normalized input features

        Returns:
            (B, n_outputs) — predicted statistics (in normalized space)
        """
        B = x.shape[0]
        device = x.device

        # 1. Tokenize features → (B, n_features, d_model)
        tokens = self.tokenizer(x)

        # 2. Prepend [CLS] → (B, n_features+1, d_model)
        cls_expanded = self.cls_token.expand(B, -1, -1)
        tokens = torch.cat([cls_expanded, tokens], dim=1)

        # 3. Add positional + type embeddings
        type_ids = self._build_type_ids(B, device)
        tokens = tokens + self.pos_embed + self.type_embed(type_ids)

        # 4. Norm + dropout
        tokens = self.embed_norm(tokens)
        tokens = self.embed_dropout(tokens)

        # 5. Transformer encoder
        tokens = self.encoder(tokens)

        # 6. Extract [CLS] representation → predict
        cls_out = tokens[:, 0]  # (B, d_model)
        cls_out = self.output_norm(cls_out)
        return self.output_head(cls_out)  # (B, n_outputs)

    def count_params(self) -> int:
        return sum(p.numel() for p in self.parameters() if p.requires_grad)


# ── MLP Baseline ─────────────────────────────────────────────────────────────

class CircuitMLP(nn.Module):
    """Simple MLP baseline for tabular regression.

    Properly sized for small datasets (~5K-55K samples).
    Default: 2 layers × 64 units = ~5K-10K params.
    """

    def __init__(
        self,
        n_features: int,
        n_outputs: int = 11,
        hidden_dims: list[int] | None = None,
        dropout: float = 0.1,
    ):
        super().__init__()
        self.n_features = n_features
        self.n_outputs = n_outputs
        hidden_dims = hidden_dims or [64, 64]

        layers = []
        in_dim = n_features
        for h_dim in hidden_dims:
            layers.extend([
                nn.Linear(in_dim, h_dim),
                nn.GELU(),
                nn.Dropout(dropout),
            ])
            in_dim = h_dim
        layers.append(nn.Linear(in_dim, n_outputs))

        self.net = nn.Sequential(*layers)
        self.apply(self._init_weights)

    def _init_weights(self, module: nn.Module):
        if isinstance(module, nn.Linear):
            nn.init.kaiming_normal_(module.weight, nonlinearity="linear")
            if module.bias is not None:
                nn.init.zeros_(module.bias)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.net(x)

    def count_params(self) -> int:
        return sum(p.numel() for p in self.parameters() if p.requires_grad)


# ── Factory functions ────────────────────────────────────────────────────────

def build_model_a(cfg, arch: str = "transformer") -> nn.Module:
    """Model A: plain HH (no ACh token)."""
    if arch == "mlp":
        return CircuitMLP(
            n_features=cfg.n_input_features_a,
            n_outputs=cfg.n_output_stats,
            hidden_dims=cfg.mlp_hidden,
            dropout=cfg.mlp_dropout,
        )
    return CircuitTransformer(
        n_features=cfg.n_input_features_a,
        n_outputs=cfg.n_output_stats,
        d_model=cfg.d_model,
        n_heads=cfg.n_heads,
        n_layers=cfg.n_layers,
        d_ff=cfg.d_ff,
        dropout=cfg.dropout,
        has_ach=False,
    )


def build_model_b(cfg, arch: str = "transformer") -> nn.Module:
    """Model B: HH + ACh modulation."""
    if arch == "mlp":
        return CircuitMLP(
            n_features=cfg.n_input_features_b,
            n_outputs=cfg.n_output_stats,
            hidden_dims=cfg.mlp_hidden,
            dropout=cfg.mlp_dropout,
        )
    return CircuitTransformer(
        n_features=cfg.n_input_features_b,
        n_outputs=cfg.n_output_stats,
        d_model=cfg.d_model,
        n_heads=cfg.n_heads,
        n_layers=cfg.n_layers,
        d_ff=cfg.d_ff,
        dropout=cfg.dropout,
        has_ach=True,
    )