#!/usr/bin/env python3
"""
helix_cdc/tiny_transformer_cdna.py

Custom TinyTransformer with explicit Linear layers that can be swapped to CDNALinear.

Unlike nn.TransformerEncoderLayer, this doesn't access .weight directly,
allowing clean CDNALinear substitution.
"""

import math
from typing import Optional

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


class SimpleAttention(nn.Module):
    """
    Simple multi-head attention with explicit Linear layers.

    No internal .weight access - fully compatible with CDNALinear swap.
    """

    def __init__(self, dim: int, n_heads: int, dropout: float = 0.0):
        super().__init__()
        self.dim = dim
        self.n_heads = n_heads
        self.head_dim = dim // n_heads
        self.scale = self.head_dim ** -0.5

        # Explicit Q, K, V projections
        self.q_proj = nn.Linear(dim, dim, bias=False)
        self.k_proj = nn.Linear(dim, dim, bias=False)
        self.v_proj = nn.Linear(dim, dim, bias=False)
        self.out_proj = nn.Linear(dim, dim, bias=False)

        self.dropout = nn.Dropout(dropout)

    def forward(self, x: torch.Tensor, mask: Optional[torch.Tensor] = None) -> torch.Tensor:
        """
        Args:
            x: [batch, seq, dim]
            mask: Optional attention mask

        Returns:
            out: [batch, seq, dim]
        """
        B, T, D = x.shape

        # Project Q, K, V
        q = self.q_proj(x)  # [B, T, D]
        k = self.k_proj(x)
        v = self.v_proj(x)

        # Reshape for multi-head attention
        q = q.view(B, T, self.n_heads, self.head_dim).transpose(1, 2)  # [B, H, T, head_dim]
        k = k.view(B, T, self.n_heads, self.head_dim).transpose(1, 2)
        v = v.view(B, T, self.n_heads, self.head_dim).transpose(1, 2)

        # Attention scores
        scores = torch.matmul(q, k.transpose(-2, -1)) * self.scale  # [B, H, T, T]

        if mask is not None:
            scores = scores.masked_fill(mask == 0, float('-inf'))

        attn = F.softmax(scores, dim=-1)
        attn = self.dropout(attn)

        # Apply attention to values
        out = torch.matmul(attn, v)  # [B, H, T, head_dim]

        # Reshape back
        out = out.transpose(1, 2).contiguous().view(B, T, D)  # [B, T, D]

        # Output projection
        out = self.out_proj(out)

        return out


class SimpleMLP(nn.Module):
    """Simple 2-layer MLP with explicit Linear layers."""

    def __init__(self, dim: int, hidden_dim: int, dropout: float = 0.0):
        super().__init__()
        self.fc1 = nn.Linear(dim, hidden_dim)
        self.fc2 = nn.Linear(hidden_dim, dim)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = self.fc1(x)
        x = F.gelu(x)
        x = self.dropout(x)
        x = self.fc2(x)
        return x


class SimpleTransformerLayer(nn.Module):
    """Transformer layer with explicit modules (no .weight access)."""

    def __init__(self, dim: int, n_heads: int, ffn_dim: int, dropout: float = 0.0):
        super().__init__()
        self.attn = SimpleAttention(dim, n_heads, dropout)
        self.mlp = SimpleMLP(dim, ffn_dim, dropout)
        self.norm1 = nn.LayerNorm(dim)
        self.norm2 = nn.LayerNorm(dim)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # Pre-norm attention
        x = x + self.dropout(self.attn(self.norm1(x)))
        # Pre-norm MLP
        x = x + self.dropout(self.mlp(self.norm2(x)))
        return x


class TinyTransformerSimple(nn.Module):
    """
    Minimal transformer with explicit Linear layers.

    Compatible with CDNALinear swap (no internal .weight access).
    """

    def __init__(
        self,
        vocab_size: int = 1000,
        dim: int = 256,
        n_heads: int = 4,
        n_layers: int = 2,
        ffn_dim: int = 512,
        dropout: float = 0.0,
    ):
        super().__init__()
        self.vocab_size = vocab_size
        self.dim = dim
        self.n_heads = n_heads
        self.n_layers = n_layers

        # Embedding
        self.embed = nn.Embedding(vocab_size, dim)

        # Transformer layers
        self.layers = nn.ModuleList([
            SimpleTransformerLayer(dim, n_heads, ffn_dim, dropout)
            for _ in range(n_layers)
        ])

        # Final norm
        self.norm = nn.LayerNorm(dim)

        # Output projection
        self.output = nn.Linear(dim, vocab_size)

        # Initialize
        self._init_weights()

    def _init_weights(self):
        """Initialize weights deterministically."""
        torch.manual_seed(42)
        for m in self.modules():
            if isinstance(m, nn.Linear):
                nn.init.xavier_uniform_(m.weight, gain=0.1)
                if m.bias is not None:
                    nn.init.zeros_(m.bias)
            elif isinstance(m, nn.Embedding):
                nn.init.normal_(m.weight, std=0.02)
            elif isinstance(m, nn.LayerNorm):
                nn.init.ones_(m.weight)
                nn.init.zeros_(m.bias)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Args:
            x: [batch, seq_len] token IDs

        Returns:
            logits: [batch, seq_len, vocab_size]
        """
        # Embed
        h = self.embed(x)  # [B, T, D]

        # Transformer layers
        for layer in self.layers:
            h = layer(h)

        # Final norm
        h = self.norm(h)

        # Output projection
        logits = self.output(h)

        return logits

    def count_linear_layers(self) -> int:
        """Count number of nn.Linear layers."""
        return sum(1 for m in self.modules() if isinstance(m, nn.Linear))


def count_parameters(model: nn.Module) -> int:
    """Count total parameters."""
    return sum(p.numel() for p in model.parameters())
