import torch
import torch.nn as nn
import torch.nn.functional as F
import math
from embeddings import TechEmbeddingLayer, create_padding_mask, create_causal_mask
class MultiHeadAttention(nn.Module):
    """Multi-head attention mechanism optimized for technical content"""    
    def __init__(self, d_model, n_heads, dropout=0.1):
        super(MultiHeadAttention, self).__init__()
        assert d_model % n_heads == 0        
        self.d_model = d_model
        self.n_heads = n_heads
        self.d_k = d_model // n_heads        
        self.w_q = nn.Linear(d_model, d_model, bias=False)
        self.w_k = nn.Linear(d_model, d_model, bias=False)
        self.w_v = nn.Linear(d_model, d_model, bias=False)
        self.w_o = nn.Linear(d_model, d_model)        
        self.dropout = nn.Dropout(dropout)
        self._init_weights()    
    def _init_weights(self):
        """Initialize weights with Xavier uniform"""
        for module in [self.w_q, self.w_k, self.w_v, self.w_o]:
            nn.init.xavier_uniform_(module.weight)    
    def forward(self, query, key, value, mask=None, pos_encoding=None):
        batch_size, seq_len, d_model = query.size()        
        Q = self.w_q(query).view(batch_size, seq_len, self.n_heads, self.d_k).transpose(1, 2)
        K = self.w_k(key).view(batch_size, seq_len, self.n_heads, self.d_k).transpose(1, 2)
        V = self.w_v(value).view(batch_size, seq_len, self.n_heads, self.d_k).transpose(1, 2)        
        if pos_encoding is not None:
            Q, K = pos_encoding(Q, K)        
        scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.d_k)        
        if mask is not None:
            mask = mask.unsqueeze(1).expand(batch_size, self.n_heads, seq_len, seq_len)
            scores.masked_fill_(mask, float('-inf'))        
        attention_weights = F.softmax(scores, dim=-1)
        attention_weights = self.dropout(attention_weights)
        attended = torch.matmul(attention_weights, V)
        attended = attended.transpose(1, 2).contiguous().view(batch_size, seq_len, d_model)
        output = self.w_o(attended)        
        return output
class FeedForward(nn.Module):
    """Position-wise feed forward network with GELU activation"""    
    def __init__(self, d_model, dim_feedforward, dropout=0.1):
        super(FeedForward, self).__init__()
        self.linear1 = nn.Linear(d_model, dim_feedforward)
        self.linear2 = nn.Linear(dim_feedforward, d_model)
        self.dropout = nn.Dropout(dropout)        
        nn.init.xavier_uniform_(self.linear1.weight)
        nn.init.xavier_uniform_(self.linear2.weight)    
    def forward(self, x):
        x = F.gelu(self.linear1(x))
        x = self.dropout(x)
        x = self.linear2(x)
        return x
class RecursionRouter(nn.Module):
    """Router to decide recursion steps for different types of technical problems"""    
    def __init__(self, d_model, max_steps=4, router_type="adaptive"):
        super(RecursionRouter, self).__init__()
        self.max_steps = max_steps
        self.router_type = router_type        
        if router_type == "adaptive":
            self.complexity_classifier = nn.Sequential(
                nn.Linear(d_model, d_model // 4),
                nn.GELU(),
                nn.Dropout(0.1),
                nn.Linear(d_model // 4, max_steps + 1),
                nn.Softmax(dim=-1)
            )
        elif router_type == "fixed":
            self.fixed_steps = max_steps    
    def forward(self, x):
        if self.router_type == "adaptive":
            seq_repr = x.mean(dim=1)
            step_probs = self.complexity_classifier(seq_repr)
            steps = torch.argmax(step_probs, dim=-1)
            return steps
        return self.fixed_steps
class RecursiveTransformerLayer(nn.Module):
    """Transformer layer with recursive computation capability"""    
    def __init__(self, d_model, n_heads, dim_feedforward, max_steps=4, 
                 dropout=0.1, router_type="adaptive"):
        super(RecursiveTransformerLayer, self).__init__()        
        self.max_steps = max_steps
        self.d_model = d_model        
        self.attention = MultiHeadAttention(d_model, n_heads, dropout)
        self.feedforward = FeedForward(d_model, dim_feedforward, dropout)
        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)
        self.dropout = nn.Dropout(dropout)
        self.router = RecursionRouter(d_model, max_steps, router_type)
        self.step_projections = nn.ModuleList([
            nn.Linear(d_model, d_model) for _ in range(max_steps)
        ])    
    def forward(self, x, mask=None, pos_encoding=None):
        steps = self.router(x)
        if isinstance(steps, int):
            num_steps = min(steps, self.max_steps)
            return self._recursive_forward_fixed(x, mask, num_steps, pos_encoding)
        return self._recursive_forward_adaptive(x, mask, steps, pos_encoding)    
    def _recursive_forward_fixed(self, x, mask, num_steps, pos_encoding):
        device = x.device
        batch_size = x.shape[0]
        computation_loss = torch.tensor(0.0, device=device)        
        for step in range(num_steps):
            step_input = self.step_projections[step](x) if step < len(self.step_projections) else x
            attended = self.attention(step_input, step_input, step_input, mask, pos_encoding)
            x = self.norm1(x + self.dropout(attended))
            fed_forward = self.feedforward(x)
            x = self.norm2(x + self.dropout(fed_forward))
            computation_loss += torch.tensor(0.1, device=device) * batch_size        
        return x, computation_loss    
    def _recursive_forward_adaptive(self, x, mask, steps, pos_encoding):
        batch_size, seq_len, d_model = x.shape
        device = x.device
        max_batch_steps = int(steps.max().item())
        computation_loss = torch.tensor(0.0, device=device)        
        active_batches = torch.ones(batch_size, device=device, dtype=torch.bool)
        for step in range(max_batch_steps):
            step_mask = (steps > step) & active_batches
            if not step_mask.any():
                break                
            step_input = self.step_projections[step](x) if step < len(self.step_projections) else x
            attended = self.attention(step_input, step_input, step_input, mask, pos_encoding)
            attended = torch.where(step_mask.unsqueeze(-1).unsqueeze(-1), attended, torch.zeros_like(attended))
            x = self.norm1(x + self.dropout(attended))
            fed_forward = self.feedforward(x)
            fed_forward = torch.where(step_mask.unsqueeze(-1).unsqueeze(-1), fed_forward, torch.zeros_like(fed_forward))
            x = self.norm2(x + self.dropout(fed_forward))
            computation_loss += torch.tensor(0.1, device=device) * step_mask.sum()
            active_batches &= (steps > step)        
        return x, computation_loss
class MixtureOfRecursions(nn.Module):
    """Main model with mixture of recursive transformer layers"""    
    def __init__(self, vocab_size, d_model=512, n_layers=6, n_heads=8, 
                 max_steps=4, dim_feedforward=2048, dropout=0.1, 
                 max_seq_len=512, router_type="adaptive", padding_idx=0):
        super(MixtureOfRecursions, self).__init__()        
        self.d_model = d_model
        self.vocab_size = vocab_size
        self.padding_idx = padding_idx        
        self.embeddings = TechEmbeddingLayer(
            vocab_size=vocab_size,
            d_model=d_model,
            max_seq_len=max_seq_len,
            dropout=dropout,
            padding_idx=padding_idx,
            pos_encoding="learned"
        )        
        self.layers = nn.ModuleList([
            RecursiveTransformerLayer(
                d_model=d_model,
                n_heads=n_heads,
                dim_feedforward=dim_feedforward,
                max_steps=max_steps,
                dropout=dropout,
                router_type=router_type
            ) for _ in range(n_layers)
        ])        
        self.final_norm = nn.LayerNorm(d_model)
        self.lm_head = nn.Linear(d_model, vocab_size, bias=False)
        self._init_weights()    
    def _init_weights(self):
        nn.init.xavier_uniform_(self.lm_head.weight)    
    def forward(self, input_ids, attention_mask=None):
        batch_size, seq_len = input_ids.shape        
        padding_mask = create_padding_mask(input_ids, self.padding_idx) if attention_mask is None else (attention_mask == 0)
        causal_mask = create_causal_mask(seq_len, input_ids.device)
        padding_mask = padding_mask.unsqueeze(1).expand(batch_size, seq_len, seq_len)
        combined_mask = padding_mask | causal_mask.unsqueeze(0)        
        x = self.embeddings(input_ids)
        pos_encoding = self.embeddings.get_positional_encoding()        
        device = x.device
        total_computation_loss = torch.tensor(0.0, device=device)
        for layer in self.layers:
            x, comp_loss = layer(x, combined_mask, pos_encoding)
            total_computation_loss += comp_loss        
        x = self.final_norm(x)
        logits = self.lm_head(x)        
        return logits, total_computation_loss    
    def generate_step(self, input_ids, temperature=1.0, top_k=None, top_p=None):
        self.eval()
        with torch.no_grad():
            logits, _ = self.forward(input_ids)
            last_logits = logits[:, -1, :] / temperature            
            if top_k is not None:
                indices_to_remove = last_logits < torch.topk(last_logits, top_k)[0][..., -1, None]
                last_logits[indices_to_remove] = float('-inf')            
            if top_p is not None:
                sorted_logits, sorted_indices = torch.sort(last_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)
                last_logits[indices_to_remove] = float('-inf')            
            probs = F.softmax(last_logits, dim=-1)
            next_token = torch.multinomial(probs, num_samples=1)
            return next_token
class TextGenerator:
    """Text generation utility for the tech model"""    
    def __init__(self, model, tokenizer, max_length=100, device=None):
        self.model = model
        self.tokenizer = tokenizer
        self.max_length = max_length
        self.device = device if device else next(model.parameters()).device
        self.model.to(self.device)
        self.eos_token_id = tokenizer.vocab.get('<|endoftext|>', -1)
        self.assistant_token_id = tokenizer.vocab.get('<|assistant|>', -1)    
    def generate(self, prompt, method="nucleus", temperature=1.0, top_k=50, top_p=0.9, max_new_tokens=None):
        if max_new_tokens is None:
            max_new_tokens = self.max_length        
        input_text = f"<|user|> {prompt}"
        input_ids = self.tokenizer.encode_ids(input_text, add_special_tokens=True)
        input_tensor = torch.tensor([input_ids], device=self.device)        
        self.model.eval()
        generated_ids = []        
        with torch.no_grad():
            for _ in range(max_new_tokens):
                if input_tensor.size(1) > self.max_length:
                    input_tensor = input_tensor[:, -self.max_length:]                
                # Generate next token
                if method == "greedy":
                    next_token = self._greedy_generate(input_tensor)
                elif method == "sample":
                    next_token = self._sample_generate(input_tensor, temperature)
                elif method == "top_k":
                    next_token = self._top_k_generate(input_tensor, temperature, top_k)
                elif method == "nucleus" or method == "top_p":
                    next_token = self._nucleus_generate(input_tensor, temperature, top_p)
                else:
                    raise ValueError(f"Unknown generation method: {method}")                
                next_token_id = next_token.item()
                generated_ids.append(next_token_id)
                input_tensor = torch.cat([input_tensor, next_token.unsqueeze(0)], dim=1)
                if next_token_id == self.eos_token_id or (self.assistant_token_id != -1 and next_token_id == self.assistant_token_id):
                    break        
        # Decode the full sequence
        full_ids = input_ids + generated_ids
        full_text = self.tokenizer.decode_ids(full_ids, skip_special_tokens=False)        
        # Extract assistant response
        if "<|assistant|>" in full_text:
            response = full_text.split("<|assistant|>")[-1].split("<|endoftext|>")[0].strip()
        else:
            response = full_text.split("<|endoftext|>")[0].strip()        
        return response if response else "No response generated."    
    def _greedy_generate(self, input_tensor):
        logits, _ = self.model(input_tensor)
        return torch.argmax(logits[:, -1, :], dim=-1, keepdim=True)    
    def _sample_generate(self, input_tensor, temperature):
        logits, _ = self.model(input_tensor)
        logits = logits[:, -1, :] / temperature
        probs = F.softmax(logits, dim=-1)
        return torch.multinomial(probs, num_samples=1)   
    def _top_k_generate(self, input_tensor, temperature, top_k):
        logits, _ = self.model(input_tensor)
        logits = logits[:, -1, :] / temperature
        top_k_logits, top_k_indices = torch.topk(logits, top_k)
        probs = F.softmax(top_k_logits, dim=-1)
        next_token_idx = torch.multinomial(probs, num_samples=1)
        return top_k_indices.gather(-1, next_token_idx)    
    def _nucleus_generate(self, input_tensor, temperature, top_p):
        return self.model.generate_step(input_tensor, temperature, top_p=top_p)
def count_parameters(model):
    total_params = sum(p.numel() for p in model.parameters())
    trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
    return total_params, trainable_params
def main():
    vocab_size = 10000
    d_model = 512
    n_layers = 6
    n_heads = 8
    seq_len = 128
    batch_size = 4    
    print("Initializing MixtureOfRecursions model...")
    model = MixtureOfRecursions(
        vocab_size=vocab_size,
        d_model=d_model,
        n_layers=n_layers,
        n_heads=n_heads,
        max_steps=4,
        dim_feedforward=2048,
        dropout=0.1,
        router_type="adaptive"
    )    
    total_params, trainable_params = count_parameters(model)
    print(f"Total parameters: {total_params:,}")
    print(f"Trainable parameters: {trainable_params:,}")    
    print("\nTesting forward pass...")
    input_ids = torch.randint(0, vocab_size, (batch_size, seq_len))
    attention_mask = torch.ones_like(input_ids)
    attention_mask[:, -10:] = 0    
    print(f"Input shape: {input_ids.shape}")
    logits, comp_loss = model(input_ids, attention_mask)    
    print(f"Output logits shape: {logits.shape}")
    print(f"Computation loss: {comp_loss}")
    print(f"Expected logits shape: ({batch_size}, {seq_len}, {vocab_size})")    
    print("\nTesting generation step...")
    next_token = model.generate_step(input_ids[:1], temperature=0.8, top_p=0.9)
    print(f"Generated next token: {next_token}")    
    print("\nModel test completed successfully!")
if __name__ == "__main__":
    main()