import math
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.checkpoint import checkpoint
from dataclasses import dataclass

@dataclass
class ModelConfig:
    """Configuration for ViuResonance100M - T4 Optimized"""
    vocab_size: int = 64000
    d_model: int = 1024
    n_layers: int = 12
    n_heads: int = 16
    dropout: float = 0.1
    chunk_size: int = 256  # T4 ke liye 256, 512 se OOM aayega
    max_seq_len: int = 4096

class RMSNorm(nn.Module):
    def __init__(self, dim: int, eps: float = 1e-6):
        super().__init__()
        self.weight = nn.Parameter(torch.ones(dim))
        self.eps = eps
        
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # T4 fp16 me stable rakhne ke liye float32 me norm
        try:
            # PyTorch 2.4+ me available
            if hasattr(F, 'rms_norm'):
                return F.rms_norm(x.float(), (x.shape[-1],), self.weight.float(), self.eps).to(x.dtype)
        except Exception:
            pass
        norm_x = x.float() * torch.rsqrt(x.float().pow(2).mean(-1, keepdim=True) + self.eps)
        return (norm_x * self.weight.float()).to(x.dtype)
        # return (norm_x * self.weight.float()).to(x.dtype)

class ResonanceLayerKaggle(nn.Module):
    def __init__(self, config: ModelConfig):
        super().__init__()
        if config.d_model % config.n_heads != 0:
            raise ValueError(f"d_model ({config.d_model}) must be divisible by n_heads ({config.n_heads})")
            
        self.n_heads = config.n_heads
        self.head_dim = config.d_model // config.n_heads
        self.chunk_size = config.chunk_size
        
        self.to_amp = nn.Linear(config.d_model, config.d_model, bias=False)
        self.to_freq = nn.Linear(config.d_model, config.n_heads, bias=False)
        self.to_phase = nn.Linear(config.d_model, config.n_heads, bias=False)
        self.to_v = nn.Linear(config.d_model, config.d_model, bias=False)
        self.to_out = nn.Linear(config.d_model, config.d_model, bias=False)
        self.norm = RMSNorm(config.d_model)
        
        # T4: pos cache - har forward me arange banane se bachao
        self.register_buffer("pos_cache", torch.arange(config.max_seq_len, dtype=torch.long), persistent=False)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        B, T, C = x.shape
        h = self.norm(x)
        A = self.to_amp(h).view(B, T, self.n_heads, self.head_dim)
        V = self.to_v(h).view(B, T, self.n_heads, self.head_dim)
        freq = torch.tanh(self.to_freq(h)) * 2.0  # B,T,H
        phase = torch.tanh(self.to_phase(h)) * math.pi

        # pos cache use karo, T tak slice
        if T > self.pos_cache.shape[0]:
            # Agar kabhi bada T aaye to dynamically banao
            pos = torch.arange(T, device=x.device, dtype=torch.long)
        else:
            pos = self.pos_cache[:T]

        out = torch.empty(B, T, C, device=x.device, dtype=x.dtype)
        
        for i in range(0, T, self.chunk_size):
            end = min(i+self.chunk_size, T)
            chunk_len = end - i
            
            phase_i = phase[:, i:end, :].permute(0,2,1).unsqueeze(-1) # B,H,chunk,1
            phase_j = phase.permute(0,2,1).unsqueeze(-2) # B,H,1,T
            
            row_idx = torch.arange(i, end, device=x.device, dtype=torch.long)[:, None]
            col_idx = pos[None, :]
            dist_long = (row_idx - col_idx).clamp(min=0)
            dist = dist_long.to(freq.dtype).view(1, 1, chunk_len, T)
            
            freq_i = freq[:, i:end, :].permute(0,2,1).unsqueeze(-1)
            angle = (phase_i - phase_j) + freq_i * dist * 0.05
            
            score = torch.cos(angle) # B,H,chunk,T
            
            causal_bool = (row_idx >= col_idx).view(1,1,chunk_len,T)
            score = score.masked_fill(~causal_bool, 0.0)
            score = score / math.sqrt(self.head_dim)
            
            amp = A.norm(dim=-1).permute(0,2,1) # B,H,T
            amp_i = amp[:,:,i:end].unsqueeze(-1)
            amp_j = amp.unsqueeze(-2)
            score = score * torch.sigmoid(amp_i * amp_j)
            
            V_t = V.permute(0,2,1,3).contiguous() # B,H,T,D
            Bh = B * self.n_heads
            
            s = score.reshape(Bh, chunk_len, T)
            v = V_t.reshape(Bh, T, self.head_dim)
            o = torch.bmm(s, v).view(B, self.n_heads, chunk_len, self.head_dim).permute(0,2,1,3).reshape(B, chunk_len, C)
            out[:, i:end] = o

        return self.to_out(out)

class ViuResonance100M(nn.Module):
    def __init__(self, config: ModelConfig):
        super().__init__()
        self.config = config
        self.emb = nn.Embedding(config.vocab_size, config.d_model)
        self.layers = nn.ModuleList([ResonanceLayerKaggle(config) for _ in range(config.n_layers)])
        self.norm = RMSNorm(config.d_model)
        self.head = nn.Linear(config.d_model, config.vocab_size, bias=False)
        self.resid_dropout = nn.Dropout(config.dropout)
        
        # Sahi order: pehle init, phir tie
        self.apply(self._init_weights)
        self.head.weight = self.emb.weight

    def _init_weights(self, module: nn.Module):
        if isinstance(module, nn.Linear):
            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
        elif isinstance(module, nn.Embedding):
            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
        elif isinstance(module, RMSNorm):
            torch.nn.init.ones_(module.weight)

    def forward(self, idx: torch.Tensor, targets: torch.Tensor = None, use_checkpoint: bool = True) -> tuple[torch.Tensor, torch.Tensor | None]:
        x = self.emb(idx)
        for layer in self.layers:
            if self.training and use_checkpoint:
                x = x + self.resid_dropout(checkpoint(layer, x, use_reentrant=False))
            else:
                x = x + self.resid_dropout(layer(x))
                
        logits = self.head(self.norm(x))
        loss = None
        if targets is not None:
            loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1), ignore_index=-100)
        return logits, loss