
import torch
import math
import torch.nn as nn
from transformers import PreTrainedModel
from transformers.modeling_outputs import CausalLMOutputWithPast
from .configuration_syko import SykoConfig


# ============================================================
# RMSNorm — FP16 STABİLİTESİ İÇİN: variance hesabı her zaman
# float32'de yapılıyor, sonra geri cast ediliyor. Bu sayede
# fp16'nın küçük sayıları sıfıra yuvarlaması engelleniyor.
# ============================================================
class SykoRMSNorm(nn.Module):
    def __init__(self, hidden_size, eps=1e-5):
        super().__init__()
        self.weight = nn.Parameter(torch.ones(hidden_size))
        self.eps = eps

    def forward(self, x):
        # Kritik: variance float32'de hesaplanıyor
        variance = x.to(torch.float32).pow(2).mean(-1, keepdim=True)
        x = x * torch.rsqrt(variance + self.eps).to(x.dtype)
        return self.weight * x


# ============================================================
# RoPE — Pozisyon encoding'i. inv_freq float32'de tutulur,
# cos/sin her seferinde doğru dtype'a cast edilir.
# ============================================================
class SykoRotaryEmbedding(nn.Module):
    def __init__(self, dim, max_seq_len=1024, theta=10000.0):
        super().__init__()
        inv_freq = 1.0 / (theta ** (torch.arange(0, dim, 2).float() / dim))
        self.register_buffer("inv_freq", inv_freq, persistent=False)

    def forward(self, x, seq_len):
        # float32'de hesapla, sonra x'in dtype'ına cast et — overflow yok
        t = torch.arange(seq_len, device=x.device, dtype=torch.float32)
        freqs = torch.outer(t, self.inv_freq.float())
        emb = torch.cat([freqs, freqs], dim=-1)
        cos = emb.cos()[None, None, :, :].to(x.dtype)
        sin = emb.sin()[None, None, :, :].to(x.dtype)
        return cos, sin


def rotate_half(x):
    x1 = x[..., : x.shape[-1] // 2]
    x2 = x[..., x.shape[-1] // 2 :]
    return torch.cat([-x2, x1], dim=-1)


def apply_rotary(q, k, cos, sin):
    q = (q * cos) + (rotate_half(q) * sin)
    k = (k * cos) + (rotate_half(k) * sin)
    return q, k


# ============================================================
# Attention — GQA destekli (num_key_value_heads < num_heads
# olunca GQA, eşit olunca MHA gibi davranır).
# ============================================================
class SykoAttention(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.num_heads = config.num_attention_heads
        self.num_kv_heads = config.num_key_value_heads
        self.head_dim = config.hidden_size // config.num_attention_heads
        self.num_kv_groups = self.num_heads // self.num_kv_heads

        self.q_proj = nn.Linear(config.hidden_size, self.num_heads * self.head_dim, bias=False)
        self.k_proj = nn.Linear(config.hidden_size, self.num_kv_heads * self.head_dim, bias=False)
        self.v_proj = nn.Linear(config.hidden_size, self.num_kv_heads * self.head_dim, bias=False)
        self.o_proj = nn.Linear(config.hidden_size, config.hidden_size, bias=False)
        self.rotary  = SykoRotaryEmbedding(self.head_dim, config.max_position_embeddings, config.rope_theta)

    def forward(self, hidden_states, attention_mask=None, past_key_value=None, use_cache=False):
        B, T, _ = hidden_states.shape

        q = self.q_proj(hidden_states).view(B, T, self.num_heads,    self.head_dim).transpose(1, 2)
        k = self.k_proj(hidden_states).view(B, T, self.num_kv_heads, self.head_dim).transpose(1, 2)
        v = self.v_proj(hidden_states).view(B, T, self.num_kv_heads, self.head_dim).transpose(1, 2)

        cos, sin = self.rotary(q, seq_len=T)
        q, k = apply_rotary(q, k, cos, sin)

        if past_key_value is not None:
            k = torch.cat([past_key_value[0], k], dim=2)
            v = torch.cat([past_key_value[1], v], dim=2)

        present = (k, v) if use_cache else None

        # GQA: k ve v'yi head sayısına göre tekrarla
        if self.num_kv_groups > 1:
            k = k.repeat_interleave(self.num_kv_groups, dim=1)
            v = v.repeat_interleave(self.num_kv_groups, dim=1)

        # PyTorch'un kendi scaled_dot_product_attention'ı fp16'da güvenli
        attn_out = torch.nn.functional.scaled_dot_product_attention(
            q, k, v,
            attn_mask=None,
            is_causal=(past_key_value is None),
            dropout_p=0.0,
        )

        attn_out = attn_out.transpose(1, 2).contiguous().view(B, T, -1)
        return self.o_proj(attn_out), present


# ============================================================
# SwiGLU MLP — Llama tarzı. gate * up * down.
# ============================================================
class SykoMLP(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.gate_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
        self.up_proj   = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
        self.down_proj = nn.Linear(config.intermediate_size, config.hidden_size, bias=False)

    def forward(self, x):
        return self.down_proj(torch.nn.functional.silu(self.gate_proj(x)) * self.up_proj(x))


# ============================================================
# Decoder Layer — Pre-norm mimarisi (norm önce, residual sonra).
# Bu Llama 3'ün stabil eğitilmesinin temel sebebi.
# ============================================================
class SykoDecoderLayer(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.attn  = SykoAttention(config)
        self.mlp   = SykoMLP(config)
        self.norm1 = SykoRMSNorm(config.hidden_size, config.rms_norm_eps)
        self.norm2 = SykoRMSNorm(config.hidden_size, config.rms_norm_eps)

    def forward(self, hidden_states, attention_mask=None, past_key_value=None, use_cache=False):
        # Attention bloğu
        residual = hidden_states
        hidden_states, present = self.attn(self.norm1(hidden_states), attention_mask, past_key_value, use_cache)
        hidden_states = residual + hidden_states

        # MLP bloğu
        residual = hidden_states
        hidden_states = residual + self.mlp(self.norm2(hidden_states))

        return hidden_states, present


# ============================================================
# Base Model
# ============================================================
class SykoBaseModel(PreTrainedModel):
    config_class = SykoConfig

    def __init__(self, config):
        super().__init__(config)
        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, padding_idx=config.pad_token_id)
        self.layers = nn.ModuleList([SykoDecoderLayer(config) for _ in range(config.num_hidden_layers)])
        self.norm   = SykoRMSNorm(config.hidden_size, config.rms_norm_eps)
        self.post_init()

    def _init_weights(self, module):
        std = 0.02
        if isinstance(module, nn.Linear):
            nn.init.normal_(module.weight, mean=0.0, std=std)
            if module.bias is not None:
                nn.init.zeros_(module.bias)
        elif isinstance(module, nn.Embedding):
            nn.init.normal_(module.weight, mean=0.0, std=std)
            if module.padding_idx is not None:
                module.weight.data[module.padding_idx].zero_()

    def forward(self, input_ids, attention_mask=None, past_key_values=None, use_cache=False, **kwargs):
        # Embedding — scaling YOK. Saf embedding, fp16 taşmaz.
        hidden_states = self.embed_tokens(input_ids)

        presents = []
        for i, layer in enumerate(self.layers):
            pkv = past_key_values[i] if past_key_values is not None else None
            hidden_states, present = layer(hidden_states, attention_mask, pkv, use_cache)
            if use_cache:
                presents.append(present)

        hidden_states = self.norm(hidden_states)
        return hidden_states, (presents if use_cache else None)


# ============================================================
# Causal LM
# ============================================================
class SykoCausalLM(PreTrainedModel):
    config_class = SykoConfig

    def __init__(self, config):
        super().__init__(config)
        self.model   = SykoBaseModel(config)
        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
        self.post_init()

        # Residual projection ağırlıklarını GPT-2 tarzı küçült.
        # post_init sonrası yapılır ki ezilmesin.
        std = 0.02 / math.sqrt(2 * config.num_hidden_layers)
        for name, param in self.named_parameters():
            if name.endswith("o_proj.weight") or name.endswith("down_proj.weight"):
                nn.init.normal_(param, mean=0.0, std=std)

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

    def forward(self, input_ids=None, attention_mask=None, past_key_values=None,
                labels=None, use_cache=None, **kwargs):

        hidden_states, presents = self.model(
            input_ids, attention_mask, past_key_values, use_cache or False
        )

        # KRİTİK: lm_head FLOAT32'de hesaplanıyor.
        # hidden_states fp16 olabilir ama 768->32000 projeksiyonu
        # fp16'da overflow yapar. float32'e çekince sorun çözülür.
        logits = self.lm_head(hidden_states.to(torch.float32))

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

        return CausalLMOutputWithPast(loss=loss, logits=logits, past_key_values=presents)

    def prepare_inputs_for_generation(self, input_ids, past_key_values=None, **kwargs):
        if past_key_values is not None:
            input_ids = input_ids[:, -1:]
        return {"input_ids": input_ids, "past_key_values": past_key_values, "use_cache": True}
