import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers import PreTrainedModel
from transformers.cache_utils import Cache, DynamicCache
from transformers.modeling_outputs import CausalLMOutputWithPast

from .configuration_g0nano import G0NanoConfig


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

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


def build_rope_cache(seq_len: int, head_dim: int, theta: float, device, dtype):
    inv_freq = 1.0 / (theta ** (torch.arange(0, head_dim, 2, device=device).float() / head_dim))
    t = torch.arange(seq_len, device=device).float()
    freqs = torch.outer(t, inv_freq)
    return freqs.cos().to(dtype), freqs.sin().to(dtype)


def apply_rope(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
    x1, x2 = x[..., ::2], x[..., 1::2]
    cos = cos[None, None, :x.size(2), :]
    sin = sin[None, None, :x.size(2), :]
    rotated1 = x1 * cos - x2 * sin
    rotated2 = x1 * sin + x2 * cos
    return torch.stack([rotated1, rotated2], dim=-1).flatten(-2).to(x.dtype)


class GQAAttention(nn.Module):
    def __init__(self, config: G0NanoConfig):
        super().__init__()
        self.n_heads = config.num_attention_heads
        self.n_kv_heads = config.num_key_value_heads
        self.head_dim = config.head_dim
        self.n_rep = self.n_heads // self.n_kv_heads

        self.wq = nn.Linear(config.hidden_size, self.n_heads * self.head_dim, bias=False)
        self.wk = nn.Linear(config.hidden_size, self.n_kv_heads * self.head_dim, bias=False)
        self.wv = nn.Linear(config.hidden_size, self.n_kv_heads * self.head_dim, bias=False)
        self.wo = nn.Linear(self.n_heads * self.head_dim, config.hidden_size, bias=False)

    def forward(self, x, cos, sin, past_key_values=None, layer_idx=None, use_cache=False):
        B, T, _ = x.shape
        q = self.wq(x).view(B, T, self.n_heads, self.head_dim).transpose(1, 2)
        k = self.wk(x).view(B, T, self.n_kv_heads, self.head_dim).transpose(1, 2)
        v = self.wv(x).view(B, T, self.n_kv_heads, self.head_dim).transpose(1, 2)

        past_len = 0 if past_key_values is None else past_key_values.get_seq_length(layer_idx)
        q = apply_rope(q, cos[past_len:past_len + T], sin[past_len:past_len + T])
        k = apply_rope(k, cos[past_len:past_len + T], sin[past_len:past_len + T])

        if past_key_values is not None:
            k, v = past_key_values.update(k, v, layer_idx)

        k = k.repeat_interleave(self.n_rep, dim=1)
        v = v.repeat_interleave(self.n_rep, dim=1)

        if past_len == 0 and T > 1:
            out = F.scaled_dot_product_attention(q, k, v, is_causal=True)
        else:
            key_len = k.size(2)
            query_positions = torch.arange(T, device=x.device) + past_len
            key_positions = torch.arange(key_len, device=x.device)
            attn_mask = key_positions.unsqueeze(0) <= query_positions.unsqueeze(1)
            out = F.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask)

        out = out.transpose(1, 2).contiguous().view(B, T, self.n_heads * self.head_dim)
        return self.wo(out)


class SwiGLU(nn.Module):
    def __init__(self, config: G0NanoConfig):
        super().__init__()
        self.w_gate = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
        self.w_up = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
        self.w_down = nn.Linear(config.intermediate_size, config.hidden_size, bias=False)

    def forward(self, x):
        return self.w_down(F.silu(self.w_gate(x)) * self.w_up(x))


class Block(nn.Module):
    def __init__(self, config: G0NanoConfig):
        super().__init__()
        self.attn_norm = RMSNorm(config.hidden_size, config.rms_norm_eps)
        self.attn = GQAAttention(config)
        self.mlp_norm = RMSNorm(config.hidden_size, config.rms_norm_eps)
        self.mlp = SwiGLU(config)

    def forward(self, x, cos, sin, past_key_values=None, layer_idx=None, use_cache=False):
        x = x + self.attn(self.attn_norm(x), cos, sin, past_key_values, layer_idx, use_cache)
        x = x + self.mlp(self.mlp_norm(x))
        return x


class G0NanoPreTrainedModel(PreTrainedModel):
    config_class = G0NanoConfig
    base_model_prefix = "g0nano"
    supports_gradient_checkpointing = False

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


class G0NanoForCausalLM(G0NanoPreTrainedModel):
    _tied_weights_keys = {"lm_head.weight": "tok_emb.weight"}

    def __init__(self, config: G0NanoConfig):
        super().__init__(config)
        self.tok_emb = nn.Embedding(config.vocab_size, config.hidden_size)
        self.blocks = nn.ModuleList([Block(config) for _ in range(config.num_hidden_layers)])
        self.final_norm = RMSNorm(config.hidden_size, config.rms_norm_eps)
        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)

        self._rope_cache = {}
        self.post_init()

    def get_input_embeddings(self):
        return self.tok_emb

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

    def get_output_embeddings(self):
        return self.lm_head

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

    def _get_rope(self, seq_len, device, dtype):
        key = (seq_len, device, dtype)
        if key not in self._rope_cache:
            self._rope_cache[key] = build_rope_cache(
                seq_len, self.config.head_dim, self.config.rope_theta, device, dtype
            )
        return self._rope_cache[key]

    def forward(
        self,
        input_ids=None,
        attention_mask=None,
        past_key_values=None,
        labels=None,
        use_cache=None,
        output_attentions=None,
        output_hidden_states=None,
        return_dict=None,
        **kwargs,
    ):
        use_cache = bool(use_cache) if use_cache is not None else False
        B, T = input_ids.shape

        if use_cache and past_key_values is None:
            past_key_values = DynamicCache(config=self.config)
        past_len = past_key_values.get_seq_length() if past_key_values is not None else 0
        assert past_len + T <= self.config.max_position_embeddings, (
            f"sequence length {past_len + T} > max_position_embeddings "
            f"{self.config.max_position_embeddings}"
        )

        x = self.tok_emb(input_ids)
        cos, sin = self._get_rope(past_len + T, input_ids.device, x.dtype)

        for i, block in enumerate(self.blocks):
            x = block(x, cos, sin, past_key_values, i, use_cache)
        x = self.final_norm(x)
        logits = self.lm_head(x)

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

        return CausalLMOutputWithPast(
            loss=loss,
            logits=logits,
            past_key_values=past_key_values if use_cache else None,
        )

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