"""Game Transformer — Decoder-only transformer for Essence Wars.

GPT-style architecture with game-specific output heads:
- Policy head: action classification over 422 actions
- Value head: scalar state value prediction
- Action-value head: per-action value prediction over 422 actions

Architecture follows LLaMA/Searchless Chess:
- Pre-normalization with RMSNorm
- SwiGLU activation in FFN
- Rotary Position Embeddings (RoPE)
- No dropout (regularize via data)
- Causal attention (decoder-only)

Model scales:
- EW-Small:  ~10M params (8 layers,  8 heads, dim=256, ffn=1024)
- EW-Medium: ~50M params (12 layers, 8 heads, dim=512, ffn=2048)
- EW-Large: ~150M params (16 layers, 12 heads, dim=768, ffn=3072)
"""

from __future__ import annotations

import math
from dataclasses import dataclass

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

ACTION_SPACE_SIZE = 422


# ============================================================================
# Configuration
# ============================================================================

@dataclass
class TransformerConfig:
    """Configuration for the Game Transformer model."""

    vocab_size: int = 631
    max_seq_len: int = 128
    n_layers: int = 8
    n_heads: int = 8
    embed_dim: int = 256
    ffn_dim: int = 1024
    action_space: int = ACTION_SPACE_SIZE
    rope_theta: float = 10000.0

    @classmethod
    def small(cls, vocab_size: int = 631) -> TransformerConfig:
        """EW-Small: ~10M params."""
        return cls(
            vocab_size=vocab_size,
            n_layers=8,
            n_heads=8,
            embed_dim=256,
            ffn_dim=1024,
            max_seq_len=128,
        )

    @classmethod
    def medium(cls, vocab_size: int = 631) -> TransformerConfig:
        """EW-Medium: ~50M params."""
        return cls(
            vocab_size=vocab_size,
            n_layers=12,
            n_heads=8,
            embed_dim=512,
            ffn_dim=2048,
            max_seq_len=256,
        )

    @classmethod
    def large(cls, vocab_size: int = 631) -> TransformerConfig:
        """EW-Large: ~150M params."""
        return cls(
            vocab_size=vocab_size,
            n_layers=16,
            n_heads=12,
            embed_dim=768,
            ffn_dim=3072,
            max_seq_len=512,
        )


# ============================================================================
# Building Blocks
# ============================================================================

class RMSNorm(nn.Module):
    """Root Mean Square Layer Normalization (LLaMA-style)."""

    def __init__(self, dim: int, eps: float = 1e-6) -> None:
        super().__init__()
        self.eps = eps
        self.weight = nn.Parameter(torch.ones(dim))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        rms = torch.rsqrt(x.float().pow(2).mean(-1, keepdim=True) + self.eps)
        return (x.float() * rms).to(x.dtype) * self.weight


def _precompute_rope_freqs(
    dim: int,
    max_seq_len: int,
    theta: float = 10000.0,
    device: torch.device | None = None,
) -> torch.Tensor:
    """Precompute RoPE frequency tensor for complex exponentials.

    Returns shape (max_seq_len, dim//2) complex tensor.
    """
    freqs = 1.0 / (theta ** (torch.arange(0, dim, 2, device=device).float() / dim))
    t = torch.arange(max_seq_len, device=device)
    freqs = torch.outer(t, freqs)
    return torch.polar(torch.ones_like(freqs), freqs)  # complex64


def _apply_rope(
    xq: torch.Tensor,
    xk: torch.Tensor,
    freqs: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Apply RoPE to query and key tensors.

    xq, xk: (batch, seq_len, n_heads, head_dim)
    freqs: (seq_len, head_dim//2) complex
    """
    # Reshape to complex pairs
    xq_complex = torch.view_as_complex(xq.float().reshape(*xq.shape[:-1], -1, 2))
    xk_complex = torch.view_as_complex(xk.float().reshape(*xk.shape[:-1], -1, 2))

    # Broadcast freqs to match batch/head dims
    freqs = freqs[: xq.shape[1]].unsqueeze(0).unsqueeze(2)  # (1, seq, 1, dim//2)

    xq_out = torch.view_as_real(xq_complex * freqs).flatten(-2)
    xk_out = torch.view_as_real(xk_complex * freqs).flatten(-2)

    return xq_out.to(xq.dtype), xk_out.to(xk.dtype)


class CausalSelfAttention(nn.Module):
    """Multi-head self-attention with RoPE and causal mask."""

    def __init__(self, config: TransformerConfig) -> None:
        super().__init__()
        assert config.embed_dim % config.n_heads == 0
        self.n_heads = config.n_heads
        self.head_dim = config.embed_dim // config.n_heads

        self.q_proj = nn.Linear(config.embed_dim, config.embed_dim, bias=False)
        self.k_proj = nn.Linear(config.embed_dim, config.embed_dim, bias=False)
        self.v_proj = nn.Linear(config.embed_dim, config.embed_dim, bias=False)
        self.o_proj = nn.Linear(config.embed_dim, config.embed_dim, bias=False)

    def forward(
        self,
        x: torch.Tensor,
        freqs: torch.Tensor,
        attention_mask: torch.Tensor | None = None,
    ) -> torch.Tensor:
        B, T, C = x.shape

        q = self.q_proj(x).view(B, T, self.n_heads, self.head_dim)
        k = self.k_proj(x).view(B, T, self.n_heads, self.head_dim)
        v = self.v_proj(x).view(B, T, self.n_heads, self.head_dim)

        # Apply RoPE
        q, k = _apply_rope(q, k, freqs)

        # Transpose for attention: (B, n_heads, T, head_dim)
        q = q.transpose(1, 2)
        k = k.transpose(1, 2)
        v = v.transpose(1, 2)

        # Build combined mask: causal + padding
        # Causal mask: lower-triangular
        causal = torch.tril(torch.ones(T, T, device=x.device, dtype=torch.bool))
        causal = causal.unsqueeze(0).unsqueeze(0)  # (1, 1, T, T)

        if attention_mask is not None:
            # attention_mask: (B, T) -> (B, 1, 1, T) for key masking
            pad_mask = attention_mask.unsqueeze(1).unsqueeze(2).bool()  # (B, 1, 1, T)
            mask = causal & pad_mask
        else:
            mask = causal

        # Scaled dot-product attention with mask
        out = F.scaled_dot_product_attention(
            q, k, v, attn_mask=mask, is_causal=False
        )

        # Merge heads
        out = out.transpose(1, 2).contiguous().view(B, T, C)
        return self.o_proj(out)  # type: ignore[no-any-return]


class SwiGLU(nn.Module):
    """SwiGLU feed-forward network (LLaMA-style).

    FFN with gated activation: output = (Swish(xW1) * xV) @ W2
    Uses 2/3 of the target FFN dim for the gate and up projections
    to match parameter count of a standard 2-layer FFN.
    """

    def __init__(self, embed_dim: int, ffn_dim: int) -> None:
        super().__init__()
        # SwiGLU has 3 weight matrices instead of 2, so we scale
        # the intermediate dim to keep param count similar
        hidden = int(2 * ffn_dim / 3)
        # Round to nearest multiple of 8 for efficiency
        hidden = ((hidden + 7) // 8) * 8

        self.gate_proj = nn.Linear(embed_dim, hidden, bias=False)
        self.up_proj = nn.Linear(embed_dim, hidden, bias=False)
        self.down_proj = nn.Linear(hidden, embed_dim, bias=False)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x))  # type: ignore[no-any-return]


class TransformerBlock(nn.Module):
    """Single transformer block with pre-norm, attention, and FFN."""

    def __init__(self, config: TransformerConfig) -> None:
        super().__init__()
        self.attn_norm = RMSNorm(config.embed_dim)
        self.attn = CausalSelfAttention(config)
        self.ffn_norm = RMSNorm(config.embed_dim)
        self.ffn = SwiGLU(config.embed_dim, config.ffn_dim)

    def forward(
        self,
        x: torch.Tensor,
        freqs: torch.Tensor,
        attention_mask: torch.Tensor | None = None,
    ) -> torch.Tensor:
        # Pre-norm + attention + residual
        x = x + self.attn(self.attn_norm(x), freqs, attention_mask)
        # Pre-norm + FFN + residual
        x = x + self.ffn(self.ffn_norm(x))
        return x


# ============================================================================
# Game Transformer Model
# ============================================================================

class GameTransformer(nn.Module):
    """Decoder-only transformer for Essence Wars game play.

    Processes tokenized game state and outputs:
    - Policy logits over 422 actions
    - Scalar state value
    - Per-action value predictions
    """

    def __init__(self, config: TransformerConfig) -> None:
        super().__init__()
        self.config = config

        # Token embedding
        self.tok_emb = nn.Embedding(config.vocab_size, config.embed_dim)

        # Transformer blocks
        self.layers = nn.ModuleList([
            TransformerBlock(config) for _ in range(config.n_layers)
        ])
        self.norm = RMSNorm(config.embed_dim)

        # Output heads
        self.policy_head = nn.Linear(config.embed_dim, config.action_space, bias=False)
        self.value_head = nn.Sequential(
            nn.Linear(config.embed_dim, config.embed_dim),
            nn.GELU(),
            nn.Linear(config.embed_dim, 1),
        )
        self.action_value_head = nn.Sequential(
            nn.Linear(config.embed_dim, config.embed_dim),
            nn.GELU(),
            nn.Linear(config.embed_dim, config.action_space),
        )

        # Precompute RoPE frequencies
        self.register_buffer(
            "rope_freqs",
            _precompute_rope_freqs(
                config.embed_dim // config.n_heads,
                config.max_seq_len,
                config.rope_theta,
            ),
            persistent=False,
        )

        # Initialize weights
        self._init_weights()

    def _init_weights(self) -> None:
        """Initialize weights following GPT-2/LLaMA convention."""
        for module in self.modules():
            if isinstance(module, nn.Linear):
                torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
                if module.bias is not None:
                    torch.nn.init.zeros_(module.bias)
            elif isinstance(module, nn.Embedding):
                torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)

        # Scale output projections by 1/sqrt(2*n_layers) for stable training
        scale = 1.0 / math.sqrt(2 * self.config.n_layers)
        for layer in self.layers:
            torch.nn.init.normal_(layer.attn.o_proj.weight, mean=0.0, std=0.02 * scale)  # type: ignore[union-attr, arg-type]
            torch.nn.init.normal_(layer.ffn.down_proj.weight, mean=0.0, std=0.02 * scale)  # type: ignore[union-attr, arg-type]

    def forward(
        self,
        input_ids: torch.Tensor,
        attention_mask: torch.Tensor | None = None,
    ) -> dict[str, torch.Tensor]:
        """Forward pass.

        Args:
            input_ids: (batch, seq_len) token IDs
            attention_mask: (batch, seq_len) 1 for real tokens, 0 for padding

        Returns:
            Dict with:
                policy_logits: (batch, 422) raw logits over action space
                value: (batch, 1) state value prediction
                action_values: (batch, 422) per-action value predictions
        """
        B, T = input_ids.shape

        # Token embeddings (no positional encoding — RoPE handles it)
        x = self.tok_emb(input_ids)

        # Pass through transformer blocks
        freqs = self.rope_freqs[:T]  # type: ignore[index]
        for layer in self.layers:
            x = layer(x, freqs, attention_mask)

        x = self.norm(x)

        # Pool: use the last non-padding token's representation
        if attention_mask is not None:
            # Find index of last real token per sequence
            lengths = attention_mask.sum(dim=1).long() - 1  # (B,)
            lengths = lengths.clamp(min=0)
            pooled = x[torch.arange(B, device=x.device), lengths]  # (B, embed_dim)
        else:
            pooled = x[:, -1]  # Last token

        # Output heads
        policy_logits = self.policy_head(pooled)
        value = self.value_head(pooled)
        action_values = self.action_value_head(pooled)

        return {
            "policy_logits": policy_logits,
            "value": value.squeeze(-1),
            "action_values": action_values,
        }

    def select_action(
        self,
        input_ids: torch.Tensor,
        attention_mask: torch.Tensor | None = None,
        legal_mask: torch.Tensor | None = None,
        mode: str = "greedy",
        temperature: float = 1.0,
    ) -> tuple[torch.Tensor, dict[str, torch.Tensor]]:
        """Select an action for play.

        Args:
            input_ids: (batch, seq_len) or (1, seq_len)
            attention_mask: padding mask
            legal_mask: (batch, 422) or (1, 422) legal action mask
            mode: "greedy" (argmax), "sample" (categorical), or "value" (max AV)
            temperature: Sampling temperature (for mode="sample")

        Returns:
            (action_indices, outputs) tuple
        """
        with torch.no_grad():
            outputs = self.forward(input_ids, attention_mask)

        logits = outputs["policy_logits"]

        if mode == "value":
            # Value-guided: pick action with highest predicted action-value
            av = outputs["action_values"]
            if legal_mask is not None:
                av = av.masked_fill(legal_mask == 0, float("-inf"))
            return av.argmax(dim=-1), outputs

        # Policy-based: mask illegal actions
        if legal_mask is not None:
            logits = logits.masked_fill(legal_mask == 0, float("-inf"))

        if mode == "greedy":
            return logits.argmax(dim=-1), outputs
        elif mode == "sample":
            probs = F.softmax(logits / temperature, dim=-1)
            return torch.multinomial(probs, 1).squeeze(-1), outputs
        else:
            raise ValueError(f"Unknown mode: {mode}")

    def count_parameters(self) -> int:
        """Count total trainable parameters."""
        return sum(p.numel() for p in self.parameters() if p.requires_grad)
