import math
from dataclasses import dataclass
from typing import Optional, Tuple, List, Union 

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

from transformers import PreTrainedModel, PretrainedConfig
from transformers import AutoConfig, AutoModelForCausalLM, AutoModel


from transformers.modeling_outputs import (
    BaseModelOutput,
    CausalLMOutput,
)

@dataclass
class BaseModelOutputWithAux(BaseModelOutput):
    aux_loss: Optional[torch.Tensor] = None
    one_hot: Optional[torch.Tensor] = None

@dataclass
class CausalLMOutputWithAux(CausalLMOutput):
    aux_loss: Optional[torch.Tensor] = None
    one_hot: Optional[torch.Tensor] = None
    


def create_causal_mask(seq_len: int, device=None, dtype=torch.float32):
    mask = torch.full((seq_len, seq_len), torch.finfo(dtype).min, device=device, dtype=dtype)
    mask = torch.triu(mask, diagonal=1)
    return mask

def create_sw_mask(seq_len: int, window_size: Optional[int], device=None, dtype=torch.float32):
    if not window_size:
        return None
    i = torch.arange(seq_len, device=device)
    j = torch.arange(seq_len, device=device)
    allowed = (j[None, :] <= i[:, None]) & (j[None, :] >= i[:, None] - (window_size - 1))
    mask = torch.zeros((seq_len, seq_len), device=device, dtype=dtype)
    mask = mask.masked_fill(~allowed, torch.finfo(dtype).min)
    return mask
    

class Linear(nn.Module):
    def __init__(self, in_dim: int, out_dim: int, b: bool = False):
        super().__init__()
        self.linear = nn.Linear(in_dim, out_dim, bias=b)

    def forward(self, x: torch.Tensor):
        return self.linear(x)
        

class RMS_norm(nn.Module):
    def __init__(self, dim: int, eps: float = 1e-5, b=True):
        super().__init__()
        self.rms_norm = nn.RMSNorm(dim, eps=eps, elementwise_affine=b)
        
    def forward(self, x: torch.Tensor):
        return self.rms_norm(x)


class GeGLU(nn.Module):
    def __init__(self, in_dim, hidden_dim, b=False):
        super().__init__()        
        self.linear1 = nn.Linear(in_dim, hidden_dim, bias=b)
        self.linear2 = nn.Linear(in_dim, hidden_dim, bias=b)
        self.linear3 = nn.Linear(hidden_dim, in_dim, bias=b)

    def forward(self, x: torch.Tensor):
        gate = F.gelu(self.linear1(x), approximate='tanh')
        return self.linear3(gate * self.linear2(x))


# --- RoPE --- #

def precompute_theta_pos_frequencies(emb_dim: int, seq_len: int, device: str, theta: float = 10000.0):
    assert emb_dim % 2 == 0, "Dimension must be divisible by 2"
    theta_numerator = torch.arange(0, emb_dim, 2, device=device).float()
    theta_values = 1.0 / (theta ** (theta_numerator / emb_dim))
    m = torch.arange(seq_len, device=device)
    freqs_matrix = torch.outer(m, theta_values).float()
    freqs_complex = torch.polar(torch.ones_like(freqs_matrix, device=device), freqs_matrix)
    return freqs_complex

def RoPE(x: torch.Tensor, freqs: torch.Tensor, index: int = -1):
    ndim = x.ndim
    assert 1 < ndim
    assert freqs.shape[-1] == x.shape[-1]
    
    if index != -1:
        # Ensure that index cannot be less than 1 --> (1 == first token position) 
        # and more than freqs.shape[-2] --> (index > max sequence length)
        index = max(1, min(index, freqs.shape[-2]))
        freqs = freqs[(index -1), :].unsqueeze(0)
        
    if freqs.shape[-2] != x.shape[-2] and index == -1:
        # if in x seq_len != max lenght, drop unneeded m frequence positions 
        freqs = freqs[:x.shape[-2], :]
    
    # from (seq_len, head_dim) to (1, 1, seq_len, head_dim)
    freqs = freqs.unsqueeze(0).unsqueeze(0)
    rotated_x = x * freqs
     
    return rotated_x

def apply_rotary_embeddings(x: torch.Tensor, freqs_complex: torch.Tensor):
    x_reshaped = x.reshape(*x.shape[:-1], -1, 2)  
    x_complex = torch.view_as_complex(x_reshaped.float())
    x_rotated = RoPE(x_complex, freqs_complex)
    x_out = torch.view_as_real(x_rotated)
    x_out = x_out.reshape(*x.shape).type_as(x)
    return x_out


class SelfAttention(nn.Module):
    def __init__(self, dim: int, n_heads: int):
        super().__init__()
        assert dim % n_heads == 0, "dim must be divisible by n_heads"
        self.dim = dim
        self.n_heads = n_heads
        self.head_dim = dim // n_heads
        self.Q = nn.Linear(dim, dim)
        self.K = nn.Linear(dim, dim)
        self.V = nn.Linear(dim, dim)
        self.O = nn.Linear(dim, dim)

    def forward(self, x: torch.Tensor, *, causal_mask: Optional[torch.Tensor] = None, key_padding_mask: Optional[torch.Tensor] = None, freqs: Optional[torch.Tensor] = None):
        # x: [B,S,D]
        B, S, D = x.shape
        H = self.n_heads
        Dh = self.head_dim

        q = self.Q(x).view(B, S, H, Dh).transpose(1, 2)  # [B,H,S,Dh]
        k = self.K(x).view(B, S, H, Dh).transpose(1, 2)
        v = self.V(x).view(B, S, H, Dh).transpose(1, 2)

        q = apply_rotary_embeddings(q, freqs)
        k = apply_rotary_embeddings(k, freqs)

        attn = (q @ k.transpose(-2, -1)) / math.sqrt(Dh)  # [B,H,S,S]
        neg_inf = torch.finfo(attn.dtype).min

        if causal_mask is not None:
            attn = attn + causal_mask[None, None, :, :]

        if key_padding_mask is not None:  # True at PAD keys
            attn = attn.masked_fill(key_padding_mask[:, None, None, :], neg_inf)

        attn = F.softmax(attn, dim=-1)
        out = (attn @ v).transpose(1, 2).contiguous().view(B, S, D)

        if key_padding_mask is not None:
            query_mask = (~key_padding_mask).to(out.dtype).unsqueeze(-1)  # [B,S,1]
            out = out * query_mask

        return self.O(out)


class Layer_Block(nn.Module):
    def __init__(self, dim: int, n_heads: int):
        super().__init__()
        self.norm_1 = RMS_norm(dim)
        self.norm_2 = RMS_norm(dim)
        self.attention = SelfAttention(dim, n_heads)
        self.hidden_dim = int(dim * 4)
        self.feedforward = GeGLU(dim, self.hidden_dim)
    
    def forward(self, x: torch.Tensor, *, causal_mask=None, key_padding_mask=None, freqs=None):
        residual = x
        x = self.norm_1(x)
        x = self.attention(x, causal_mask=causal_mask, key_padding_mask=key_padding_mask, freqs=freqs)
        x = residual + x

        residual = x
        x = self.norm_2(x)
        x = self.feedforward(x)
        x = residual + x
        return x


class Top_k_Router(nn.Module):
    def __init__(self, dim: int, n_experts: int, k: int):
        super().__init__()
        self.n_experts = n_experts
        self.k = k
        self.router = Linear(dim, n_experts)

    def forward(self, x: torch.Tensor):
        # x: [T, D]
        logits = self.router(x)
        probs = F.softmax(logits, dim=-1)
        topk_prob, topk_idx = torch.topk(probs, k=self.k, dim=-1)
        topk_prob = topk_prob / (topk_prob.sum(dim=-1, keepdim=True) + 1e-9)

        p_mean = probs.mean(dim=0)
        T, E = probs.size()
        one_hot = F.one_hot(topk_idx, num_classes=E).sum(dim=1)
        f_mean = one_hot.float().mean(dim=0) / float(self.k)
        aux_loss = (E * (p_mean * f_mean).sum())

        return topk_idx, topk_prob, aux_loss, one_hot


class MoE_Block(nn.Module):
    def __init__(self, dim_in: int, n_experts: int, k: int, dim_hidden: int, dim_out: Optional[int] = None):
        super().__init__()
        self.n_experts = n_experts
        self.k = k
        self.dim_out = dim_in if dim_out is None else dim_out
        self.router = Top_k_Router(dim_in, n_experts, k)
        self.experts = nn.ModuleList([nn.Linear(dim_in, self.dim_out, bias=True) for _ in range(n_experts)])

    def forward(self, x: torch.Tensor):
        # x: [B,S,D]
        B, S, D = x.shape
        T = B * S
        x_flat = x.view(T, D)

        topk_idx, topk_prob, aux_loss, one_hot = self.router(x_flat)  # [T,k], [T,k]
        y_flat = x_flat.new_zeros((T, self.dim_out))

        for slot in range(self.k):
            e_ids = topk_idx[:, slot]
            e_prob = topk_prob[:, slot].unsqueeze(-1)  # [T,1]
            for e in range(self.n_experts):
                idx = (e_ids == e).nonzero(as_tuple=False).squeeze(1)
                if idx.numel() == 0:
                    continue
                y_e = self.experts[e](x_flat.index_select(0, idx))  # [Ne, dim_out]
                y_flat.index_add_(0, idx, y_e * e_prob.index_select(0, idx))

        return y_flat.view(B, S, self.dim_out), aux_loss, one_hot.view(B, S, self.n_experts)


class ParallelLayer(nn.Module):
    def __init__(self, dim: int, n_heads: int, n_blocks: int, k: int):
        super().__init__()
        self.dim = dim
        self.n_blocks = n_blocks
        self.blocks = nn.ModuleList([Layer_Block(dim, n_heads) for _ in range(n_blocks)])
        self.router = Top_k_Router(dim, n_blocks, k)

    def forward(self, x: torch.Tensor, *, causal_mask=None, key_padding_mask=None, freqs=None):
        B, S, D = x.shape
        T = B * S
        outs = [blk(x, causal_mask=causal_mask, key_padding_mask=key_padding_mask, freqs=freqs) for blk in self.blocks]
        
        Y = torch.stack([y.view(T, D) for y in outs], dim=1)  # [T, n_blocks, D]

        x_flat = x.view(T, D)
        
        topk_idx, topk_prob, aux_loss, one_hot = self.router(x_flat)

        gather_idx = topk_idx.unsqueeze(-1).expand(-1, -1, D)  # [T,k,D]
        Y_topk = torch.gather(Y, dim=1, index=gather_idx)      # [T,k,D]
        y_flat = (Y_topk * topk_prob.unsqueeze(-1)).sum(dim=1) # [T,D]

        return y_flat.view(B, S, D), aux_loss, one_hot.view(B, S, self.n_blocks)


class MoEPgemmaLMConfig(PretrainedConfig):
    model_type = "moep_gemma_lm"

    def __init__(
        self,
        vocab_size: int = 16384,
        bos_token_id = 1,
        eos_token_id = 2,
        pad_token_id = 3,
        mask_token_id = 4,
        max_context_length: int = 1024,
        n_layers: int = 12,
        ffn_multp: int = 4,
        dim: int = 348,
        n_heads: int = 6,
        # Parallel specifics
        n_parallel_blocks: int = 4,
        parallel_dim: int = 192,
        parallel_n_heads: int = 3,
        parallel_k: int = 2,
        sliding_window: Optional[int] = None,
        # MoE specifics
        n_experts: int = 4,
        moe_k: int = 2,
        # Norms
        layer_norm_epsilon: float = 1e-5,
        tie_word_embeddings: bool = True,
        **kwargs,
    ):
        super().__init__(
            vocab_size=vocab_size,
            pad_token_id=pad_token_id,
            bos_token_id=bos_token_id,
            eos_token_id=eos_token_id,
            **kwargs,
        )

        # Special tokens
        self.mask_token_id = mask_token_id

        # Core dims
        self.vocab_size = vocab_size
        self.max_context_length = max_context_length
        self.n_layers = n_layers
        self.ffn_multp = ffn_multp
        self.dim = dim
        self.n_heads = n_heads

        # Parallel specifics
        self.n_parallel_blocks = n_parallel_blocks
        self.parallel_dim = parallel_dim
        self.parallel_n_heads = parallel_n_heads
        self.parallel_k = parallel_k
        self.sliding_window = sliding_window

        # MoE specifics
        self.n_experts = n_experts
        self.moe_k = moe_k

        # norms
        self.layer_norm_epsilon = layer_norm_epsilon

        # Tying
        self.tie_word_embeddings = tie_word_embeddings

class MoEPgemmaModel(PreTrainedModel):
    config_class = MoEPgemmaLMConfig
    base_model_prefix = "model"

    def __init__(self, config: MoEPgemmaLMConfig):
        super().__init__(config)
        self.config = config

        # --- embeddings ---
        self.embed_tokens = nn.Embedding(config.vocab_size, config.dim, padding_idx=config.pad_token_id)
     
        self.layers = nn.ModuleList()

        # Large Layer
        self.layers.append(Layer_Block(config.dim, 
                                       config.n_heads))

        # MoE Linear
        self.layers.append(MoE_Block(config.dim, 
                                     config.n_experts, 
                                     config.moe_k, 
                                     dim_hidden=config.dim * 2,
                                     dim_out=config.parallel_dim))

        # Parallel Layers
        for _ in range(config.n_layers - 2):
            self.layers.append(ParallelLayer(config.parallel_dim, 
                                             config.parallel_n_heads,
                                             config.n_parallel_blocks, 
                                             config.parallel_k))

        # MoE Linear
        self.layers.append(MoE_Block(config.parallel_dim, 
                                     config.n_experts, 
                                     config.moe_k, 
                                     dim_hidden=config.parallel_dim * config.ffn_multp,
                                     dim_out=config.dim))
        
        # Large Layer
        self.layers.append(Layer_Block(config.dim, 
                                       config.n_heads))

        self.norm = RMS_norm(config.dim)
        self.post_init()

        self.sliding_window = config.sliding_window

        self.freqs_p = None

        self.freqs = None


    def get_input_embeddings(self):
        return self.embed_tokens

    def set_input_embeddings(self, new_embeddings):
        self.embed_tokens = new_embeddings

    def _init_weights(self, module):
        if isinstance(module, (nn.Linear, nn.Embedding)):
            nn.init.normal_(module.weight, mean=0.0, std=0.02)
        if isinstance(module, nn.Linear) and module.bias is not None:
            nn.init.zeros_(module.bias)

    def forward(
        self,
        input_ids: torch.LongTensor,
        attention_mask: Optional[torch.Tensor] = None,
        return_dict: bool = True,
        return_aux: Optional[bool] = False,
        return_aux_dict: Optional[bool] = False,
        **unused,
    ) -> Union[BaseModelOutputWithAux, BaseModelOutput, Tuple]:

        device = input_ids.device
        B, S = input_ids.shape

        if self.freqs is None or self.freqs.device != device:
            self.freqs = precompute_theta_pos_frequencies(
                int(self.config.dim / self.config.n_heads),
                self.config.max_context_length,
                device=device
            )
        
        if self.freqs_p is None or self.freqs_p.device != device:
            self.freqs_p = precompute_theta_pos_frequencies(
                int(self.config.parallel_dim / self.config.parallel_n_heads),
                self.config.max_context_length,
                device=device
            )

        key_padding = input_ids.eq(self.config.pad_token_id) if attention_mask is None else (attention_mask == 0)

        # embeddings
        x = self.embed_tokens(input_ids)

        mask_dtype = x.dtype
        causal = create_causal_mask(S, device=device, dtype=mask_dtype)
        sw = create_sw_mask(S, self.sliding_window, device=self.device, dtype=mask_dtype)
        sw_mask = None if sw is None else torch.minimum(causal, sw)
        
      
        aux_total = None
        all_one_hots = []
        
        
        freqs = self.freqs

        # First large
        x = self.layers[0](x, causal_mask=causal, key_padding_mask=key_padding, freqs=freqs)

        # MoE shrink
        x, aux, one_hot = self.layers[1](x)
 
        if aux is not None:
            aux_total = aux

        if one_hot is not None:
            all_one_hots.append(one_hot)

        freqs_p = self.freqs_p
        
        idx = 2
        for i in range(self.config.n_layers - 2):
            
            use_sw = (self.sliding_window is not None) and (i % 2 == 1)
            mask = sw_mask if (use_sw and sw_mask is not None) else causal
            
            x, aux, one_hot = self.layers[idx](x, causal_mask=mask, key_padding_mask=key_padding, freqs=freqs_p)
            if aux is not None:
                aux_total += aux
            if one_hot is not None:
                all_one_hots.append(one_hot)
            idx += 1

        # MoE grow
        x, aux, one_hot = self.layers[idx](x)
        
        if aux is not None:
                aux_total += aux

        if one_hot is not None:
            all_one_hots.append(one_hot)
            
        idx += 1

        # Second large
        x = self.layers[idx](x, causal_mask=causal, key_padding_mask=key_padding, freqs=freqs)

        last_hidden = self.norm(x)

        if not return_dict:
            return (last_hidden, aux_total) if return_aux else (last_hidden,)

        if return_aux_dict:
            return BaseModelOutputWithAux(
                last_hidden_state=last_hidden,
                hidden_states=None,
                attentions=None,
                aux_loss=aux_total,
                one_hot=all_one_hots,
            )
        else:
            return BaseModelOutput(
                last_hidden_state=last_hidden,
                hidden_states=None,
                attentions=None,
            )


class MoEPgemmaForCausalLM(PreTrainedModel):
    config_class = MoEPgemmaLMConfig
    base_model_prefix = "model"

    def __init__(self, config: MoEPgemmaLMConfig):
        super().__init__(config)
        self.config = config
        self.model = MoEPgemmaModel(config)
        self.lm_head = nn.Linear(config.dim, config.vocab_size, bias=False)
        self.post_init()
        if config.tie_word_embeddings:
            self.lm_head.weight = self.model.embed_tokens.weight

    def get_input_embeddings(self):
        return self.model.get_input_embeddings()

    def set_input_embeddings(self, new_embeddings):
        self.model.set_input_embeddings(new_embeddings)
 
    def _init_weights(self, module):
        if isinstance(module, (nn.Linear, nn.Embedding)):
            nn.init.normal_(module.weight, mean=0.0, std=0.02)
        if isinstance(module, nn.Linear) and module.bias is not None:
            nn.init.zeros_(module.bias)
            
    def forward(
        self,
        input_ids: torch.LongTensor,
        attention_mask: Optional[torch.Tensor] = None,
        labels: Optional[torch.LongTensor] = None,
        return_dict: bool = True,
        return_aux: Optional[bool] = False,
        aux_loss_weight: Optional[float] = 0.2,
        **unused,
    ) -> Union[CausalLMOutput, Tuple]:
    

        if return_aux:
            base_out = self.model(input_ids=input_ids, 
                                  attention_mask=attention_mask, 
                                  return_aux_dict=True,
                                  return_aux=True)
        else: 
            base_out = self.model(input_ids=input_ids, 
                                  attention_mask=attention_mask)
    
        logits = self.lm_head(base_out.last_hidden_state).float()

        loss = None
        
        if labels is not None:
            
            shift_logits = logits[:, :-1, :].contiguous()
            shift_labels = labels[:, 1:].contiguous()

            pad_id = self.config.pad_token_id
            if pad_id is not None:
                shift_labels = shift_labels.masked_fill(shift_labels == pad_id, -100)

            loss = F.cross_entropy(
                shift_logits.view(-1, shift_logits.size(-1)),
                shift_labels.view(-1),
                ignore_index=-100,
            )

            aux = getattr(base_out, "aux_loss", None)
            if aux is not None:
                loss = loss + aux_loss_weight * aux

        if return_aux:
            return CausalLMOutputWithAux(
                loss=loss,
                logits=logits,
                hidden_states=base_out.hidden_states,
                attentions=base_out.attentions,
                aux_loss=getattr(base_out, "aux_loss", None),
                one_hot=getattr(base_out, "one_hot", None),
            )
        else:
            return CausalLMOutput(
                loss=loss,
                logits=logits,
                hidden_states=base_out.hidden_states,
                attentions=base_out.attentions,
            )

AutoConfig.register("moep_gemma_lm", MoEPgemmaLMConfig)
AutoModelForCausalLM.register(MoEPgemmaLMConfig, MoEPgemmaForCausalLM)
AutoModel.register(MoEPgemmaLMConfig, MoEPgemmaModel)