"""ARCHON-Brain — 300M Transformer with Multi-Token Prediction.
Custom architecture optimized for ARCHON's domain: code, systems, trading, security.
"""
import torch
import torch.nn as nn
import torch.nn.functional as F
import math
from config import ArchonBrainConfig


class RMSNorm(nn.Module):
    """Root Mean Square Layer Normalization (more efficient than LayerNorm)."""
    def __init__(self, dim: int, eps: float = 1e-6):
        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


class RotaryEmbedding(nn.Module):
    """Rotary Position Embedding (RoPE) — encodes position via rotation."""
    def __init__(self, dim: int, max_seq_len: int = 2048, theta: float = 10_000.0):
        super().__init__()
        inv_freq = 1.0 / (theta ** (torch.arange(0, dim, 2).float() / dim))
        self.register_buffer("inv_freq", inv_freq)

        t = torch.arange(max_seq_len)
        freqs = torch.outer(t, self.inv_freq)
        self.register_buffer("cos_cached", freqs.cos(), persistent=False)
        self.register_buffer("sin_cached", freqs.sin(), persistent=False)

    def forward(self, seq_len: int):
        return self.cos_cached[:seq_len], self.sin_cached[:seq_len]


def apply_rotary_emb(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
    """Apply RoPE rotation to query/key tensors."""
    d = x.shape[-1] // 2
    x1, x2 = x[..., :d], x[..., d:]
    cos = cos[:x.shape[-2]].unsqueeze(0).unsqueeze(0)
    sin = sin[:x.shape[-2]].unsqueeze(0).unsqueeze(0)
    return torch.cat([x1 * cos - x2 * sin, x2 * cos + x1 * sin], dim=-1)


class Attention(nn.Module):
    """Multi-Head Self-Attention with RoPE and causal mask."""
    def __init__(self, config: ArchonBrainConfig):
        super().__init__()
        self.num_heads = config.num_heads
        self.head_dim = config.head_dim

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

        self.rotary = RotaryEmbedding(config.head_dim, config.max_seq_len, config.rope_theta)

    def forward(self, x: torch.Tensor, mask: torch.Tensor = None) -> torch.Tensor:
        B, T, C = x.shape
        q = self.q_proj(x).view(B, T, self.num_heads, self.head_dim).transpose(1, 2)
        k = self.k_proj(x).view(B, T, self.num_heads, self.head_dim).transpose(1, 2)
        v = self.v_proj(x).view(B, T, self.num_heads, self.head_dim).transpose(1, 2)

        cos, sin = self.rotary(T)
        q = apply_rotary_emb(q, cos, sin)
        k = apply_rotary_emb(k, cos, sin)

        # Scaled dot-product attention with causal mask
        y = F.scaled_dot_product_attention(q, k, v, is_causal=True)
        y = y.transpose(1, 2).contiguous().view(B, T, C)
        return self.o_proj(y)


class SwiGLUFFN(nn.Module):
    """SwiGLU Feed-Forward Network (LLaMA style, better than standard FFN)."""
    def __init__(self, config: ArchonBrainConfig):
        super().__init__()
        self.gate = nn.Linear(config.hidden_dim, config.intermediate_dim, bias=False)
        self.up = nn.Linear(config.hidden_dim, config.intermediate_dim, bias=False)
        self.down = nn.Linear(config.intermediate_dim, config.hidden_dim, bias=False)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.down(F.silu(self.gate(x)) * self.up(x))


class TransformerBlock(nn.Module):
    """Pre-norm Transformer block: RMSNorm → Attention → Residual → RMSNorm → FFN → Residual."""
    def __init__(self, config: ArchonBrainConfig):
        super().__init__()
        self.attn_norm = RMSNorm(config.hidden_dim, config.norm_eps)
        self.attn = Attention(config)
        self.ffn_norm = RMSNorm(config.hidden_dim, config.norm_eps)
        self.ffn = SwiGLUFFN(config)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = x + self.attn(self.attn_norm(x))
        x = x + self.ffn(self.ffn_norm(x))
        return x


class MTPHead(nn.Module):
    """Multi-Token Prediction head — projects hidden state to predict future token.
    Each head predicts token at offset +k (k=1,2,3,4,5)."""
    def __init__(self, hidden_dim: int, vocab_size: int):
        super().__init__()
        self.proj = nn.Linear(hidden_dim, hidden_dim, bias=False)
        self.norm = RMSNorm(hidden_dim)

    def forward(self, x: torch.Tensor, shared_embed_weight: torch.Tensor) -> torch.Tensor:
        """Returns logits by projecting through learned transform then shared embedding."""
        h = self.norm(self.proj(x))
        return F.linear(h, shared_embed_weight)  # [B, T, vocab]


class ArchonBrain(nn.Module):
    """ARCHON-Brain: 300M parameter transformer with MTP=5.

    Architecture:
    - RoPE positional encoding (no learned positions)
    - Pre-norm with RMSNorm
    - SwiGLU FFN
    - Tied input/output embeddings
    - 5 Multi-Token Prediction heads
    """
    def __init__(self, config: ArchonBrainConfig = None):
        super().__init__()
        if config is None:
            config = ArchonBrainConfig()
        self.config = config

        # Token embedding (shared with output)
        self.embed = nn.Embedding(config.vocab_size, config.hidden_dim)

        # Transformer layers
        self.layers = nn.ModuleList([
            TransformerBlock(config) for _ in range(config.num_layers)
        ])

        # Final norm
        self.norm = RMSNorm(config.hidden_dim, config.norm_eps)

        # Main LM head (next token prediction) — tied with embedding
        self.lm_head = nn.Linear(config.hidden_dim, config.vocab_size, bias=False)
        self.lm_head.weight = self.embed.weight  # Weight tying

        # MTP heads (predict +2, +3, +4, +5 tokens ahead)
        # Head 0 = main LM head (+1), heads 1-4 = MTP
        self.mtp_heads = nn.ModuleList([
            MTPHead(config.hidden_dim, config.vocab_size)
            for _ in range(config.mtp_heads - 1)  # -1 because main head is +1
        ])

        # Initialize weights
        self.apply(self._init_weights)

    def _init_weights(self, module):
        if isinstance(module, nn.Linear):
            nn.init.normal_(module.weight, mean=0.0, std=0.02)
            if module.bias is not None:
                nn.init.zeros_(module.bias)
        elif isinstance(module, nn.Embedding):
            nn.init.normal_(module.weight, mean=0.0, std=0.02)

    def forward(self, input_ids: torch.Tensor, targets: torch.Tensor = None):
        """
        Args:
            input_ids: [B, T] token IDs
            targets: [B, T] target token IDs (shifted by 1 for next-token)

        Returns:
            logits: [B, T, vocab] main predictions
            loss: scalar if targets provided (includes MTP loss)
            mtp_logits: list of [B, T, vocab] for each MTP head
        """
        B, T = input_ids.shape

        # Embed tokens
        h = self.embed(input_ids)

        # Run through transformer layers
        for layer in self.layers:
            h = layer(h)

        # Final norm
        h = self.norm(h)

        # Main LM head — next token prediction
        logits = self.lm_head(h)  # [B, T, vocab]

        # Compute loss if targets provided
        loss = None
        mtp_logits = []
        if targets is not None:
            # Main loss: cross-entropy for next token
            main_loss = F.cross_entropy(
                logits[:, :-1].reshape(-1, self.config.vocab_size),
                targets[:, 1:].reshape(-1),
                ignore_index=-100,
            )
            loss = self.config.mtp_loss_weights[0] * main_loss

            # MTP losses: predict +2, +3, +4, +5
            for i, mtp_head in enumerate(self.mtp_heads):
                offset = i + 2  # +2, +3, +4, +5
                mtp_logit = mtp_head(h, self.embed.weight)
                mtp_logits.append(mtp_logit)

                if T > offset:
                    mtp_loss = F.cross_entropy(
                        mtp_logit[:, :-offset].reshape(-1, self.config.vocab_size),
                        targets[:, offset:].reshape(-1),
                        ignore_index=-100,
                    )
                    loss = loss + self.config.mtp_loss_weights[i + 1] * mtp_loss

        return logits, loss, mtp_logits

    @torch.no_grad()
    def generate(self, input_ids: torch.Tensor, max_new_tokens: int = 256,
                 temperature: float = 0.7, top_k: int = 50, top_p: float = 0.9) -> torch.Tensor:
        """Autoregressive text generation."""
        for _ in range(max_new_tokens):
            # Crop to max seq len
            idx_cond = input_ids[:, -self.config.max_seq_len:]
            logits, _, _ = self(idx_cond)
            logits = logits[:, -1, :] / temperature

            # Top-k filtering
            if top_k > 0:
                v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
                logits[logits < v[:, [-1]]] = float('-inf')

            # Top-p (nucleus) filtering
            if top_p < 1.0:
                sorted_logits, sorted_indices = torch.sort(logits, descending=True)
                cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1)
                sorted_indices_to_remove = cumulative_probs > top_p
                sorted_indices_to_remove[:, 1:] = sorted_indices_to_remove[:, :-1].clone()
                sorted_indices_to_remove[:, 0] = 0
                indices_to_remove = sorted_indices_to_remove.scatter(1, sorted_indices, sorted_indices_to_remove)
                logits[indices_to_remove] = float('-inf')

            probs = F.softmax(logits, dim=-1)
            next_token = torch.multinomial(probs, num_samples=1)
            input_ids = torch.cat([input_ids, next_token], dim=1)

            # Stop on EOS (token 2 by convention)
            if next_token.item() == 2:
                break

        return input_ids

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


if __name__ == "__main__":
    config = ArchonBrainConfig()
    model = ArchonBrain(config)
    print(f"ARCHON-Brain initialized")
    print(f"  Actual parameters: {model.count_parameters():,}")
    print(f"  Estimated: {config.param_count:,}")
    print(f"  Config estimate: {config.param_count_human}")

    # Test forward pass
    x = torch.randint(0, config.vocab_size, (2, 128))
    logits, loss, mtp = model(x, x)
    print(f"\n  Forward pass OK:")
    print(f"  Input: {x.shape}")
    print(f"  Logits: {logits.shape}")
    print(f"  Loss: {loss.item():.4f}")
    print(f"  MTP heads: {len(mtp)}")

    # Test generation
    prompt = torch.randint(0, config.vocab_size, (1, 10))
    out = model.generate(prompt, max_new_tokens=20)
    print(f"\n  Generation OK: {prompt.shape} -> {out.shape}")