import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers import PreTrainedModel, GenerationMixin
from transformers.modeling_outputs import CausalLMOutput
from .configuration_sovythos import SovythosConfig

class RMSNorm(nn.Module):
    def __init__(self, dim, eps=1e-6):
        super().__init__()
        self.w = nn.Parameter(torch.ones(dim))
        self.eps = eps

    def forward(self, x):
        rms = x.pow(2).mean(-1, keepdim=True)
        return self.w * x * torch.rsqrt(rms + self.eps)

class RoPE(nn.Module):
    def __init__(self, head_dim):
        super().__init__()
        inv_freq = 1.0 / (10000 ** (torch.arange(0, head_dim, 2).float() / head_dim))
        self.register_buffer("inv_freq", inv_freq, persistent=False)
        self._cos_cache = None
        self._sin_cache = None

    def _build_cache(self, seq_len, device):
        if self._cos_cache is not None and self._cos_cache.shape[0] >= seq_len:
            return
        t = torch.arange(seq_len, device=device).type_as(self.inv_freq)
        freqs = torch.einsum("i,j->ij", t, self.inv_freq)
        emb = torch.cat((freqs, freqs), dim=-1)
        self._cos_cache = emb.cos()
        self._sin_cache = emb.sin()

    def forward(self, x, seq_len):
        self._build_cache(seq_len, x.device)
        cos = self._cos_cache[:seq_len][None, None, :, :]
        sin = self._sin_cache[:seq_len][None, None, :, :]
        x1, x2 = x[..., :x.shape[-1]//2], x[..., x.shape[-1]//2:]
        return (x * cos) + (torch.cat((-x2, x1), dim=-1) * sin)

class TitanAttention(nn.Module):
    def __init__(self, dim, heads):
        super().__init__()
        self.heads = heads
        self.head_dim = dim // heads
        self.q_proj = nn.Linear(dim, dim, bias=False)
        self.k_proj = nn.Linear(dim, dim, bias=False)
        self.v_proj = nn.Linear(dim, dim, bias=False)
        self.o_proj = nn.Linear(dim, dim, bias=False)
        self.q_norm = RMSNorm(self.head_dim)
        self.k_norm = RMSNorm(self.head_dim)
        self.rope = RoPE(self.head_dim)

    def forward(self, x, is_causal=True):
        B, T, C = x.shape
        q = self.q_proj(x).view(B, T, self.heads, self.head_dim).transpose(1, 2)
        k = self.k_proj(x).view(B, T, self.heads, self.head_dim).transpose(1, 2)
        v = self.v_proj(x).view(B, T, self.heads, self.head_dim).transpose(1, 2)
        
        q = self.rope(self.q_norm(q), T)
        k = self.rope(self.k_norm(k), T)
        
        out = F.scaled_dot_product_attention(q, k, v, attn_mask=None, is_causal=is_causal)
        out = out.transpose(1, 2).contiguous().view(B, T, C)
        return self.o_proj(out)

class Block(nn.Module):
    def __init__(self, dim, heads):
        super().__init__()
        self.n1 = RMSNorm(dim)
        self.attn = TitanAttention(dim, heads)
        self.n2 = RMSNorm(dim)
        self.w1 = nn.Linear(dim, 4 * dim, bias=False)
        self.w2 = nn.Linear(dim, 4 * dim, bias=False)
        self.w3 = nn.Linear(4 * dim, dim, bias=False)

    def forward(self, x):
        x = x + self.attn(self.n1(x))
        h = self.n2(x)
        x = x + self.w3(F.silu(self.w1(h)) * self.w2(h))
        return x

class SovythosModel(PreTrainedModel, GenerationMixin):
    config_class = SovythosConfig

    def __init__(self, config):
        # حل مشكلة num_hidden_layers بربطها ديناميكياً بالمتغير الحالي
        if not hasattr(config, "num_hidden_layers"):
            config.num_hidden_layers = getattr(config, "layers", 12)
            
        super().__init__(config)
        self.emb = nn.Embedding(config.vocab_size, config.dim)
        self.blocks = nn.ModuleList([Block(config.dim, config.heads) for _ in range(config.layers)])
        self.norm = RMSNorm(config.dim)
        self.fc = nn.Linear(config.dim, config.vocab_size, bias=False)
        
        self.fc.weight = self.emb.weight
        self.post_init()

    def get_input_embeddings(self):
        return self.emb

    def set_input_embeddings(self, value):
        self.emb = value

    def get_output_embeddings(self):
        return self.fc

    def set_output_embeddings(self, new_embeddings):
        self.fc = new_embeddings

    def forward(self, input_ids, attention_mask=None, labels=None, **kwargs):
        x = self.emb(input_ids)
        for blk in self.blocks:
            x = blk(x)
        x = self.norm(x)
        logits = self.fc(x)

        loss = None
        if labels is not None:
            loss_fn = nn.CrossEntropyLoss()
            loss = loss_fn(logits.view(-1, self.config.vocab_size), labels.view(-1))

        return CausalLMOutput(loss=loss, logits=logits)

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