"""
Intelligence_OS — 1.5 Billion Parameter Sparse Decoder-Only Architecture
Implemented 100% from scratch in Python / PyTorch, matching the exact C++ specification.
Exact Theoretical Parameter Count: 1,518,659,584 parameters.
"""

import math
import torch
import torch.nn as nn
import torch.nn.functional as F
from dataclasses import dataclass
from typing import Optional, Tuple, List, Dict


@dataclass
class ModelConfig:
    vocab_size: int = 49152
    d_model: int = 2048
    n_layers: int = 26
    n_q_heads: int = 16
    n_kv_heads: int = 4
    head_dim: int = 128  # d_model // n_q_heads
    n_experts: int = 4
    top_k_experts: int = 2
    expert_hidden: int = 1792
    max_seq_len: int = 4096
    rms_eps: float = 1e-5
    adaptive_routing: bool = False
    single_expert_confidence: float = 0.80
    tied_embeddings: bool = True

    @classmethod
    def intelligence_os_1p5b(cls) -> "ModelConfig":
        """Returns the official 1.519B parameter production configuration."""
        return cls()

    @classmethod
    def smoke_test(cls) -> "ModelConfig":
        """Returns a micro configuration for rapid local testing and verification."""
        return cls(
            vocab_size=256,
            d_model=64,
            n_layers=2,
            n_q_heads=4,
            n_kv_heads=2,
            head_dim=16,
            n_experts=4,
            top_k_experts=2,
            expert_hidden=128,
            max_seq_len=128,
            single_expert_confidence=0.80,
            tied_embeddings=True,
        )

    def validate(self):
        assert self.vocab_size > 0 and self.d_model > 0 and self.n_layers > 0
        assert self.d_model % self.n_q_heads == 0, "d_model must divide evenly by n_q_heads"
        assert self.n_q_heads % self.n_kv_heads == 0, "n_q_heads must be divisible by n_kv_heads"
        assert self.n_experts >= self.top_k_experts > 0, "n_experts >= top_k_experts > 0"
        assert self.head_dim % 2 == 0, "Head dimension must be even for RoPE"

    def theoretical_parameter_count(self) -> int:
        """
        Calculates theoretical parameter count without instantiating tensors.
        Bias-free projections and tied embeddings:
        - Embeddings: vocab_size * d_model
        - Attention: Q (d_model * d_model) + K (d_model * kv_dim) + V (d_model * kv_dim) + O (d_model * d_model)
        - Experts: n_experts * 3 * d_model * expert_hidden (gate, up, down)
        - Router: d_model * n_experts
        - LayerNorms: 2 * d_model per layer
        - Final LayerNorm: d_model
        """
        self.validate()
        kv_dim = self.n_kv_heads * self.head_dim
        attention_params = (
            self.d_model * self.d_model
            + self.d_model * kv_dim
            + self.d_model * kv_dim
            + self.d_model * self.d_model
        )
        one_expert_params = 3 * self.d_model * self.expert_hidden
        router_params = self.d_model * self.n_experts
        per_layer = attention_params + (self.n_experts * one_expert_params) + router_params + (2 * self.d_model)
        total = (self.vocab_size * self.d_model) + (self.n_layers * per_layer) + self.d_model
        return total


class RMSNorm(nn.Module):
    """Pre-norm Root Mean Square Layer Normalization."""
    def __init__(self, dim: int, eps: float = 1e-5):
        super().__init__()
        self.eps = eps
        self.weight = nn.Parameter(torch.ones(dim))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        orig_dtype = x.dtype
        # Calculate variance in float32 for numerical stability
        variance = x.to(torch.float32).pow(2).mean(-1, keepdim=True)
        norm_x = x * torch.rsqrt(variance + self.eps).to(orig_dtype)
        return norm_x * self.weight


def precompute_rope_freqs(head_dim: int, seq_len: int, theta: float = 10000.0) -> Tuple[torch.Tensor, torch.Tensor]:
    """Precomputes rotary position embeddings (cos and sin cache)."""
    assert head_dim % 2 == 0, "head_dim must be even"
    indices = torch.arange(0, head_dim, 2, dtype=torch.float32)
    inv_freq = 1.0 / (theta ** (indices / head_dim))
    t = torch.arange(seq_len, dtype=torch.float32)
    freqs = torch.outer(t, inv_freq)  # [seq_len, head_dim / 2]
    cos = freqs.cos()
    sin = freqs.sin()
    return cos, sin


def apply_rope(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
    """
    Applies rotary position embedding to input tensor.
    x shape: [batch, heads, seq_len, head_dim]
    cos, sin shape: [seq_len, head_dim / 2]
    """
    b, h, t, d = x.shape
    cos_t = cos[:t].to(x.device, dtype=x.dtype).unsqueeze(0).unsqueeze(0)  # [1, 1, t, d/2]
    sin_t = sin[:t].to(x.device, dtype=x.dtype).unsqueeze(0).unsqueeze(0)  # [1, 1, t, d/2]

    # Split into even and odd components
    x_pairs = x.view(b, h, t, d // 2, 2)
    even = x_pairs[..., 0]
    odd = x_pairs[..., 1]

    # Rotate
    first = even * cos_t - odd * sin_t
    second = even * sin_t + odd * cos_t

    return torch.stack([first, second], dim=-1).flatten(-2, -1)


class GroupedQueryAttention(nn.Module):
    """
    Grouped-Query Attention (GQA): 16 query heads, 4 KV heads.
    Uses native PyTorch 2.x scaled_dot_product_attention for optimal speed/memory.
    """
    def __init__(self, config: ModelConfig):
        super().__init__()
        self.d_model = config.d_model
        self.n_q_heads = config.n_q_heads
        self.n_kv_heads = config.n_kv_heads
        self.head_dim = config.head_dim
        self.num_groups = self.n_q_heads // self.n_kv_heads

        self.q_proj = nn.Linear(self.d_model, self.d_model, bias=False)
        self.k_proj = nn.Linear(self.d_model, self.n_kv_heads * self.head_dim, bias=False)
        self.v_proj = nn.Linear(self.d_model, self.n_kv_heads * self.head_dim, bias=False)
        self.o_proj = nn.Linear(self.d_model, self.d_model, bias=False)

    def forward(
        self,
        x: torch.Tensor,
        cos: torch.Tensor,
        sin: torch.Tensor,
        mask: Optional[torch.Tensor] = None
    ) -> torch.Tensor:
        b, t, _ = x.shape

        q = self.q_proj(x).view(b, t, self.n_q_heads, self.head_dim).transpose(1, 2)
        k = self.k_proj(x).view(b, t, self.n_kv_heads, self.head_dim).transpose(1, 2)
        v = self.v_proj(x).view(b, t, self.n_kv_heads, self.head_dim).transpose(1, 2)

        # Apply RoPE
        q = apply_rope(q, cos, sin)
        k = apply_rope(k, cos, sin)

        # Expand KV heads to match Q heads (Grouped-Query Attention)
        if self.num_groups > 1:
            k = k.repeat_interleave(self.num_groups, dim=1)
            v = v.repeat_interleave(self.num_groups, dim=1)

        # PyTorch F.scaled_dot_product_attention automatically uses FlashAttention/MemEfficient kernels
        is_causal = mask is None
        attn_out = F.scaled_dot_product_attention(
            q, k, v,
            attn_mask=mask,
            dropout_p=0.0,
            is_causal=is_causal
        )

        out = attn_out.transpose(1, 2).contiguous().view(b, t, self.d_model)
        return self.o_proj(out)


class SwiGLU(nn.Module):
    """SwiGLU feed-forward expert: down_proj(silu(gate_proj(x)) * up_proj(x))."""
    def __init__(self, d_model: int, expert_hidden: int):
        super().__init__()
        self.gate_proj = nn.Linear(d_model, expert_hidden, bias=False)
        self.up_proj = nn.Linear(d_model, expert_hidden, bias=False)
        self.down_proj = nn.Linear(expert_hidden, d_model, bias=False)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x))


class SparseMoE(nn.Module):
    """
    Sparse Mixture-of-Experts with 4 SwiGLU experts per block.
    Supports standard Top-2 routing and novel Energy-Budgeted Adaptive Routing.
    """
    def __init__(self, config: ModelConfig):
        super().__init__()
        self.d_model = config.d_model
        self.n_experts = config.n_experts
        self.top_k = config.top_k_experts
        self.adaptive_routing = config.adaptive_routing
        self.single_expert_confidence = config.single_expert_confidence

        self.router = nn.Linear(self.d_model, self.n_experts, bias=False)
        self.experts = nn.ModuleList([
            SwiGLU(self.d_model, config.expert_hidden) for _ in range(self.n_experts)
        ])

    def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, Dict[str, torch.Tensor]]:
        orig_shape = x.shape
        flat = x.contiguous().view(-1, self.d_model)  # [num_tokens, d_model]
        num_tokens = flat.size(0)

        # Router logits
        logits = self.router(flat).to(torch.float32)  # [num_tokens, n_experts]
        top_weights, top_indices = torch.topk(logits, self.top_k, dim=-1)
        selected_weights = F.softmax(top_weights, dim=-1).to(flat.dtype)

        # Energy-Budgeted Adaptive Routing:
        # If the primary expert router probability exceeds threshold, skip secondary expert
        if self.adaptive_routing:
            use_second = selected_weights[:, 0] < self.single_expert_confidence
            primary_w = torch.where(use_second, selected_weights[:, 0], torch.ones_like(selected_weights[:, 0]))
            secondary_w = selected_weights[:, 1] * use_second.to(flat.dtype)
            selected_weights = torch.stack([primary_w, secondary_w], dim=1)
        else:
            use_second = torch.ones(num_tokens, dtype=torch.bool, device=flat.device)

        # Dispatch tokens only to selected experts
        output = torch.zeros_like(flat)

        for e in range(self.n_experts):
            for slot in range(self.top_k):
                active = (top_indices[:, slot] == e)
                if slot == 1:
                    active = active & use_second
                token_ids = torch.nonzero(active).squeeze(-1)
                if token_ids.numel() == 0:
                    continue
                token_x = flat.index_select(0, token_ids)
                expert_y = self.experts[e](token_x)
                route_w = selected_weights[:, slot].index_select(0, token_ids).unsqueeze(1)
                output.index_add_(0, token_ids, expert_y * route_w)

        # Auxiliary telemetry
        aux = {
            "router_logits": logits,
            "use_second_ratio": use_second.float().mean().detach() if self.adaptive_routing else torch.tensor(1.0)
        }

        return output.view(orig_shape), aux


class DecoderBlock(nn.Module):
    """Pre-norm Decoder Block with Grouped-Query Attention and Sparse MoE."""
    def __init__(self, config: ModelConfig):
        super().__init__()
        self.norm_attn = RMSNorm(config.d_model, config.rms_eps)
        self.attention = GroupedQueryAttention(config)
        self.norm_moe = RMSNorm(config.d_model, config.rms_eps)
        self.moe = SparseMoE(config)

    def forward(
        self,
        x: torch.Tensor,
        cos: torch.Tensor,
        sin: torch.Tensor,
        mask: Optional[torch.Tensor] = None
    ) -> Tuple[torch.Tensor, Dict[str, torch.Tensor]]:
        # Pre-norm Attention with residual
        normed_attn = self.norm_attn(x)
        attn_out = self.attention(normed_attn, cos, sin, mask=mask)
        x = x + attn_out

        # Pre-norm MoE with residual
        normed_moe = self.norm_moe(x)
        moe_out, aux = self.moe(normed_moe)
        x = x + moe_out

        return x, aux


class IntelligenceOS(nn.Module):
    """
    Intelligence_OS: 1.5 Billion Parameter Decoder-Only Sparse Transformer.
    Features:
    - 26 Decoder Layers
    - 2048 Hidden Width
    - 16 Query Heads / 4 KV Heads (GQA)
    - RoPE Positional Encoding
    - 4 SwiGLU Experts per layer with Top-2 or Adaptive Routing
    - Tied Input/Output Embedding Head
    """
    def __init__(self, config: ModelConfig):
        super().__init__()
        self.config = config
        self.config.validate()

        self.token_embedding = nn.Embedding(config.vocab_size, config.d_model)
        self.layers = nn.ModuleList([
            DecoderBlock(config) for _ in range(config.n_layers)
        ])
        self.final_norm = RMSNorm(config.d_model, config.rms_eps)

        # Precompute RoPE frequencies cache
        cos, sin = precompute_rope_freqs(config.head_dim, config.max_seq_len)
        self.register_buffer("rope_cos", cos, persistent=False)
        self.register_buffer("rope_sin", sin, persistent=False)

    def forward(
        self,
        token_ids: torch.Tensor,
        targets: Optional[torch.Tensor] = None,
        return_loss: bool = True
    ) -> Tuple[torch.Tensor, Optional[torch.Tensor], List[Dict[str, torch.Tensor]]]:
        """
        Forward pass for next-token prediction.
        token_ids shape: [batch_size, seq_len]
        targets shape: [batch_size, seq_len] (optional)
        """
        b, t = token_ids.shape
        assert t <= self.config.max_seq_len, f"Sequence length {t} exceeds maximum {self.config.max_seq_len}"

        x = self.token_embedding(token_ids)
        cos = self.rope_cos[:t]
        sin = self.rope_sin[:t]

        all_aux = []
        for layer in self.layers:
            x, aux = layer(x, cos, sin)
            all_aux.append(aux)

        x = self.final_norm(x)

        # Tied language-model head projection:
        # Reuses token_embedding.weight for output projection: [batch, seq, d_model] @ [d_model, vocab]
        logits = F.linear(x, self.token_embedding.weight)

        loss = None
        if targets is not None and return_loss:
            loss = F.cross_entropy(
                logits.view(-1, self.config.vocab_size),
                targets.view(-1),
                ignore_index=-100
            )

        return logits, loss, all_aux

    def count_parameters(self) -> int:
        """Returns the actual count of instantiated unique trainable parameters."""
        return sum(p.numel() for p in self.parameters() if p.requires_grad)

    @torch.no_grad()
    def generate(
        self,
        prompt_ids: torch.Tensor,
        max_new_tokens: int = 50,
        temperature: float = 0.8,
        top_k: int = 50,
        eos_token_id: Optional[int] = None
    ) -> torch.Tensor:
        """Autoregressive text generation."""
        self.eval()
        generated = prompt_ids.clone()

        for _ in range(max_new_tokens):
            seq_len = generated.size(1)
            if seq_len > self.config.max_seq_len:
                generated_input = generated[:, -self.config.max_seq_len:]
            else:
                generated_input = generated

            logits, _, _ = self.forward(generated_input, return_loss=False)
            next_token_logits = logits[:, -1, :] / max(temperature, 1e-5)

            if top_k > 0:
                v, _ = torch.topk(next_token_logits, min(top_k, next_token_logits.size(-1)))
                next_token_logits[next_token_logits < v[:, [-1]]] = -float("Inf")

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

            if eos_token_id is not None and (next_token == eos_token_id).all():
                break

        return generated