import torch
import math

def RoPE(batch):
    B, L, D = batch.shape
    x1 = batch[..., 0::2]
    x2 = batch[..., 1::2]

    pos = torch.arange(L, device=batch.device)

    freqs = 10000 ** (-torch.arange(0, D, 2, device=batch.device) / D)

    angles = pos[:, None] * freqs[None, :]

    cos = torch.cos(angles).unsqueeze(0)
    sin = torch.sin(angles).unsqueeze(0)

    x1_rot = x1 * cos - x2 * sin
    x2_rot = x1 * sin + x2 * cos

    return torch.stack([x1_rot, x2_rot], dim=-1).flatten(-2)

class SelfAttention(torch.nn.Module):
    def __init__(self, emb_size, head_size):
        super().__init__()
        self.C = math.sqrt(head_size)
        self.Query = torch.nn.Linear(emb_size, head_size, bias=False)
        self.Key = torch.nn.Linear(emb_size, head_size, bias=False)
        self.Value = torch.nn.Linear(emb_size, head_size, bias=False)

    def forward(self, x, mask=None): # x[batch_size, tok_size, emb_size], 
        query, key, value = self.Query(x), self.Key(x), self.Value(x) #[batch, tok_size, head_size]
        query, key = RoPE(query), RoPE(key)
        
        internal_state = query @ key.transpose(-2, -1) #[batch, tok_size, tok_size]
        if mask is not None: internal_state = internal_state.masked_fill(~mask, float('-inf'))
        weights = torch.nn.functional.softmax(internal_state / self.C, dim=-1) #[batch, tok_size, tok_size]
        out = weights @ value
        return out # [batch, tok_size, head_size]

class MultiHeadSelfAttention(torch.nn.Module):
    def __init__(self, emb_size, n_heads):
        super().__init__()
        head_size = emb_size // n_heads
        self.heads = torch.nn.ModuleList([SelfAttention(emb_size, head_size) for head in range(n_heads)])
        self.Output = torch.nn.Linear(emb_size, emb_size, bias=False)

    def forward(self, x, mask=None):
        temporary_states = [head(x, mask) for head in self.heads]
        temporary_states = torch.cat(temporary_states, dim=-1) # [batch, tok_size, head_size] -> [batch, tok_size, emb_size]
        output = self.Output(temporary_states)
        return output # [batch, tok_size, emb_size]

class FeedForward(torch.nn.Module):
    def __init__(self, emb_size, out_size=None):
        super().__init__()
        neurons = int((8 / 3) * emb_size) - (int((8 / 3) * emb_size) % 64)
        if out_size is None: out_size = emb_size
        self.W1 = torch.nn.Linear(emb_size, neurons, bias=False)
        self.SiLU = torch.nn.SiLU()
        self.W2 = torch.nn.Linear(emb_size, neurons, bias=False)
        self.W3 = torch.nn.Linear(neurons, out_size, bias=False)

    def forward(self, x): # [batch, tok_size, emb_size]
        hidden_state = self.SiLU(self.W1(x)) * self.W2(x) # [batch, tok_size, neurons]
        output = self.W3(hidden_state)
        return output # [batch, tok_size, out_size]

class EncoderBlock(torch.nn.Module):
    def __init__(self, emb_size, n_heads):
        super().__init__()
        self.MultiHeadSelfAttention = MultiHeadSelfAttention(emb_size, n_heads)
        self.RMSNorm1 = torch.nn.RMSNorm(emb_size)
        self.FeedForward = FeedForward(emb_size)
        self.RMSNorm2 = torch.nn.RMSNorm(emb_size)

    def forward(self, context, mask): # [batch_size, context_tok_size, emb_size]

        # [batch_size, context_tok_size, emb_size]
        residual_connection1 = self.MultiHeadSelfAttention(self.RMSNorm1(context), mask) + context
        
        # [batch_size, context_tok_size, emb_size]
        residual_connection2 = self.FeedForward(self.RMSNorm2(residual_connection1)) + residual_connection1

        return residual_connection2 # [batch_size, context_tok_size, emb_size]

class Encoder(torch.nn.Module):
    def __init__(self, n_blocks, emb_size, n_heads):
        super().__init__()
        self.blocks = torch.nn.ModuleList([EncoderBlock(emb_size, n_heads) for block in range(n_blocks)])
        self.RMSNorm = torch.nn.RMSNorm(emb_size)

    def forward(self, context, mask):
        internal_state = context
        for encoder in self.blocks:
            internal_state = encoder(internal_state, mask)
        output = self.RMSNorm(internal_state)
        return output # [batch_size, context_tok_size, emb_size]

class Thinking(torch.nn.Module):
    def __init__(self, emb_size, internal_state_size=128, vocab_size=12288, momentum=0.1):
        super().__init__()
        self.momentum = momentum
        self.to_vocab = torch.nn.Linear(emb_size, vocab_size, bias=False)
        self.alpha = torch.nn.Parameter(torch.ones(vocab_size))
        self.SiLU = torch.nn.SiLU()

        self.to_internal = torch.nn.Linear(vocab_size, internal_state_size, bias=False)
        self.ln = torch.nn.LayerNorm(internal_state_size)
        self.to_emb = torch.nn.Linear(internal_state_size, emb_size, bias=False)

        self.register_buffer("e_mean", torch.zeros(vocab_size))
        self.t_in = torch.nn.Linear(vocab_size, internal_state_size, bias=False)
        self.t_ln = torch.nn.LayerNorm(internal_state_size)
        self.t_out = torch.nn.Linear(internal_state_size, emb_size, bias=False)

    def forward(self, x, mask_for_thinking, context_ids):
        # x [B, Cs, Emb], mask_for_thinking [B, Cs, 1], context_ids [B, Cs]
        x_unit = x / x.norm(dim=-1, keepdim=True).clamp(min=1e-6)
        vocab_unit = self.to_vocab.weight / self.to_vocab.weight.norm(dim=-1, keepdim=True).clamp(min=1e-6)
        energy = x_unit @ vocab_unit.T # [B, Cs, V], from -1 to 1

        own = torch.zeros_like(energy, dtype=torch.bool).scatter_(2, context_ids[..., None], True)
        drop = (~mask_for_thinking) | own # padding + the position's own token
        kept = energy.masked_fill(drop, 0.0) # [B, Cs, V]

        # words channel
        internal_state = kept * self.alpha
        internal_state = self.to_internal(internal_state)
        internal_state = self.ln(internal_state)
        internal_state = self.SiLU(internal_state)
        words = self.to_emb(internal_state) # [B, Cs, Emb]

        # thought channel
        e = kept.sum(dim=1).float() # [B, V]: sum of energies over positions
        if self.training:
            mu = e.mean(dim=0) # mean over texts in the batch
            with torch.no_grad():
                self.e_mean.mul_(1 - self.momentum).add_(self.momentum * mu)
        else:
            mu = self.e_mean # inference: running mean
        e = e - mu # how this text differs from the others
        e = e / e.pow(2).mean(dim=-1, keepdim=True).sqrt().clamp(min=1e-6) # unit scale

        thought = self.t_in(e)
        thought = self.t_ln(thought)
        thought = self.SiLU(thought)
        thought = self.t_out(thought) # [B, Emb]

        thinking_out = words + thought.unsqueeze(1) # thought is added to every position
        return thinking_out

class CrossAttention(torch.nn.Module):
    def __init__(self, emb_size, head_size):
        super().__init__()
        self.C = math.sqrt(head_size)
        self.Query = torch.nn.Linear(emb_size, head_size, bias=False)
        self.Key = torch.nn.Linear(emb_size, head_size, bias=False)
        self.Value = torch.nn.Linear(emb_size, head_size, bias=False)

    def forward(self, data_for_query, data_for_key_value, mask=None):
        # data_for_query[batch, answer_tok_size, emb_size]
        # data_for_key_value[batch, context_tok_size, emb_size] 
        
        query, key, value = self.Query(data_for_query), self.Key(data_for_key_value), self.Value(data_for_key_value)
        # query[batch, answer_tok_size, head_size]
        # key, value[batch, context_tok_size, head_size]
        
        internal_state = query @ key.transpose(-2, -1) #[batch, answer_tok_size, context_tok_size]
        if mask is not None: internal_state = internal_state.masked_fill(~mask, float('-inf'))
        weights = torch.nn.functional.softmax(internal_state / self.C, dim=-1)
        out = weights @ value
        return out #[batch, answer_tok_size, head_size]

class MultiHeadCrossAttention(torch.nn.Module):
    def __init__(self, emb_size, n_heads):
        super().__init__()
        head_size = emb_size // n_heads
        self.heads = torch.nn.ModuleList([CrossAttention(emb_size, head_size) for head in range(n_heads)])
        self.Output = torch.nn.Linear(emb_size, emb_size, bias=False)

    def forward(self, data_for_query, data_for_key_value, mask=None):
        temporary_states = [head(data_for_query, data_for_key_value, mask) for head in self.heads]
        temporary_states = torch.cat(temporary_states, dim=-1) # [batch, answer_tok_size, head_size] -> [batch, answer_tok_size, emb_size]
        output = self.Output(temporary_states)
        return output # [batch, answer_tok_size, emb_size]

class DecoderBlock(torch.nn.Module):
    def __init__(self, emb_size, n_heads):
        super().__init__()
        self.MultiHeadSelfAttention = MultiHeadSelfAttention(emb_size, n_heads)
        self.RMSNorm1 = torch.nn.RMSNorm(emb_size)
        self.MultiHeadCrossAttention = MultiHeadCrossAttention(emb_size, n_heads)
        self.RMSNorm2 = torch.nn.RMSNorm(emb_size)
        self.RMSNorm3 = torch.nn.RMSNorm(emb_size)
        self.FeedForward = FeedForward(emb_size)

    def forward(self, encoder_output, answer, encoder_mask=None):
        
        # [batch, answer_tok_size, emb_size]
        internal_state = self.MultiHeadSelfAttention(self.RMSNorm1(answer))
        residual_connection1 =  internal_state + answer
        internal_state = self.MultiHeadCrossAttention(
            self.RMSNorm2(residual_connection1), encoder_output, encoder_mask) # [batch, answer_tok_size, emb_size]
        residual_connection2 = internal_state + residual_connection1
        internal_state = self.FeedForward(self.RMSNorm3(residual_connection2))
        residual_connection3 = internal_state + residual_connection2 # [batch, answer_tok_size, emb_size]
        return residual_connection3 # [batch, answer_tok_size, emb_size]
        
class Decoder(torch.nn.Module):
    def __init__(self, n_blocks, emb_size, n_heads):
        super().__init__()
        self.blocks = torch.nn.ModuleList([DecoderBlock(emb_size, n_heads) for block in range(n_blocks)])
        self.RMSNorm = torch.nn.RMSNorm(emb_size)
        self.final = torch.nn.Linear(emb_size, 12288, bias=False)

    def forward(self, answer, encoder_output, mask):
        internal_state = answer
        for decoder in self.blocks:
            internal_state = decoder(encoder_output, internal_state, mask) # [batch_size, answer_size, emb_size]
        internal_state = self.RMSNorm(internal_state)
        internal_state = self.final(internal_state)
        return internal_state # [batch_size, answer_size, dict_size]

class NN_Thinking(torch.nn.Module):
    def _init_weights(self, m):
        if isinstance(m, (torch.nn.Linear, torch.nn.Embedding)):
            torch.nn.init.normal_(m.weight, mean=0.0, std=0.02)
        
    def __init__(self, 
                 embedding_size, 
                 encoder_num, 
                 heads_encoder,
                 thinking_internal_state_size, 
                 decoder_num, 
                 heads_decoder):
        super().__init__()
        
        self.Embedding = torch.nn.Embedding(12288, embedding_size)
        self.Encoder = Encoder(encoder_num, embedding_size, heads_encoder)
        self.Thinking = Thinking(embedding_size, thinking_internal_state_size)
        self.Thinking.to_vocab.weight = self.Embedding.weight
        self.out_norm = torch.nn.RMSNorm(embedding_size)
        self.Decoder = Decoder(decoder_num, embedding_size, heads_decoder)
        self.Decoder.final.weight = self.Embedding.weight
        self.apply(self._init_weights)
        torch.nn.init.zeros_(self.Thinking.t_out.weight)

    def encode(self, context_ids):
        # context_ids [B, Cs]
        PAD_ID = 0
        mask = (context_ids != PAD_ID)[:, None, :] # [B, 1, Cs]
        memory = self.Encoder(self.Embedding(context_ids), mask) # [B, Cs, Emb]
        memory = self.Thinking(memory, mask.transpose(1, 2), context_ids)
        memory = self.out_norm(memory)
        return memory, mask
    
    def decode(self, answer_ids, memory, mask):
        return self.Decoder(self.Embedding(answer_ids), memory, mask) # [B, La, V]
    
    def forward(self, context_ids, answer_ids):
        memory, mask = self.encode(context_ids)
        return self.decode(answer_ids, memory, mask)