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

class PositionalEncoding(nn.Module):
    def __init__(self, d_model, max_len=5000):
        super().__init__()
        pe = torch.zeros(max_len, d_model)
        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
        div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))
        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)
        self.register_buffer('pe', pe.unsqueeze(0))

    def forward(self, x):
        # x shape: (batch_size, seq_len, d_model)
        return x + self.pe[:, :x.size(1)]


class ManasRouter(nn.Module):
    """
    The Surface Layer (Manas).
    Fast, shallow processing. Decides if a query needs deep reasoning.
    """
    def __init__(self, vocab_size, d_model, n_heads, n_layers=2):
        super().__init__()
        self.embedding = nn.Embedding(vocab_size, d_model)
        self.pos_encoder = PositionalEncoding(d_model)
        
        encoder_layer = nn.TransformerEncoderLayer(
            d_model=d_model, 
            nhead=n_heads, 
            dim_feedforward=d_model*4,
            batch_first=True
        )
        self.transformer_encoder = nn.TransformerEncoder(encoder_layer, num_layers=n_layers)
        
        # Router head to determine if we need Mahat/Buddhi (0 to 1)
        self.router_head = nn.Linear(d_model, 1)

    def forward(self, x_idx):
        # x_idx: (batch, seq_len)
        emb = self.embedding(x_idx)
        emb = self.pos_encoder(emb)
        manas_state = self.transformer_encoder(emb) # (batch, seq_len, d_model)
        
        # Pool the sequence to make a routing decision (e.g. using the mean)
        pooled = manas_state.mean(dim=1) # (batch, d_model)
        routing_prob = torch.sigmoid(self.router_head(pooled)) # (batch, 1)
        
        return manas_state, routing_prob


class MahatMemoryInterface(nn.Module):
    """
    The Memory Layer (Mahat).
    A non-parametric (or frozen/separately updated) associative memory.
    Mocked here as a trainable memory bank that the query attends to.
    """
    def __init__(self, d_model, memory_size=1024):
        super().__init__()
        # In a real model, this is an interface to graph_manifest.json (RAG)
        self.memory_bank = nn.Parameter(torch.randn(memory_size, d_model))
        
        # Multi-head attention to query the memory bank
        self.cross_attn = nn.MultiheadAttention(embed_dim=d_model, num_heads=4, batch_first=True)
        
    def forward(self, query_state):
        # query_state: (batch, seq_len, d_model)
        batch_size = query_state.size(0)
        
        # Expand memory bank for the batch: (batch, memory_size, d_model)
        memory = self.memory_bank.unsqueeze(0).expand(batch_size, -1, -1)
        
        # Query the memory bank
        # query is manas_state, key/value is memory_bank
        context_state, _ = self.cross_attn(query_state, memory, memory)
        
        return context_state


class BuddhiSynthesizer(nn.Module):
    """
    The Reasoning Layer (Buddhi).
    A dense transformer block that activates for synthesis, utilizing Mahat's context.
    """
    def __init__(self, vocab_size, d_model, n_heads, n_layers=4):
        super().__init__()
        decoder_layer = nn.TransformerDecoderLayer(
            d_model=d_model,
            nhead=n_heads,
            dim_feedforward=d_model*4,
            batch_first=True
        )
        self.transformer_decoder = nn.TransformerDecoder(decoder_layer, num_layers=n_layers)
        self.fc_out = nn.Linear(d_model, vocab_size)
        
    def forward(self, tgt_emb, memory_context, tgt_mask=None):
        # tgt_emb: (batch, tgt_seq_len, d_model)
        # memory_context: output from Mahat (batch, src_seq_len, d_model)
        
        out_state = self.transformer_decoder(
            tgt=tgt_emb,
            memory=memory_context,
            tgt_mask=tgt_mask
        )
        
        logits = self.fc_out(out_state) # (batch, tgt_seq_len, vocab_size)
        return logits


class VedicMindLLM(nn.Module):
    """
    The complete Vedic Mind Architecture: Manas -> Mahat -> Buddhi.
    """
    def __init__(self, vocab_size, d_model=512, n_heads=8):
        super().__init__()
        self.vocab_size = vocab_size
        self.d_model = d_model
        
        self.manas = ManasRouter(vocab_size, d_model, n_heads, n_layers=2)
        self.mahat = MahatMemoryInterface(d_model, memory_size=2048)
        self.buddhi = BuddhiSynthesizer(vocab_size, d_model, n_heads, n_layers=6)
        
    def generate_square_subsequent_mask(self, sz):
        mask = (torch.triu(torch.ones(sz, sz)) == 1).transpose(0, 1)
        mask = mask.float().masked_fill(mask == 0, float('-inf')).masked_fill(mask == 1, float(0.0))
        return mask
        
    def forward(self, src_idx, tgt_idx):
        # 1. Surface Processing (Manas)
        manas_state, routing_prob = self.manas(src_idx)
        
        # 2. Context Retrieval (Mahat)
        # In this POC, we always retrieve context. In a full version, 
        # we might gate this using `routing_prob` to save compute.
        context_state = self.mahat(manas_state)
        
        # 3. Synthesis (Buddhi)
        # We need to embed the target tokens for teacher forcing in training
        tgt_emb = self.manas.embedding(tgt_idx)
        tgt_emb = self.manas.pos_encoder(tgt_emb)
        
        tgt_mask = self.generate_square_subsequent_mask(tgt_idx.size(1)).to(tgt_idx.device)
        
        logits = self.buddhi(tgt_emb, context_state, tgt_mask=tgt_mask)
        return logits, routing_prob