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

from transformers import PreTrainedModel, GenerationMixin
from transformers.modeling_outputs import CausalLMOutputWithPast

from .configuration_veyra import VeyraConfig


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

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


class RotaryEmbedding(nn.Module):
    def __init__(self, dim, max_seq_len, base=10000.0):
        super().__init__()
        self.dim = dim
        self.max_seq_len = max_seq_len
        self.base = base

    def forward(self, x, positions):
        # No buffers. Compute everything fresh to avoid HF remote-code buffer corruption.
        half_dim = self.dim // 2

        inv_freq = 1.0 / (
            self.base ** (
                torch.arange(0, self.dim, 2, device=x.device, dtype=torch.float32) / self.dim
            )
        )

        positions = positions.to(device=x.device, dtype=torch.float32)
        freqs = torch.outer(positions, inv_freq)

        cos = torch.cos(freqs).to(dtype=x.dtype)[None, None, :, :]
        sin = torch.sin(freqs).to(dtype=x.dtype)[None, None, :, :]

        x_even = x[..., 0::2]
        x_odd = x[..., 1::2]

        out_even = x_even * cos - x_odd * sin
        out_odd = x_even * sin + x_odd * cos

        return torch.stack((out_even, out_odd), dim=-1).flatten(-2)


class SwiGLU(nn.Module):
    def __init__(self, d_model, intermediate_size):
        super().__init__()
        self.gate_proj = nn.Linear(d_model, intermediate_size, bias=False)
        self.up_proj = nn.Linear(d_model, intermediate_size, bias=False)
        self.down_proj = nn.Linear(intermediate_size, d_model, bias=False)

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


class GQAAttention(nn.Module):
    def __init__(self, config):
        super().__init__()

        assert config.d_model % config.n_q_heads == 0
        assert config.n_q_heads % config.n_kv_heads == 0

        self.d_model = config.d_model
        self.n_q_heads = config.n_q_heads
        self.n_kv_heads = config.n_kv_heads
        self.head_dim = config.d_model // config.n_q_heads
        self.n_rep = config.n_q_heads // config.n_kv_heads

        self.q_proj = nn.Linear(config.d_model, config.n_q_heads * self.head_dim, bias=False)
        self.k_proj = nn.Linear(config.d_model, config.n_kv_heads * self.head_dim, bias=False)
        self.v_proj = nn.Linear(config.d_model, config.n_kv_heads * self.head_dim, bias=False)
        self.o_proj = nn.Linear(config.n_q_heads * self.head_dim, config.d_model, bias=False)

        self.rope = RotaryEmbedding(self.head_dim, config.max_seq_len, config.rope_theta)

    def _shape_q(self, x):
        b, t, _ = x.shape
        return x.view(b, t, self.n_q_heads, self.head_dim).transpose(1, 2)

    def _shape_kv(self, x):
        b, t, _ = x.shape
        return x.view(b, t, self.n_kv_heads, self.head_dim).transpose(1, 2)

    def _repeat_kv(self, x):
        if self.n_rep == 1:
            return x
        return x.repeat_interleave(self.n_rep, dim=1)

    def forward(self, x):
        b, t, _ = x.shape
        positions = torch.arange(t, device=x.device)

        q = self._shape_q(self.q_proj(x))
        k = self._shape_kv(self.k_proj(x))
        v = self._shape_kv(self.v_proj(x))

        q = self.rope(q, positions)
        k = self.rope(k, positions)

        k = self._repeat_kv(k)
        v = self._repeat_kv(v)

        y = F.scaled_dot_product_attention(
            q,
            k,
            v,
            attn_mask=None,
            dropout_p=0.0,
            is_causal=True,
        )

        y = y.transpose(1, 2).contiguous().view(b, t, self.d_model)

        return self.o_proj(y)


class AttentionBlock(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.input_norm = RMSNorm(config.d_model, config.rms_norm_eps)
        self.attn = GQAAttention(config)
        self.post_attn_norm = RMSNorm(config.d_model, config.rms_norm_eps)
        self.mlp = SwiGLU(config.d_model, config.intermediate_size)

    def forward(self, x):
        x = x + self.attn(self.input_norm(x))
        x = x + self.mlp(self.post_attn_norm(x))
        return x


class MLPOnlyBlock(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.norm = RMSNorm(config.d_model, config.rms_norm_eps)
        self.mlp = SwiGLU(config.d_model, config.intermediate_size)

    def forward(self, x):
        return x + self.mlp(self.norm(x))


class VeyraPreTrainedModel(PreTrainedModel):
    config_class = VeyraConfig
    base_model_prefix = ""
    supports_gradient_checkpointing = False
    _no_split_modules = ["AttentionBlock", "MLPOnlyBlock"]


class VeyraForCausalLM(VeyraPreTrainedModel, GenerationMixin):
    config_class = VeyraConfig
    _tied_weights_keys = {"lm_head.weight": "token_emb.weight"}

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

        self.token_emb = nn.Embedding(config.vocab_size, config.d_model)

        blocks = []
        for kind in config.layer_pattern:
            if kind == "A":
                blocks.append(AttentionBlock(config))
            elif kind == "M":
                blocks.append(MLPOnlyBlock(config))
            else:
                raise ValueError(f"Unknown layer kind: {kind}")

        self.blocks = nn.ModuleList(blocks)
        self.norm = RMSNorm(config.d_model, config.rms_norm_eps)
        self.lm_head = nn.Linear(config.d_model, config.vocab_size, bias=False)

        if config.tie_word_embeddings:
            self.lm_head.weight = self.token_emb.weight

        # Needed by newer Transformers internals, but we intentionally avoid post_init().
        self.all_tied_weights_keys = {"lm_head.weight": "token_emb.weight"}

    def tie_weights(self, *args, **kwargs):
        if self.config.tie_word_embeddings:
            self.lm_head.weight = self.token_emb.weight

    def get_input_embeddings(self):
        return self.token_emb

    def set_input_embeddings(self, value):
        self.token_emb = value
        if self.config.tie_word_embeddings:
            self.lm_head.weight = self.token_emb.weight

    def get_output_embeddings(self):
        return self.lm_head

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

    def forward(
        self,
        input_ids=None,
        attention_mask=None,
        labels=None,
        return_dict=True,
        **kwargs,
    ):
        if input_ids is None:
            raise ValueError("input_ids must be provided")

        x = self.token_emb(input_ids)

        for block in self.blocks:
            x = block(x)

        logits = self.lm_head(self.norm(x))

        loss = None
        if labels is not None:
            loss = F.cross_entropy(
                logits[:, :-1, :].contiguous().view(-1, logits.size(-1)),
                labels[:, 1:].contiguous().view(-1),
                ignore_index=-100,
            )

        if not return_dict:
            output = (logits,)
            return ((loss,) + output) if loss is not None else output

        return CausalLMOutputWithPast(
            loss=loss,
            logits=logits,
            past_key_values=None,
            hidden_states=None,
            attentions=None,
        )

    def prepare_inputs_for_generation(self, input_ids, **kwargs):
        max_len = self.config.max_seq_len
        if input_ids.shape[1] > max_len:
            input_ids = input_ids[:, -max_len:]
        return {"input_ids": input_ids}