from __future__ import annotations

from typing import Optional

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch import Tensor

from config import ModelConfig


class RMSNorm(nn.Module):
    def __init__(self, dim: int, eps: float = 1.0e-5):
        super().__init__()
        self.weight = nn.Parameter(torch.ones(dim))
        self.eps = eps

    def forward(self, x: Tensor) -> Tensor:
        dtype = x.dtype
        x32 = x.float()
        rms = x32.pow(2).mean(dim=-1, keepdim=True).add_(self.eps).rsqrt_()
        out = (x32 * rms).to(dtype)
        return out * self.weight.to(dtype)


class RotaryEmbedding(nn.Module):
    def __init__(self, head_dim: int, max_seq_len: int, theta: float = 10_000.0):
        super().__init__()
        self.head_dim = head_dim
        self.max_seq_len = max_seq_len
        self.theta = theta
        self._cached_len: int = 0
        self._cos_cache: Optional[Tensor] = None
        self._sin_cache: Optional[Tensor] = None

    def _build_cache(self, seq_len: int, device, dtype):
        inv_freq = 1.0 / (
            self.theta ** (torch.arange(0, self.head_dim, 2, dtype=torch.float32, device=device) / self.head_dim)
        )
        t = torch.arange(seq_len, dtype=torch.float32, device=device)
        freqs = torch.outer(t, inv_freq)
        emb = torch.cat([freqs, freqs], dim=-1)
        self._cos_cache = emb.cos().to(dtype)
        self._sin_cache = emb.sin().to(dtype)
        self._cached_len = seq_len

    def forward(self, seq_len: int, device, dtype) -> tuple[Tensor, Tensor]:
        if (
            self._cos_cache is None
            or seq_len > self._cached_len
            or self._cos_cache.device != device
            or self._cos_cache.dtype != dtype
        ):
            self._build_cache(max(seq_len, self.max_seq_len), device, dtype)
        return self._cos_cache[:seq_len], self._sin_cache[:seq_len]


def _rotate_half(x: Tensor) -> Tensor:
    x1, x2 = x.chunk(2, dim=-1)
    return torch.cat([-x2, x1], dim=-1)


def apply_rope(q: Tensor, k: Tensor, cos: Tensor, sin: Tensor) -> tuple[Tensor, Tensor]:
    cos = cos.unsqueeze(0).unsqueeze(0)
    sin = sin.unsqueeze(0).unsqueeze(0)
    q_rot = (q * cos) + (_rotate_half(q) * sin)
    k_rot = (k * cos) + (_rotate_half(k) * sin)
    return q_rot, k_rot


class Attention(nn.Module):
    def __init__(self, cfg: ModelConfig, layer_idx: int):
        super().__init__()
        self.cfg = cfg
        self.layer_idx = layer_idx
        self.num_heads = cfg.num_heads
        self.num_kv_heads = cfg.num_kv_heads
        self.head_dim = cfg.head_dim
        self.kv_groups = cfg.kv_groups
        self.scale = self.head_dim ** -0.5

        h, hd = cfg.hidden_size, self.head_dim
        self.q_proj = nn.Linear(h, self.num_heads * hd, bias=False)
        self.k_proj = nn.Linear(h, self.num_kv_heads * hd, bias=False)
        self.v_proj = nn.Linear(h, self.num_kv_heads * hd, bias=False)
        self.o_proj = nn.Linear(self.num_heads * hd, h, bias=False)

        if cfg.qk_norm:
            self.q_norm = RMSNorm(hd, eps=cfg.rms_norm_eps)
            self.k_norm = RMSNorm(hd, eps=cfg.rms_norm_eps)
        else:
            self.q_norm = nn.Identity()
            self.k_norm = nn.Identity()

        self.attn_dropout = cfg.attn_dropout

    def forward(self, x: Tensor, cos: Tensor, sin: Tensor) -> Tensor:
        B, S, _ = x.shape
        q = self.q_proj(x).view(B, S, self.num_heads, self.head_dim)
        k = self.k_proj(x).view(B, S, self.num_kv_heads, self.head_dim)
        v = self.v_proj(x).view(B, S, self.num_kv_heads, self.head_dim)

        q = self.q_norm(q)
        k = self.k_norm(k)

        q = q.transpose(1, 2)
        k = k.transpose(1, 2)
        v = v.transpose(1, 2)

        q, k = apply_rope(q, k, cos, sin)

        if self.kv_groups > 1:
            k = k.repeat_interleave(self.kv_groups, dim=1)
            v = v.repeat_interleave(self.kv_groups, dim=1)

        out = F.scaled_dot_product_attention(
            q, k, v,
            attn_mask=None,
            dropout_p=self.attn_dropout if self.training else 0.0,
            is_causal=True,
        )
        out = out.transpose(1, 2).contiguous().view(B, S, self.num_heads * self.head_dim)
        return self.o_proj(out)


class SwiGLU(nn.Module):
    def __init__(self, cfg: ModelConfig):
        super().__init__()
        h, i = cfg.hidden_size, cfg.intermediate_size
        self.gate_proj = nn.Linear(h, i, bias=False)
        self.up_proj = nn.Linear(h, i, bias=False)
        self.down_proj = nn.Linear(i, h, bias=False)

    def forward(self, x: Tensor) -> Tensor:
        return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x))


class Block(nn.Module):
    def __init__(self, cfg: ModelConfig, layer_idx: int):
        super().__init__()
        self.input_norm = RMSNorm(cfg.hidden_size, cfg.rms_norm_eps)
        self.attn = Attention(cfg, layer_idx)
        self.post_attn_norm = RMSNorm(cfg.hidden_size, cfg.rms_norm_eps)
        self.mlp = SwiGLU(cfg)
        self.resid_drop = nn.Dropout(cfg.resid_dropout) if cfg.resid_dropout > 0 else nn.Identity()

    def forward(self, x: Tensor, cos: Tensor, sin: Tensor) -> Tensor:
        x = x + self.resid_drop(self.attn(self.input_norm(x), cos, sin))
        x = x + self.resid_drop(self.mlp(self.post_attn_norm(x)))
        return x


class MarulLLM(nn.Module):
    def __init__(self, cfg: ModelConfig):
        super().__init__()
        self.cfg = cfg

        self.embed_tokens = nn.Embedding(cfg.vocab_size, cfg.hidden_size)
        self.rotary = RotaryEmbedding(cfg.head_dim, cfg.max_seq_len, cfg.rope_theta)
        self.layers = nn.ModuleList(Block(cfg, i) for i in range(cfg.num_layers))
        self.final_norm = RMSNorm(cfg.hidden_size, cfg.rms_norm_eps)

        if cfg.tie_word_embeddings:
            self.lm_head = None
        else:
            self.lm_head = nn.Linear(cfg.hidden_size, cfg.vocab_size, bias=False)

        self.apply(self._init_weights)
        self._scale_residual_inits()

        self.num_params = sum(p.numel() for p in self.parameters())
        self.num_params_trainable = sum(p.numel() for p in self.parameters() if p.requires_grad)
        embed_params = cfg.vocab_size * cfg.hidden_size
        self.num_params_non_embed = self.num_params - embed_params

    def _init_weights(self, module: nn.Module) -> None:
        std = self.cfg.initializer_range
        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 _scale_residual_inits(self) -> None:
        scale = (2 * self.cfg.num_layers) ** -0.5
        with torch.no_grad():
            for block in self.layers:
                block.attn.o_proj.weight.mul_(scale)
                block.mlp.down_proj.weight.mul_(scale)

    def forward(
        self,
        input_ids: Tensor,
        targets: Optional[Tensor] = None,
        return_logits: bool = True,
    ) -> tuple[Optional[Tensor], Optional[Tensor]]:
        B, S = input_ids.shape
        assert S <= self.cfg.max_seq_len, (
            f"dizi uzunluğu {S}, modelin bağlam sınırı {self.cfg.max_seq_len}")

        x = self.embed_tokens(input_ids)
        cos, sin = self.rotary(S, x.device, x.dtype)

        for block in self.layers:
            x = block(x, cos, sin)
        x = self.final_norm(x)

        if self.cfg.tie_word_embeddings:
            logits = F.linear(x, self.embed_tokens.weight)
        else:
            logits = self.lm_head(x)

        loss: Optional[Tensor] = None
        if targets is not None:
            flat_logits = logits.view(-1, logits.size(-1))
            flat_targets = targets.view(-1)
            ce = F.cross_entropy(
                flat_logits, flat_targets, ignore_index=-100, reduction="mean"
            )
            loss = ce
            if self.cfg.z_loss_coef > 0:
                mask = flat_targets != -100
                lse = torch.logsumexp(flat_logits, dim=-1)
                if mask.any():
                    z = (lse[mask].float().pow(2)).mean()
                    loss = loss + self.cfg.z_loss_coef * z

        if not return_logits and targets is not None:
            logits = None
        return logits, loss

    @torch.no_grad()
    def generate(
        self,
        input_ids: Tensor,
        max_new_tokens: int = 128,
        temperature: float = 0.6,
        top_k: int = 40,
        top_p: float = 0.88,
        repetition_penalty: float = 1.20,
        no_repeat_ngram_size: int = 4,
        min_p: float = 0.05,
        eos_token_id: Optional[int] = None,
    ) -> Tensor:
        self.eval()
        eos = eos_token_id if eos_token_id is not None else self.cfg.eos_token_id
        out = input_ids.clone()
        device = out.device

        for _ in range(max_new_tokens):
            ctx = out[:, -self.cfg.max_seq_len:]
            logits, _ = self.forward(ctx)
            logits = logits[:, -1, :].float()

            if repetition_penalty is not None and repetition_penalty != 1.0:
                for b in range(out.size(0)):
                    seen = torch.unique(out[b])
                    vals = logits[b, seen]
                    logits[b, seen] = torch.where(
                        vals > 0, vals / repetition_penalty, vals * repetition_penalty
                    )

            if no_repeat_ngram_size and no_repeat_ngram_size > 0:
                n = no_repeat_ngram_size
                if out.size(1) >= n - 1:
                    for b in range(out.size(0)):
                        seq = out[b].tolist()
                        ngrams: dict = {}
                        for i in range(len(seq) - n + 1):
                            prefix = tuple(seq[i : i + n - 1])
                            ngrams.setdefault(prefix, set()).add(seq[i + n - 1])
                        curr = tuple(seq[-(n - 1):])
                        if curr in ngrams:
                            banned = torch.tensor(list(ngrams[curr]), device=device, dtype=torch.long)
                            logits[b, banned] = float("-inf")

            if temperature is not None and temperature != 1.0:
                logits = logits / max(temperature, 1.0e-6)

            if top_k is not None and top_k > 0:
                v, _ = torch.topk(logits, k=min(top_k, logits.size(-1)))
                logits[logits < v[:, -1:]] = float("-inf")

            if min_p is not None and min_p > 0.0:
                probs_tmp = F.softmax(logits, dim=-1)
                max_probs, _ = probs_tmp.max(dim=-1, keepdim=True)
                logits = logits.masked_fill(probs_tmp < (max_probs * min_p), float("-inf"))

            if top_p is not None and 0.0 < top_p < 1.0:
                sorted_logits, sorted_idx = torch.sort(logits, descending=True, dim=-1)
                probs = F.softmax(sorted_logits, dim=-1)
                cumprobs = probs.cumsum(dim=-1)
                mask = cumprobs > top_p
                mask[..., 1:] = mask[..., :-1].clone()
                mask[..., 0] = False
                sorted_logits = sorted_logits.masked_fill(mask, float("-inf"))
                logits = torch.full_like(logits, float("-inf")).scatter(-1, sorted_idx, sorted_logits)

            probs = F.softmax(logits, dim=-1)
            next_token = torch.multinomial(probs, num_samples=1)
            out = torch.cat([out, next_token], dim=1)

            if eos is not None and (next_token == eos).all():
                break
        return out