import math
import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers import PreTrainedModel
from transformers.modeling_outputs import CausalLMOutput
from .configuration_dumbmini import DumbMiniConfig

class DumbLinearNext(nn.Module):
    def __init__(self, in_dim, out_dim):
        super().__init__()
        self.main_proj = nn.Linear(in_dim, out_dim)
        self.gate_proj = nn.Linear(in_dim, out_dim)
        self.extra_bias = nn.Parameter(torch.zeros(out_dim))

    def forward(self, x):
        return (self.main_proj(x) * torch.sigmoid(self.gate_proj(x))) + self.extra_bias

class TriDumbSwiGLU(nn.Module):
    def __init__(self, config, is_semantic_layer=False):
        super().__init__()
        mult = 1.5 if is_semantic_layer else 1.0
        hidden_dim = int(config.intermediate_size * mult)
        self.w1 = nn.Linear(config.hidden_size, hidden_dim, bias=False)
        self.w2 = nn.Linear(config.hidden_size, hidden_dim, bias=False)
        self.w3 = nn.Linear(config.hidden_size, hidden_dim, bias=False)
        self.w_out = nn.Linear(hidden_dim, config.hidden_size, bias=False)

    def forward(self, x):
        return self.w_out(F.silu(self.w1(x)) * self.w2(x) * torch.tanh(self.w3(x)))

class DumbGatedAttention(nn.Module):
    def __init__(self, config, layer_idx):
        super().__init__()
        self.num_heads = config.num_attention_heads
        self.head_dim = config.hidden_size // config.num_attention_heads
        self.is_structural = layer_idx < 8
        self.qkv = nn.Linear(config.hidden_size, config.hidden_size * 3, bias=False)
        self.gate = nn.Linear(config.hidden_size, config.hidden_size)
        self.o_proj = nn.Linear(config.hidden_size, config.hidden_size, bias=False)

    def forward(self, x):
        b, s, d = x.shape
        qkv = self.qkv(x).reshape(b, s, 3, self.num_heads, self.head_dim).permute(2, 0, 3, 1, 4)
        q, k, v = qkv[0], qkv[1], qkv[2]
        attn = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(self.head_dim))
        mask = torch.triu(torch.ones(s, s, device=x.device), 1).bool()
        attn = attn.masked_fill(mask, float('-inf'))
        attn = F.softmax(attn, dim=-1)
        out = (attn @ v).transpose(1, 2).reshape(b, s, d)
        g = torch.sigmoid(self.gate(x))
        out = out * g if self.is_structural else out * (0.5 + 0.5 * g)
        return self.o_proj(out)

class DumbMiniLayer(nn.Module):
    def __init__(self, config, layer_idx):
        super().__init__()
        self.ln1 = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
        self.attn = DumbGatedAttention(config, layer_idx)
        self.ln2 = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
        self.ffn = TriDumbSwiGLU(config, is_semantic_layer=(layer_idx >= 8))

    def forward(self, x):
        x = x + self.attn(self.ln1(x))
        x = x + self.ffn(self.ln2(x))
        return x

class DumbMiniModel(PreTrainedModel):
    config_class = DumbMiniConfig

    def __init__(self, config):
        super().__init__(config)
        self.embed_in = nn.Embedding(config.vocab_size, config.emb_split_dim)
        self.up_proj = DumbLinearNext(config.emb_split_dim, config.hidden_size)
        self.pos_embed = nn.Embedding(config.max_position_embeddings, config.hidden_size)
        self.layers = nn.ModuleList([DumbMiniLayer(config, i) for i in range(config.num_hidden_layers)])
        self.ln_f = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
        self.down_proj = nn.Linear(config.hidden_size, config.emb_split_dim, bias=False)
        
        self.post_init()

    def _init_weights(self, module):
        std = self.config.initializer_range
        if isinstance(module, nn.Linear):
            module.weight.data.normal_(mean=0.0, std=std)
            if module.bias is not None:
                module.bias.data.zero_()
        elif isinstance(module, nn.Embedding):
            module.weight.data.normal_(mean=0.0, std=std)
        elif isinstance(module, nn.LayerNorm):
            module.bias.data.zero_()
            module.weight.data.fill_(1.0)

    def forward(self, input_ids=None, labels=None, **kwargs):
        s_len = input_ids.shape[1]
        pos = torch.arange(s_len, device=input_ids.device).unsqueeze(0)
        x = self.up_proj(self.embed_in(input_ids)) + self.pos_embed(pos)
        for layer in self.layers:
            x = layer(x)
        x = self.down_proj(self.ln_f(x))
        logits = F.linear(x, self.embed_in.weight)

        loss = None
        if labels is not None:
            shift_logits = logits[..., :-1, :].contiguous()
            shift_labels = labels[..., 1:].contiguous()
            loss = F.cross_entropy(shift_logits.view(-1, self.config.vocab_size), shift_labels.view(-1))

        return CausalLMOutput(
            loss=loss,
            logits=logits
        )

    def prepare_inputs_for_generation(self, input_ids, **kwargs):
        return {"input_ids": input_ids}