import importlib
import math

import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers import PreTrainedModel
from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast

from .configuration_tiny_gpt import TinyGPTConfig


_FLASH2_KERNEL = None
_FLASH3_KERNEL = None


def _get_flash2_kernel():
    global _FLASH2_KERNEL

    if _FLASH2_KERNEL is None:
        kernels = importlib.import_module("kernels")
        _FLASH2_KERNEL = kernels.get_kernel("kernels-community/flash-attn2", version=1)

    return _FLASH2_KERNEL


def _get_flash3_kernel():
    global _FLASH3_KERNEL

    if _FLASH3_KERNEL is None:
        kernels = importlib.import_module("kernels")
        _FLASH3_KERNEL = kernels.get_kernel("kernels-community/flash-attn3", version=1)

    return _FLASH3_KERNEL


def _get_sageattn():
    module = importlib.import_module("sageattention")
    return module.sageattn


def rotate_half(x):
    x_even = x[..., ::2]
    x_odd = x[..., 1::2]
    x_rot = torch.stack((-x_odd, x_even), dim=-1)
    return x_rot.flatten(start_dim=-2)


def apply_rope(x, cos, sin):
    return (x * cos) + (rotate_half(x) * sin)


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

        if dim % 2 != 0:
            raise ValueError(f"RoPE dim must be even, got {dim}")

        self.dim = int(dim)
        self.max_position_embeddings = int(max_position_embeddings)
        self.base = float(base)

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

        self.register_buffer(
            "inv_freq",
            inv_freq,
            persistent=False,
        )

        self._cos_cached = None
        self._sin_cached = None
        self._seq_len_cached = 0
        self._device_cached = None
        self._dtype_cached = None

    def _build_cache(self, seq_len, device, dtype):
        t = torch.arange(
            seq_len,
            device=device,
            dtype=torch.float32,
        )

        freqs = torch.einsum(
            "i,j->ij",
            t,
            self.inv_freq.to(device=device, dtype=torch.float32),
        )

        emb = torch.repeat_interleave(freqs, repeats=2, dim=-1)

        cos = emb.cos().to(dtype=dtype).view(1, 1, seq_len, self.dim)
        sin = emb.sin().to(dtype=dtype).view(1, 1, seq_len, self.dim)

        self._cos_cached = cos
        self._sin_cached = sin
        self._seq_len_cached = int(seq_len)
        self._device_cached = device
        self._dtype_cached = dtype

    def forward(self, seq_len, device, dtype):
        if (
            self._cos_cached is None
            or self._sin_cached is None
            or self._seq_len_cached < seq_len
            or self._device_cached != device
            or self._dtype_cached != dtype
        ):
            self._build_cache(
                seq_len=seq_len,
                device=device,
                dtype=dtype,
            )

        return (
            self._cos_cached[:, :, :seq_len, :],
            self._sin_cached[:, :, :seq_len, :],
        )


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

        if config.n_embd % config.n_head != 0:
            raise ValueError("n_embd must be divisible by n_head")

        self.n_head = int(config.n_head)
        self.head_dim = int(config.n_embd // config.n_head)
        self.attention_backend = str(getattr(config, "attention_backend", "torch"))
        self.torch_fallback = bool(getattr(config, "torch_fallback", False))
        self.dropout_p = float(config.dropout)

        if self.head_dim % 2 != 0:
            raise ValueError(f"RoPE requires even head_dim, got {self.head_dim}")

        if self.attention_backend not in ("sage", "torch", "flash2", "flash3"):
            raise ValueError("attention_backend must be sage, torch, flash2 or flash3")

        if self.attention_backend == "sage" and self.head_dim not in (64, 96, 128):
            raise ValueError(f"SageAttention requires head_dim in [64, 96, 128], got {self.head_dim}")

        if self.attention_backend == "sage" and self.dropout_p != 0.0:
            raise ValueError("SageAttention requires dropout=0.0")

        if self.attention_backend == "flash3" and self.dropout_p != 0.0:
            raise ValueError("FlashAttention3 requires dropout=0.0")

        if self.attention_backend in ("flash2", "flash3") and self.head_dim % 8 != 0:
            raise ValueError(f"FlashAttention requires head_dim multiple of 8, got {self.head_dim}")

        self.qkv = nn.Linear(config.n_embd, 3 * config.n_embd, bias=False)
        self.proj = nn.Linear(config.n_embd, config.n_embd, bias=False)
        self.dropout = nn.Dropout(config.dropout)

        self.rope = RotaryEmbedding(
            dim=self.head_dim,
            max_position_embeddings=config.ctx_len,
            base=float(getattr(config, "rope_base", 10000.0)),
        )

        mask = torch.tril(torch.ones(config.ctx_len, config.ctx_len, dtype=torch.bool))
        self.register_buffer(
            "mask",
            mask.view(1, 1, config.ctx_len, config.ctx_len),
            persistent=False,
        )

        self.sageattn = None
        self.flash_kernel = None

        if self.attention_backend == "sage":
            try:
                self.sageattn = _get_sageattn()
            except Exception:
                if self.torch_fallback:
                    self.attention_backend = "torch"
                else:
                    raise

        if self.attention_backend == "flash2":
            try:
                self.flash_kernel = _get_flash2_kernel()
            except Exception:
                if self.torch_fallback:
                    self.attention_backend = "torch"
                else:
                    raise

        if self.attention_backend == "flash3":
            try:
                self.flash_kernel = _get_flash3_kernel()
            except Exception:
                if self.torch_fallback:
                    self.attention_backend = "torch"
                else:
                    raise

    def _torch_attention(self, q, k, v, t):
        scores = (q @ k.transpose(-2, -1)) / math.sqrt(self.head_dim)
        scores = scores.masked_fill(self.mask[:, :, :t, :t] == 0, float("-inf"))
        att = F.softmax(scores.float(), dim=-1).to(q.dtype)
        att = self.dropout(att)
        return att @ v

    def _sage_attention(self, q, k, v):
        if self.sageattn is None or not q.is_cuda:
            if self.torch_fallback:
                return None
            raise RuntimeError("SageAttention requires CUDA + sageattention")

        return self.sageattn(
            q.contiguous(),
            k.contiguous(),
            v.contiguous(),
            tensor_layout="HND",
            is_causal=True,
        )

    def _flash2_attention(self, q, k, v):
        if self.flash_kernel is None or not q.is_cuda:
            if self.torch_fallback:
                return None
            raise RuntimeError("FlashAttention2 requires CUDA + kernels")

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

        dropout_p = self.dropout_p if self.training else 0.0

        y = self.flash_kernel.flash_attn_func(
            q,
            k,
            v,
            dropout_p=dropout_p,
            causal=True,
        )

        return y.transpose(1, 2).contiguous()

    def _flash3_attention(self, q, k, v):
        if self.flash_kernel is None or not q.is_cuda:
            if self.torch_fallback:
                return None
            raise RuntimeError("FlashAttention3 requires CUDA + kernels")

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

        y = self.flash_kernel.flash_attn_func(
            q,
            k,
            v,
            causal=True,
        )

        return y.transpose(1, 2).contiguous()

    def forward(self, x):
        b, t, c = x.shape

        qkv = self.qkv(x)
        q, k, v = qkv.chunk(3, dim=-1)

        q = q.view(b, t, self.n_head, self.head_dim).transpose(1, 2).contiguous()
        k = k.view(b, t, self.n_head, self.head_dim).transpose(1, 2).contiguous()
        v = v.view(b, t, self.n_head, self.head_dim).transpose(1, 2).contiguous()

        cos, sin = self.rope(
            seq_len=t,
            device=q.device,
            dtype=q.dtype,
        )

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

        if self.attention_backend == "sage":
            y = self._sage_attention(q, k, v)
            if y is None:
                y = self._torch_attention(q, k, v, t)
        elif self.attention_backend == "flash2":
            y = self._flash2_attention(q, k, v)
            if y is None:
                y = self._torch_attention(q, k, v, t)
        elif self.attention_backend == "flash3":
            y = self._flash3_attention(q, k, v)
            if y is None:
                y = self._torch_attention(q, k, v, t)
        else:
            y = self._torch_attention(q, k, v, t)

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

        return self.proj(y)


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

        intermediate_size = int(getattr(config, "intermediate_size", 4 * config.n_embd))

        self.fc = nn.Linear(config.n_embd, intermediate_size, bias=False)
        self.proj = nn.Linear(intermediate_size, config.n_embd, bias=False)
        self.dropout = nn.Dropout(config.dropout)

    def forward(self, x):
        x = self.fc(x)
        x = F.gelu(x)
        x = self.proj(x)
        x = self.dropout(x)

        return x


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

        self.ln1 = nn.LayerNorm(config.n_embd)
        self.attn = CausalSelfAttention(config)
        self.ln2 = nn.LayerNorm(config.n_embd)
        self.mlp = MLP(config)

    def forward(self, x):
        x = x + self.attn(self.ln1(x))
        x = x + self.mlp(self.ln2(x))

        return x


class TinyGPTPreTrainedModel(PreTrainedModel):
    config_class = TinyGPTConfig
    base_model_prefix = "tiny_gpt"
    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)
            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=0.02)


class TinyGPTModel(TinyGPTPreTrainedModel):
    _tied_weights_keys = ["head.weight"]

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

        self.tok_emb = nn.Embedding(config.vocab_size, config.n_embd)
        self.drop = nn.Dropout(config.dropout)

        self.blocks = nn.ModuleList(
            [Block(config) for _ in range(config.n_layer)]
        )

        self.ln_f = nn.LayerNorm(config.n_embd)
        self.head = nn.Linear(config.n_embd, config.vocab_size, bias=False)

        self.post_init()
        self.tie_weights()

    def get_input_embeddings(self):
        return self.tok_emb

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

    def get_output_embeddings(self):
        return self.head

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

    def tie_weights(self):
        self._tie_or_clone_weights(self.head, self.tok_emb)

    def forward(
        self,
        input_ids,
        attention_mask=None,
        return_dict=True,
        return_logits=False,
        **kwargs,
    ):
        b, t = input_ids.shape

        if t > self.config.ctx_len:
            raise ValueError(
                f"Input length {t} > ctx_len {self.config.ctx_len}. "
                "Truncate before calling the model."
            )

        x = self.tok_emb(input_ids)
        x = self.drop(x)

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

        hidden = self.ln_f(x)
        logits = self.head(hidden) if return_logits else None

        if not return_dict:
            return (hidden, logits) if return_logits else (hidden,)

        if return_logits:
            return hidden, logits

        return BaseModelOutputWithPast(
            last_hidden_state=hidden,
            past_key_values=None,
            hidden_states=None,
            attentions=None,
        )


class TinyGPTForCausalLM(TinyGPTPreTrainedModel):
    _tied_weights_keys = ["tiny_gpt.head.weight"]

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

        self.tiny_gpt = TinyGPTModel(config)

        self.post_init()
        self.tie_weights()

    def get_input_embeddings(self):
        return self.tiny_gpt.tok_emb

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

    def get_output_embeddings(self):
        return self.tiny_gpt.head

    def set_output_embeddings(self, new_embeddings):
        self.tiny_gpt.head = new_embeddings

    def tie_weights(self):
        self._tie_or_clone_weights(
            self.tiny_gpt.head,
            self.tiny_gpt.tok_emb,
        )

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

    def forward(
        self,
        input_ids,
        attention_mask=None,
        labels=None,
        return_dict=True,
        **kwargs,
    ):
        hidden, logits = self.tiny_gpt(
            input_ids=input_ids,
            attention_mask=attention_mask,
            return_dict=True,
            return_logits=True,
        )

        loss = None

        if labels is not None:
            shift_logits = logits[:, :-1, :].contiguous()
            shift_labels = labels[:, 1:].contiguous()

            loss = F.cross_entropy(
                shift_logits.view(-1, shift_logits.size(-1)).float(),
                shift_labels.view(-1),
            )

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

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