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

class TokenImportanceNetwork(nn.Module):
    """
    Computes importance scores for each token based on:
    1. Local context patterns
    2. Token frequency
    3. Position information
    """
    def __init__(self, config):
        super().__init__()
        self.n_embd = config.n_embd
        
        # Local context processing
        self.context_net = nn.Sequential(
            nn.Conv1d(config.n_embd, config.n_embd, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.Conv1d(config.n_embd, config.n_embd, kernel_size=1)
        )
        
        # Frequency awareness
        self.freq_embedding = nn.Embedding(256, config.n_embd)
        
        # Position awareness
        self.pos_embedding = nn.Embedding(config.block_size, config.n_embd)
        
        # Feature fusion
        self.fusion = nn.Sequential(
            nn.LayerNorm(config.n_embd * 3),
            nn.Linear(config.n_embd * 3, config.n_embd),
            nn.Dropout(config.importance_dropout),  
            nn.GELU(),
            nn.Linear(config.n_embd, 1),
            nn.Dropout(config.importance_dropout),  
            nn.Sigmoid()
        )
        
    def forward(self, x, freq_table, positions):
        B, T, C = x.shape
        
        # Ensure inputs are on the correct device
        freq_table = freq_table.to(x.device)
        positions = positions.to(x.device)
        
        # Process local context
        x_local = self.context_net(x.transpose(1, 2))  # [B, C, T]
        x_local = x_local.transpose(1, 2)  # [B, T, C]
        
        # Get embeddings
        freq_emb = self.freq_embedding(freq_table)  # [B, T, C]
        pos_emb = self.pos_embedding(positions)  # [B, T, C]
        
        # Concatenate features
        combined = torch.cat([x_local, freq_emb, pos_emb], dim=-1)  # [B, T, 3C]
        
        # Compute importance scores
        importance = self.fusion(combined)  # [B, T, 1]
        
        return importance

class SparseDenseAttention(nn.Module):
    """
    Ultra memory-efficient hybrid attention using:
    - Flash attention style computation
    - Gradient checkpointing
    - Aggressive memory management
    """
    def __init__(self, config):
        super().__init__()
        assert config.n_embd % config.n_head == 0
        
        self.n_head = config.n_head
        self.n_embd = config.n_embd
        self.head_size = config.n_embd // config.n_head
        self.dropout = config.dropout
        
        # Key, Query, Value projections for all heads
        self.c_attn = nn.Linear(config.n_embd, 3 * config.n_embd, bias=config.bias)
        self.c_proj = nn.Linear(config.n_embd, config.n_embd, bias=config.bias)
        
        # Dropouts
        self.attn_dropout = nn.Dropout(config.dropout)
        self.resid_dropout = nn.Dropout(config.dropout)
        
        # Sparse attention parameters
        self.sparse_topk = min(
            getattr(config, 'sparse_topk', 32),
            config.block_size // 4
        )
        
        # Initialize scale
        self.register_buffer("scale", torch.tensor(1.0 / math.sqrt(self.head_size)))
        
    def _chunk_attention(self, q, k, v, importance_chunk, chunk_size):
        B, H, L, D = q.shape
        
        # Process in smaller sub-chunks for memory efficiency
        sub_chunk_size = min(chunk_size, 256)
        
        out = torch.zeros_like(v[:, :, :chunk_size])
        normalizer = torch.zeros((B, H, chunk_size, 1), device=q.device)
        
        # Process key-value pairs in sub-chunks
        for i in range(0, L, sub_chunk_size):
            end_idx = min(i + sub_chunk_size, L)
            
            # Get current key/value sub-chunk
            k_sub = k[:, :, i:end_idx]
            v_sub = v[:, :, i:end_idx]
            
            # Compute attention scores for this sub-chunk
            scores = torch.matmul(q[:, :, :chunk_size], k_sub.transpose(-2, -1))
            scores = scores * self.scale
            
            # Apply sparse attention based on importance scores
            if importance_chunk is not None:
                # Properly reshape importance scores for broadcasting
                imp = importance_chunk.view(B, 1, -1, 1)  # [B, 1, chunk_size, 1]
                imp = imp.expand(-1, H, -1, end_idx - i)  # [B, H, chunk_size, current_chunk_size]
                
                mask = (imp < 0.5)
                if mask.any():
                    scores_masked = scores.masked_fill(~mask, float('-inf'))
                    topk_values, _ = torch.topk(
                        scores_masked,
                        k=min(self.sparse_topk, scores_masked.size(-1)),
                        dim=-1,
                        sorted=False
                    )
                    threshold = topk_values[..., -1:]
                    scores = scores.masked_fill((scores < threshold) & mask, float('-inf'))
            
            # Apply softmax in a memory-efficient way
            scores_max = torch.max(scores, dim=-1, keepdim=True)[0]
            exp_scores = torch.exp(scores - scores_max)
            
            # Update output and normalization factor
            out += torch.matmul(exp_scores, v_sub)
            normalizer += exp_scores.sum(dim=-1, keepdim=True)
            
            # Free memory
            del scores, exp_scores
            torch.cuda.empty_cache()
        
        # Normalize the output
        out = out / (normalizer + 1e-6)
        return out
    
    def forward(self, x, importance_scores):
        B, T, C = x.shape
        
        # Project and split heads
        qkv = self.c_attn(x)
        q, k, v = qkv.chunk(3, dim=-1)
        
        # Reshape to [B, H, T, D]
        q = q.view(B, T, self.n_head, self.head_size).transpose(1, 2)
        k = k.view(B, T, self.n_head, self.head_size).transpose(1, 2)
        v = v.view(B, T, self.n_head, self.head_size).transpose(1, 2)
        
        # Process attention in chunks
        chunk_size = min(T, 128)
        num_chunks = (T + chunk_size - 1) // chunk_size
        
        # Initialize output tensor
        output = torch.zeros_like(x)
        
        for chunk_idx in range(num_chunks):
            start_idx = chunk_idx * chunk_size
            end_idx = min(start_idx + chunk_size, T)
            
            # Get importance scores for current chunk
            imp_chunk = importance_scores[:, start_idx:end_idx]
            
            # Process chunk with mixed precision
            with torch.amp.autocast(device_type='cuda', dtype=torch.float16):
                chunk_output = self._chunk_attention(
                    q[:, :, start_idx:end_idx],
                    k,
                    v,
                    imp_chunk,
                    end_idx - start_idx
                )
            
            # Reshape and store chunk output
            chunk_output = chunk_output.transpose(1, 2).contiguous().view(B, end_idx - start_idx, C)
            output[:, start_idx:end_idx] = chunk_output
            
            # Free memory
            del chunk_output
            torch.cuda.empty_cache()
        
        # Final projection and dropout
        output = self.resid_dropout(self.c_proj(output))
        return output

class Block(nn.Module):
    """
    Transformer block with importance-aware processing
    """
    def __init__(self, config):
        super().__init__()
        self.ln_1 = nn.LayerNorm(config.n_embd)
        self.attn = SparseDenseAttention(config)
        self.ln_2 = nn.LayerNorm(config.n_embd)
        
        self.mlp = nn.Sequential(
            nn.Linear(config.n_embd, 4 * config.n_embd),
            nn.GELU(),
            nn.Linear(4 * config.n_embd, config.n_embd),
            nn.Dropout(config.dropout),
        )
        
        # Feature amplification
        self.feature_gate = nn.Sequential(
            nn.Linear(config.n_embd, config.n_embd),
            nn.Sigmoid()
        )
        
    def forward(self, x, importance_scores):
        # Self-attention with importance awareness
        attn_output = self.attn(self.ln_1(x), importance_scores)
        x = x + attn_output
        
        # Feature amplification based on importance
        gate = self.feature_gate(x)
        x = x * (1 + importance_scores * gate)
        
        # MLP block
        x = x + self.mlp(self.ln_2(x))
        return x

class DTATTransformer(nn.Module):
    """
    Dynamic Token-Aware Transformer (DTAT) for character-level language modeling
    """
    def __init__(self, config):
        super().__init__()
        assert config.vocab_size is not None
        assert config.block_size is not None
        self.config = config
        
        self.transformer = nn.ModuleDict(dict(
            wte = nn.Embedding(config.vocab_size, config.n_embd),
            wpe = nn.Embedding(config.block_size, config.n_embd),
            drop = nn.Dropout(config.dropout),
            h = nn.ModuleList([Block(config) for _ in range(config.n_layer)]),
            ln_f = nn.LayerNorm(config.n_embd)
        ))
        
        # Token importance network
        self.importance_net = TokenImportanceNetwork(config)
        
        # Output head
        self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False)
        
        # Initialize weights
        self.apply(self._init_weights)
        # Apply special scaled init to the residual projections, per GPT-2 paper
        for pn, p in self.named_parameters():
            if pn.endswith('c_proj.weight'):
                torch.nn.init.normal_(p, mean=0.0, std=0.02/math.sqrt(2 * config.n_layer))
        
        # Report number of parameters
        print("number of parameters: %.2fM" % (self.get_num_params()/1e6,))
        
    def get_num_params(self, non_embedding=True):
        """
        Return the number of parameters in the model.
        For non-embedding count (default), the position embeddings get subtracted.
        """
        n_params = sum(p.numel() for p in self.parameters())
        if non_embedding:
            n_params -= self.transformer.wpe.weight.numel()
        return n_params
    
    def _init_weights(self, module):
        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)
    
    def forward(self, idx, targets=None):
        device = idx.device
        b, t = idx.size()
        assert t <= self.config.block_size, f"Cannot forward sequence of length {t}, block size is only {self.config.block_size}"
        pos = torch.arange(0, t, dtype=torch.long, device=device).unsqueeze(0) # shape (1, t)
        
        # Get token embeddings
        tok_emb = self.transformer.wte(idx) # token embeddings of shape (b, t, n_embd)
        pos_emb = self.transformer.wpe(pos) # position embeddings of shape (1, t, n_embd)
        x = self.transformer.drop(tok_emb + pos_emb)
        
        # Create frequency table for importance calculation
        freq_table = idx.clone()  # Use token indices as frequency table for now
        
        # Calculate token importance scores
        importance_scores = self.importance_net(x, freq_table, pos.expand(b, -1))
        
        # Forward through transformer blocks
        for block in self.transformer.h:
            x = block(x, importance_scores)
        x = self.transformer.ln_f(x)
        
        # Get logits and loss
        logits = self.lm_head(x)
        
        loss = None
        if targets is not None:
            # Reshape logits to (B*T, vocab_size)
            B, T, C = logits.shape
            logits = logits.view(B*T, C)
            targets = targets.view(B*T)
            # Calculate loss directly in BPC instead of nats
            loss = F.cross_entropy(logits, targets) / math.log(2)
        
        return logits, loss, importance_scores
    
    @torch.no_grad()
    def generate(self, idx, max_new_tokens, temperature=1.0, top_k=None):
        """
        Take a conditioning sequence of indices idx (LongTensor of shape (b,t)) and complete
        the sequence max_new_tokens times, feeding the predictions back into the model each time.
        Most likely you'll want to make sure to be in model.eval() mode of operation for this.
        """
        for _ in range(max_new_tokens):
            # if the sequence context is growing too long we must crop it at block_size
            idx_cond = idx if idx.size(1) <= self.config.block_size else idx[:, -self.config.block_size:]
            # forward the model to get the logits for the index in the sequence
            logits, _, _ = self(idx_cond)
            # pluck the logits at the final step and scale by desired temperature
            logits = logits[:, -1, :] / temperature
            # optionally crop the logits to only the top k options
            if top_k is not None:
                v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
                logits[logits < v[:, [-1]]] = -float('Inf')
            # apply softmax to convert logits to (normalized) probabilities
            probs = F.softmax(logits, dim=-1)
            # sample from the distribution
            idx_next = torch.multinomial(probs, num_samples=1)
            # append sampled index to the running sequence
            idx = torch.cat((idx, idx_next), dim=1)

        return idx