from __future__ import annotations

import math

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

from tiny_transformer.config import ModelConfig


class CausalSelfAttention(nn.Module):
    def __init__(self, config: ModelConfig) -> None:
        super().__init__()
        self.n_head = config.n_head
        self.head_dim = config.n_embd // config.n_head
        self.qkv = nn.Linear(config.n_embd, 3 * config.n_embd)
        self.proj = nn.Linear(config.n_embd, config.n_embd)
        self.attn_dropout = nn.Dropout(config.dropout)
        self.resid_dropout = nn.Dropout(config.dropout)
        mask = torch.tril(torch.ones(config.block_size, config.block_size))
        self.register_buffer("causal_mask", mask.view(1, 1, config.block_size, config.block_size))

    def forward(
        self, x: torch.Tensor, return_attention: bool = False
    ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
        batch, seq_len, channels = x.shape
        qkv = self.qkv(x)
        query, key, value = qkv.split(channels, dim=2)

        query = query.view(batch, seq_len, self.n_head, self.head_dim).transpose(1, 2)
        key = key.view(batch, seq_len, self.n_head, self.head_dim).transpose(1, 2)
        value = value.view(batch, seq_len, self.n_head, self.head_dim).transpose(1, 2)

        scores = query @ key.transpose(-2, -1) / math.sqrt(self.head_dim)
        scores = scores.masked_fill(self.causal_mask[:, :, :seq_len, :seq_len] == 0, float("-inf"))
        weights = F.softmax(scores, dim=-1)
        weights = self.attn_dropout(weights)
        out = weights @ value
        out = out.transpose(1, 2).contiguous().view(batch, seq_len, channels)
        out = self.resid_dropout(self.proj(out))
        if return_attention:
            return out, weights
        return out

    def forward_cached(
        self,
        x: torch.Tensor,
        past_key_value: tuple[torch.Tensor, torch.Tensor] | None = None,
    ) -> tuple[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]:
        batch, seq_len, channels = x.shape
        query, key, value = self.qkv(x).split(channels, dim=2)
        query = query.view(batch, seq_len, self.n_head, self.head_dim).transpose(1, 2)
        key = key.view(batch, seq_len, self.n_head, self.head_dim).transpose(1, 2)
        value = value.view(batch, seq_len, self.n_head, self.head_dim).transpose(1, 2)

        past_len = 0
        if past_key_value is not None:
            past_key, past_value = past_key_value
            past_len = past_key.size(2)
            key = torch.cat((past_key, key), dim=2)
            value = torch.cat((past_value, value), dim=2)

        total_len = key.size(2)
        scores = query @ key.transpose(-2, -1) / math.sqrt(self.head_dim)
        mask = self.causal_mask[:, :, past_len : past_len + seq_len, :total_len]
        scores = scores.masked_fill(mask == 0, float("-inf"))
        weights = self.attn_dropout(F.softmax(scores, dim=-1))
        out = weights @ value
        out = out.transpose(1, 2).contiguous().view(batch, seq_len, channels)
        return self.resid_dropout(self.proj(out)), (key, value)


class FeedForward(nn.Module):
    def __init__(self, config: ModelConfig) -> None:
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(config.n_embd, 4 * config.n_embd),
            nn.GELU(),
            nn.Linear(4 * config.n_embd, config.n_embd),
            nn.Dropout(config.dropout),
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.net(x)


class TransformerBlock(nn.Module):
    def __init__(self, config: ModelConfig) -> None:
        super().__init__()
        self.ln_1 = nn.LayerNorm(config.n_embd)
        self.attn = CausalSelfAttention(config)
        self.ln_2 = nn.LayerNorm(config.n_embd)
        self.ffwd = FeedForward(config)

    def forward(
        self, x: torch.Tensor, return_attention: bool = False
    ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
        if return_attention:
            attn_out, weights = self.attn(self.ln_1(x), return_attention=True)
            x = x + attn_out
            x = x + self.ffwd(self.ln_2(x))
            return x, weights

        x = x + self.attn(self.ln_1(x))
        x = x + self.ffwd(self.ln_2(x))
        return x

    def forward_cached(
        self,
        x: torch.Tensor,
        past_key_value: tuple[torch.Tensor, torch.Tensor] | None = None,
    ) -> tuple[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]:
        attn_out, present = self.attn.forward_cached(self.ln_1(x), past_key_value)
        x = x + attn_out
        x = x + self.ffwd(self.ln_2(x))
        return x, present


class TinyTransformer(nn.Module):
    def __init__(self, config: ModelConfig) -> None:
        super().__init__()
        self.config = config
        self.token_embedding = nn.Embedding(config.vocab_size, config.n_embd)
        self.position_embedding = nn.Embedding(config.block_size, config.n_embd)
        self.dropout = nn.Dropout(config.dropout)
        self.blocks = nn.ModuleList([TransformerBlock(config) for _ in range(config.n_layer)])
        self.ln_f = nn.LayerNorm(config.n_embd)
        self.lm_head = nn.Linear(config.n_embd, config.vocab_size)
        self.apply(self._init_weights)
        self.lm_head.weight = self.token_embedding.weight

    def _init_weights(self, module: nn.Module) -> None:
        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)

    def forward(
        self, idx: torch.Tensor, targets: torch.Tensor | None = None
    ) -> tuple[torch.Tensor, torch.Tensor | None]:
        logits, loss, _ = self._forward(idx, targets, capture_attention=False)
        return logits, loss

    def _forward(
        self,
        idx: torch.Tensor,
        targets: torch.Tensor | None = None,
        capture_attention: bool = False,
    ) -> tuple[torch.Tensor, torch.Tensor | None, list[torch.Tensor]]:
        batch, seq_len = idx.shape
        if seq_len > self.config.block_size:
            raise ValueError("Sequence length exceeds block_size")

        attentions: list[torch.Tensor] = []
        positions = torch.arange(seq_len, device=idx.device)
        x = self.token_embedding(idx) + self.position_embedding(positions)
        x = self.dropout(x)
        for block in self.blocks:
            if capture_attention:
                x, weights = block(x, return_attention=True)
                attentions.append(weights)
            else:
                x = block(x)
        x = self.ln_f(x)
        logits = self.lm_head(x)

        loss = None
        if targets is not None:
            loss = F.cross_entropy(logits.view(batch * seq_len, -1), targets.view(batch * seq_len))
        return logits, loss, attentions

    def _forward_cached(
        self,
        idx: torch.Tensor,
        past_key_values: list[tuple[torch.Tensor, torch.Tensor]] | None = None,
    ) -> tuple[torch.Tensor, list[tuple[torch.Tensor, torch.Tensor]]]:
        _, seq_len = idx.shape
        past_len = 0 if past_key_values is None else past_key_values[0][0].size(2)
        if past_len + seq_len > self.config.block_size:
            raise ValueError("Cached sequence length exceeds block_size")

        positions = torch.arange(past_len, past_len + seq_len, device=idx.device)
        x = self.dropout(self.token_embedding(idx) + self.position_embedding(positions))
        presents = []
        for layer, block in enumerate(self.blocks):
            past = None if past_key_values is None else past_key_values[layer]
            x, present = block.forward_cached(x, past)
            presents.append(present)
        return self.lm_head(self.ln_f(x)), presents

    def parameter_count(self, trainable_only: bool = True) -> int:
        parameters = self.parameters()
        if trainable_only:
            parameters = (parameter for parameter in parameters if parameter.requires_grad)
        return sum(parameter.numel() for parameter in parameters)

    @torch.no_grad()
    def attention_maps(self, idx: torch.Tensor) -> list[torch.Tensor]:
        self.eval()
        _, _, attentions = self._forward(idx, capture_attention=True)
        return attentions

    @torch.no_grad()
    def generate(
        self,
        idx: torch.Tensor,
        max_new_tokens: int,
        temperature: float = 1.0,
        top_k: int | None = None,
        use_cache: bool = True,
    ) -> torch.Tensor:
        if max_new_tokens < 0:
            raise ValueError("max_new_tokens must be non-negative")
        if temperature <= 0:
            raise ValueError("temperature must be positive")
        if top_k is not None and top_k <= 0:
            raise ValueError("top_k must be positive")

        past_key_values = None
        for _ in range(max_new_tokens):
            if not use_cache:
                idx_cond = idx[:, -self.config.block_size :]
                logits, _ = self(idx_cond)
            elif past_key_values is None:
                idx_cond = idx[:, -self.config.block_size :]
                logits, past_key_values = self._forward_cached(idx_cond)
            else:
                cached_length = past_key_values[0][0].size(2)
                if cached_length >= self.config.block_size:
                    idx_cond = idx[:, -self.config.block_size :]
                    logits, past_key_values = self._forward_cached(idx_cond)
                else:
                    logits, past_key_values = self._forward_cached(
                        idx[:, -1:], past_key_values
                    )
            logits = logits[:, -1, :] / temperature
            if top_k is not None:
                values, _ = torch.topk(logits, min(top_k, logits.size(-1)))
                logits[logits < values[:, [-1]]] = -float("inf")
            probs = F.softmax(logits, dim=-1)
            next_idx = torch.multinomial(probs, num_samples=1)
            idx = torch.cat((idx, next_idx), dim=1)
        return idx