import math
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.checkpoint import checkpoint
from config import ViuAIConfig

class RMSNorm(nn.Module):
    def __init__(self, dim, eps=1e-5):
        super().__init__()
        self.eps = eps
        self.weight = nn.Parameter(torch.ones(dim))

    def forward(self, x):
        x_fp32 = x.float()
        norm = x_fp32 * torch.rsqrt(x_fp32.pow(2).mean(-1, keepdim=True) + self.eps)
        # Optimized: removed redundant self.weight.float() cast for bfloat16/float16 throughput
        return (norm * self.weight).type_as(x)

def precompute_rope(head_dim, max_len, theta=10000.0, device="cpu"):
    freqs = 1.0 / (theta ** (torch.arange(0, head_dim, 2, device=device).float() / head_dim))
    t = torch.arange(max_len, device=device).float()
    freqs = torch.outer(t, freqs)
    return torch.cos(freqs), torch.sin(freqs)

def apply_rope(x, cos, sin):
    cos = cos.to(x.dtype)
    sin = sin.to(x.dtype)
    x1, x2 = x[..., ::2], x[..., 1::2]
    # x is [B, n_heads, T, head_dim]
    # cos, sin are [max_len, head_dim//2]
    T = x.size(2)
    cos = cos[:T][None, None, :, :]
    sin = sin[:T][None, None, :, :]
    rotated = torch.stack([x1 * cos - x2 * sin, x1 * sin + x2 * cos], dim=-1)
    return rotated.flatten(-2)

def repeat_kv(x, n_rep):
    if n_rep == 1:
        return x
    return x.repeat_interleave(n_rep, dim=1)

class Attention(nn.Module):
    def __init__(self, cfg: ViuAIConfig):
        super().__init__()
        self.n_heads = cfg.n_heads
        self.n_kv_heads = cfg.n_kv_heads
        self.head_dim = cfg.d_model // cfg.n_heads
        self.n_rep = self.n_heads // self.n_kv_heads

        self.q_proj = nn.Linear(cfg.d_model, cfg.n_heads * self.head_dim, bias=False)
        self.k_proj = nn.Linear(cfg.d_model, cfg.n_kv_heads * self.head_dim, bias=False)
        self.v_proj = nn.Linear(cfg.d_model, cfg.n_kv_heads * self.head_dim, bias=False)
        self.o_proj = nn.Linear(cfg.n_heads * self.head_dim, cfg.d_model, bias=False)
        self.o_proj._is_residual = True  # flag for depth-scaled init

        self.q_norm = RMSNorm(self.head_dim, cfg.norm_eps)
        self.k_norm = RMSNorm(self.head_dim, cfg.norm_eps)
        self.attn_dropout = cfg.attn_dropout

    def forward(self, x, cos, sin, attn_mask=None):   
        B, T, C = x.shape
        q = self.q_proj(x).view(B, T, self.n_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)

        q, k = self.q_norm(q), self.k_norm(k)
        q, k = apply_rope(q, cos, sin), apply_rope(k, cos, sin)

        dropout_p = self.attn_dropout if self.training else 0.0
        try:
            if attn_mask is not None:
                out = F.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask, dropout_p=dropout_p, enable_gqa=True)
            else:
                out = F.scaled_dot_product_attention(q, k, v, is_causal=True, dropout_p=dropout_p, enable_gqa=True)
        except TypeError:
            if self.n_rep > 1:
                k = repeat_kv(k, self.n_rep)
                v = repeat_kv(v, self.n_rep)
            if attn_mask is not None:
                out = F.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask, dropout_p=dropout_p)
            else:
                out = F.scaled_dot_product_attention(q, k, v, is_causal=True, dropout_p=dropout_p)

        out = out.transpose(1, 2).contiguous().view(B, T, -1)
        return self.o_proj(out)

class SwiGLU(nn.Module):
    def __init__(self, cfg: ViuAIConfig):
        super().__init__()
        self.gate_proj = nn.Linear(cfg.d_model, cfg.ffn_hidden, bias=False)
        self.up_proj = nn.Linear(cfg.d_model, cfg.ffn_hidden, bias=False)
        self.down_proj = nn.Linear(cfg.ffn_hidden, cfg.d_model, bias=False)
        self.down_proj._is_residual = True  # flag for depth-scaled init

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

class Block(nn.Module):
    def __init__(self, cfg: ViuAIConfig):
        super().__init__()
        self.attn_norm = RMSNorm(cfg.d_model, cfg.norm_eps)
        self.attn = Attention(cfg)
        self.ffn_norm = RMSNorm(cfg.d_model, cfg.norm_eps)
        self.ffn = SwiGLU(cfg)
        self.resid_dropout = nn.Dropout(cfg.resid_dropout) if cfg.resid_dropout > 0 else nn.Identity()

    def forward(self, x, cos, sin, attn_mask=None):   
        x = x + self.resid_dropout(self.attn(self.attn_norm(x), cos, sin, attn_mask))
        x = x + self.resid_dropout(self.ffn(self.ffn_norm(x)))
        return x

class ViuAI(nn.Module):
    def __init__(self, cfg: ViuAIConfig):
        super().__init__()
        assert cfg.d_model % cfg.n_heads == 0, "d_model must be divisible by n_heads"
        assert cfg.n_heads % cfg.n_kv_heads == 0, "n_heads must be divisible by n_kv_heads"
        assert (cfg.d_model // cfg.n_heads) % 2 == 0, "head_dim must be even for RoPE"
        
        self.cfg = cfg
        self.tok_emb = nn.Embedding(cfg.vocab_size, cfg.d_model)
        self.blocks = nn.ModuleList([Block(cfg) for _ in range(cfg.n_layers)])
        self.final_norm = RMSNorm(cfg.d_model, cfg.norm_eps)
        self.head = nn.Linear(cfg.d_model, cfg.vocab_size, bias=False)
        self.head.weight = self.tok_emb.weight

        head_dim = cfg.d_model // cfg.n_heads
        cos, sin = precompute_rope(head_dim, cfg.context_length, cfg.rope_theta)
        self.register_buffer("rope_cos", cos, persistent=False)
        self.register_buffer("rope_sin", sin, persistent=False)

        self.apply(self._init_weights)

    def _init_weights(self, module):
        if isinstance(module, nn.Linear):
            if hasattr(self, 'head') and module is self.head:
                return
            std = 0.02
            if getattr(module, '_is_residual', False):
                std *= (2 * self.cfg.n_layers) ** -0.5
            nn.init.normal_(module.weight, mean=0.0, std=std)
        elif isinstance(module, nn.Embedding):
            nn.init.normal_(module.weight, mean=0.0, std=0.02)

    def forward(self, idx, targets=None, pad_id=None, shift_labels=False):   
        B, T = idx.shape
        assert T <= self.cfg.context_length, f"Sequence length {T} exceeds context_length {self.cfg.context_length}"
        x = self.tok_emb(idx)
        
        # NEFTune: Noisy Embedding Fine-Tuning for conversational quality boost
        if self.training and getattr(self.cfg, "neftune_alpha", 0.0) > 0.0:
            dims = x.size(1) * x.size(2)
            mag = self.cfg.neftune_alpha / math.sqrt(max(1, dims))
            noise = torch.zeros_like(x).uniform_(-1.0, 1.0) * mag
            x = x + noise

        cos = self.rope_cos.to(x.device)
        sin = self.rope_sin.to(x.device)

        # BUG 2 FIX: Additive Float Attention Mask (0.0 = attend, -inf = mask out)
        attn_mask = None
        if pad_id is not None:
            causal = torch.tril(torch.ones(T, T, dtype=torch.bool, device=x.device))
            key_mask = (idx != pad_id).unsqueeze(1).unsqueeze(2)     # [B, 1, 1, T]
            bool_mask = causal.unsqueeze(0).unsqueeze(0) & key_mask   # [B, 1, T, T] True = allowed
            attn_mask = torch.zeros(B, 1, T, T, dtype=x.dtype, device=x.device)
            attn_mask = attn_mask.masked_fill(~bool_mask, float("-inf"))

        for block in self.blocks:
            if self.training and self.cfg.use_checkpoint:
                if attn_mask is not None:
                    x = checkpoint(block, x, cos, sin, attn_mask, use_reentrant=False)
                else:
                    x = checkpoint(block, x, cos, sin, use_reentrant=False)
            else:
                x = block(x, cos, sin, attn_mask)
        x = self.final_norm(x)

        # =====================================================================
        # ⚡ Chunked Cross-Entropy: Compute loss in 512-token chunks to avoid
        #    materializing the full [B, T, V=64003] logits tensor (~16 GB).
        #    Mathematically IDENTICAL to standard CE. Saves ~12 GB VRAM.
        # =====================================================================
        loss = None
        if targets is not None:
            V = self.cfg.vocab_size
            ce_chunk = 512  # tokens per chunk

            if shift_labels:
                # Shifted: predict token[i+1] from hidden[i]
                hidden = x[:, :-1, :].contiguous()   # [B, T-1, D]
                tgt = targets[:, 1:].contiguous()     # [B, T-1]
            else:
                hidden = x
                tgt = targets

            total_loss = torch.zeros(1, device=x.device, dtype=torch.float32)
            total_active = (tgt != -100).sum().clamp(min=1)
            seq_len = hidden.size(1)

            for i in range(0, seq_len, ce_chunk):
                end = min(i + ce_chunk, seq_len)
                chunk_logits = self.head(hidden[:, i:end, :])       # [B, chunk, V]
                chunk_tgt = tgt[:, i:end]                           # [B, chunk]
                chunk_loss = F.cross_entropy(
                    chunk_logits.reshape(-1, V), chunk_tgt.reshape(-1),
                    ignore_index=-100, reduction="sum"
                )
                total_loss = total_loss + chunk_loss

            loss = total_loss / total_active

            # z-loss: chunked like CE above so pretraining never materializes
            # the full [B, T, V] logits tensor (~16 GB). Mathematically identical.
            if self.cfg.z_loss_weight > 0:
                z_sum = torch.zeros(1, device=x.device, dtype=torch.float32)
                z_count = 0
                x_flat_len = x.size(1)
                for i in range(0, x_flat_len, ce_chunk):
                    end = min(i + ce_chunk, x_flat_len)
                    chunk_logits = self.head(x[:, i:end, :])
                    z_sum = z_sum + (torch.logsumexp(chunk_logits.float(), dim=-1) ** 2).sum()
                    z_count += chunk_logits.size(0) * chunk_logits.size(1)
                z_loss = self.cfg.z_loss_weight * (z_sum / max(1, z_count))
                loss = loss + z_loss

            # During training, return None logits to save ~16 GB VRAM
            # Caller must NOT use logits when targets are provided
            return None, loss

        # Inference path: return full logits for generation
        logits = self.head(x)
        return logits, loss

    def num_params(self):
        """Count parameters, deduplicating tied weights."""
        seen = set()
        total = 0
        for p in self.parameters():
            if p.data_ptr() not in seen:
                seen.add(p.data_ptr())
                total += p.numel()
        return total

    @torch.no_grad()
    def generate(self, idx, max_new_tokens, temperature=1.0, top_k=50, top_p=0.9,
                 eos_token_id=None, repetition_penalty=1.0):
        self.eval()
        idx = idx.to(next(self.parameters()).device)
        def _is_all_eos(next_tokens):
            if eos_token_id is None:
                return False
            if isinstance(eos_token_id, (list, tuple, set)):
                is_eos = torch.zeros(next_tokens.shape[0], dtype=torch.bool, device=next_tokens.device)
                for eid in eos_token_id:
                    is_eos |= (next_tokens.squeeze(-1) == eid)
                return is_eos.all().item()
            return (next_tokens == eos_token_id).all().item()

        for _ in range(max_new_tokens):
            # Crop to context window
            idx_cond = idx if idx.size(1) <= self.cfg.context_length else idx[:, -self.cfg.context_length:]
            logits, _ = self(idx_cond)
            logits = logits[:, -1, :]  # only last position

            # Repetition penalty (Standard Hugging Face formulation)
            if repetition_penalty != 1.0:
                for b in range(idx.size(0)):
                    prev_tokens = idx[b].unique()
                    score = logits[b, prev_tokens]
                    logits[b, prev_tokens] = torch.where(
                        score < 0, score * repetition_penalty, score / repetition_penalty
                    )

            if temperature <= 0:
                idx_next = logits.argmax(dim=-1, keepdim=True)
                idx = torch.cat([idx, idx_next], dim=1)
                if _is_all_eos(idx_next):
                    break
                continue

            # Temperature scaling
            logits = logits / max(temperature, 1e-8)
            if not torch.isfinite(logits).all():
                raise ValueError("Model logits became non-finite (NaN/Inf) during generation. Ensure stable precision or checkpoint validity.")

            # Top-k filtering
            if top_k > 0:
                top_k_val = min(top_k, logits.size(-1))
                kth_vals, _ = torch.topk(logits, top_k_val)
                logits[logits < kth_vals[:, [-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] = False
                indices_to_remove = torch.zeros_like(logits, dtype=torch.bool).scatter_(
                    1, sorted_indices, sorted_indices_to_remove
                )
                logits[indices_to_remove] = float('-inf')

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

            if _is_all_eos(idx_next):
                break

        return idx

# Aliases for universal compatibility
Transformer = ViuAI
ModelArgs = ViuAIConfig