import torch
import torch.nn as nn
import torch.nn.functional as F
import math
from dataclasses import dataclass
from typing import Optional


@dataclass
class ModelConfig:
    vocab_size: int
    hidden_size: int
    n_heads: int
    n_kv_heads: int
    n_kv_groups: int
    head_dim: int
    n_layers: int
    attention_bias: bool
    intermediate_size: int
    mlp_bias: bool
    eps: float
    dropout: float
    max_position_embeddings: int
    pre_norm: bool
    tie_weights: bool
    max_seq_len: int


class RMSNorm(nn.Module):
    def __init__(self, config: ModelConfig):
        super().__init__()
        self.eps = config.eps
        self.weight = nn.Parameter(torch.ones(config.hidden_size))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        rms = torch.sqrt(torch.mean(x * x, dim=-1, keepdim=True) + self.eps)
        return (x / rms) * self.weight


class RotaryEmbedding(nn.Module):
    def __init__(self, head_dim, max_position_embeddings=2048):
        super().__init__()
        inv_freq = 1.0 / (10000 ** (torch.arange(0, head_dim, 2).float() / head_dim))
        t = torch.arange(max_position_embeddings, dtype=torch.float32)
        freqs = torch.einsum("i,j->ij", t, inv_freq)
        emb = torch.cat((freqs, freqs), dim=-1)
        self.register_buffer("cos", emb.cos()[None, None, :, :], persistent=False)
        self.register_buffer("sin", emb.sin()[None, None, :, :], persistent=False)

    def forward(self, seq_len: int, device: torch.device, dtype: torch.dtype):
        cos = self.cos[:, :, :seq_len, :].to(device=device, dtype=dtype)
        sin = self.sin[:, :, :seq_len, :].to(device=device, dtype=dtype)
        return cos, sin


def apply_rotary(x, cos, sin):
    x1, x2 = x[..., ::2], x[..., 1::2]
    x_rot = torch.stack([-x2, x1], dim=-1).reshape_as(x)
    return (x * cos) + (x_rot * sin)


class GroupedMultiQueryAttention(nn.Module):
    def __init__(self, config: ModelConfig):
        super().__init__()
        self.hidden_size = config.hidden_size
        self.n_heads = config.n_heads
        self.n_kv_heads = config.n_kv_heads
        self.head_dim = config.head_dim
        self.attention_bias = config.attention_bias
        self.dropout = nn.Dropout(config.dropout)

        if self.n_heads * self.head_dim != self.hidden_size:
            raise ValueError("hidden_size must equal n_heads * head_dim")

        # derive n_kv_groups if None
        if config.n_kv_groups is None:
            if self.n_kv_heads == 0:
                raise ValueError("n_kv_heads must be > 0")
            self.n_kv_groups = self.n_heads // self.n_kv_heads
            if self.n_heads % self.n_kv_heads != 0:
                raise ValueError("n_heads must be divisible by n_kv_heads to derive groups")
        else:
            self.n_kv_groups = config.n_kv_groups
            if self.n_kv_heads * self.n_kv_groups != self.n_heads:
                raise ValueError("n_heads must equal n_kv_heads * n_kv_groups")

        self.q_proj = nn.Linear(self.hidden_size, self.n_heads * self.head_dim, bias=self.attention_bias)
        self.k_proj = nn.Linear(self.hidden_size, self.n_kv_heads * self.head_dim, bias=self.attention_bias)
        self.v_proj = nn.Linear(self.hidden_size, self.n_kv_heads * self.head_dim, bias=self.attention_bias)
        self.w_o = nn.Linear(self.hidden_size, self.hidden_size, bias=False)
        self.rope = RotaryEmbedding(self.head_dim, config.max_position_embeddings)

    def forward(self, x):
        B, T, _ = x.shape
        device = x.device
        dtype = x.dtype

        q = self.q_proj(x).view(B, T, self.n_heads, self.head_dim).transpose(1, 2)
        k = self.k_proj(x).view(B, T, self.n_kv_heads, self.head_dim).transpose(1, 2)
        v = self.v_proj(x).view(B, T, self.n_kv_heads, self.head_dim).transpose(1, 2)

        cos, sin = self.rope(T, device=device, dtype=dtype)
        q = apply_rotary(q, cos, sin)
        k = apply_rotary(k, cos, sin)

        if self.n_kv_groups != 1:
            k = k.repeat_interleave(self.n_kv_groups, dim=1)
            v = v.repeat_interleave(self.n_kv_groups, dim=1)

        scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.head_dim)

        # causal mask
        mask = torch.triu(torch.full((T, T), float("-inf"), device=device, dtype=dtype), diagonal=1)
        scores = scores + mask.unsqueeze(0).unsqueeze(0)
        attn = torch.softmax(scores, dim=-1)
        attn = self.dropout(attn)

        out = torch.matmul(attn, v)
        out = out.transpose(1, 2).contiguous().view(B, T, self.hidden_size)
        return self.w_o(out)


class SwiGLUFeedForward(nn.Module):
    def __init__(self, config: ModelConfig):
        super().__init__()
        self.hidden_size = config.hidden_size
        self.intermediate_size = config.intermediate_size
        self.dropout = nn.Dropout(config.dropout)

        self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size)
        self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size)
        self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size)
        self.act = nn.SiLU()

    def forward(self, x):
        x = self.act(self.gate_proj(x)) * self.up_proj(x)
        x = self.down_proj(self.dropout(x))
        return x


class TransformerBlock(nn.Module):
    def __init__(self, config: ModelConfig):
        super().__init__()
        self.attention = GroupedMultiQueryAttention(config)
        self.feed_forward = SwiGLUFeedForward(config)
        self.attn_norm = RMSNorm(config)
        self.ffn_norm = RMSNorm(config)
        self.dropout = nn.Dropout(config.dropout)
        self.pre_norm = config.pre_norm

    def forward(self, x):
        if self.pre_norm:
            x = x + self.dropout(self.attention(self.attn_norm(x)))
            x = x + self.dropout(self.feed_forward(self.ffn_norm(x)))
        else:
            x = self.attn_norm(x + self.dropout(self.attention(x)))
            x = self.ffn_norm(x + self.dropout(self.feed_forward(x)))
        return x


class Transformer(nn.Module):
    def __init__(self, config: ModelConfig):
        super().__init__()
        self.config = config
        self.token_embedding = nn.Embedding(config.vocab_size, config.hidden_size)
        self.blocks = nn.ModuleList([TransformerBlock(config) for _ in range(config.n_layers)])
        self.embedding_dropout = nn.Dropout(config.dropout)
        self.final_norm = RMSNorm(config)
        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
        if config.tie_weights:
            self.lm_head.weight = self.token_embedding.weight

        self.apply(self._init_weights)

    def _init_weights(self, module):
        if isinstance(module, nn.Linear):
            nn.init.normal_(module.weight, mean=0.0, std=0.02 / math.sqrt(max(1, self.config.n_layers)))
            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, input_ids: torch.Tensor, targets: Optional[torch.Tensor] = None):
        x = self.token_embedding(input_ids) * math.sqrt(self.config.hidden_size)
        x = self.embedding_dropout(x)
        for block in self.blocks:
            x = block(x)
        x = self.final_norm(x)
        logits = self.lm_head(x)
        return logits

def top_k_top_p_filtering(logits: torch.Tensor, top_k: int = 0, top_p: float = 0.0, filter_value: float = -float('Inf')) -> torch.Tensor:
    """
    Filter a distribution of logits using top-k and/or nucleus (top-p) filtering.
    This is taken from common implementations (Hugging Face transformers style).
    Args:
        logits: logits distribution shape (batch, vocab)
        top_k: keep only top k tokens with highest probability (0 = no top-k)
        top_p: keep the top tokens with cumulative probability >= top_p (0.0 = no nucleus)
        filter_value: value to set for filtered logits
    Returns:
        filtered logits with the same shape
    """
    top_k = max(top_k, 0)
    batch_size, vocab_size = logits.size()

    if top_k > 0:
        # Remove all tokens with a probability less than the top-k tokens
        top_k = min(max(top_k, 1), vocab_size)
        values_to_keep, _ = torch.topk(logits, top_k)
        min_values = values_to_keep[:, -1].unsqueeze(1).expand_as(logits)
        logits = torch.where(logits < min_values, torch.full_like(logits, filter_value), logits)

    if top_p > 0.0:
        sorted_logits, sorted_indices = torch.sort(logits, descending=True, dim=-1)
        sorted_probs = F.softmax(sorted_logits, dim=-1)
        cumulative_probs = torch.cumsum(sorted_probs, dim=-1)

        # Remove tokens with cumulative probability above the threshold
        sorted_mask = cumulative_probs > top_p

        # Shift the mask right to keep at least one token
        sorted_mask[..., 1:] = sorted_mask[..., :-1].clone()
        sorted_mask[..., 0] = False

        indices_to_remove = sorted_mask.scatter(1, sorted_indices, sorted_mask)
        logits = logits.masked_fill(indices_to_remove, filter_value)

    return logits


@torch.no_grad()
def generate(
    model: Transformer,
    input_ids: torch.LongTensor,
    max_new_tokens: int = 50,
    temperature: float = 1.0,
    top_k: int = 0,
    top_p: float = 0.0,
    do_sample: bool = True,
    eos_token_id: Optional[int] = None,
    pad_token_id: Optional[int] = None,
    device: Optional[torch.device] = None,
):
    """
    Autoregressive generation helper for the model. This implementation does NOT use KV cache
    (the model defined in this file does not implement a cache), so generation is performed
    by repeatedly calling the model on the growing sequence. It supports temperature,
    top-k and nucleus (top-p) sampling, greedy decoding, and optional early stopping
    on an `eos_token_id`.

    Args:
        model: the Transformer instance
        input_ids: (batch, seq_len) input token ids
        max_new_tokens: number of tokens to generate
        temperature: sampling temperature (<=0 or do_sample=False => greedy)
        top_k: top-k filtering (0 disables)
        top_p: nucleus/top-p filtering (0.0 disables)
        do_sample: whether to sample (True) or do greedy decoding (False)
        eos_token_id: optional EOS id to stop generation for individual sequences
        pad_token_id: optional pad id to use for finished sequences
        device: optional torch.device to run on; if None uses model's device

    Returns:
        tensor of shape (batch, seq_len + generated) with generated tokens appended
    """
    model.eval()
    if device is None:
        # try to infer device
        try:
            device = next(model.parameters()).device
        except StopIteration:
            device = torch.device('cpu')

    input_ids = input_ids.to(device)

    batch_size, seq_len = input_ids.shape
    generated = 0
    unfinished = torch.ones(batch_size, dtype=torch.bool, device=device)

    for _ in range(max_new_tokens):
        logits = model(input_ids)
        # logits shape: (batch, seq_len_total, vocab)
        next_token_logits = logits[:, -1, :]

        if temperature <= 0 or not do_sample:
            # Greedy
            next_tokens = torch.argmax(next_token_logits, dim=-1)
        else:
            logits_proc = next_token_logits / max(temperature, 1e-8)
            logits_proc = top_k_top_p_filtering(logits_proc, top_k=top_k, top_p=top_p)
            probs = F.softmax(logits_proc, dim=-1)
            next_tokens = torch.multinomial(probs, num_samples=1).squeeze(-1)

        # If EOS is provided, update finished sequences and pad further tokens
        if eos_token_id is not None:
            is_eos = next_tokens.eq(eos_token_id)
            # sequences that have just finished
            just_finished = unfinished & is_eos
            unfinished = unfinished & (~is_eos)

        # For sequences already finished, append pad_token_id (if provided), otherwise keep EOS or sampled token
        if pad_token_id is not None and not unfinished.all():
            finished_mask = ~unfinished
            if finished_mask.any():
                next_tokens = next_tokens.masked_fill(finished_mask, pad_token_id)

        # append
        input_ids = torch.cat([input_ids, next_tokens.unsqueeze(-1)], dim=1)
        generated += 1

        if eos_token_id is not None and not unfinished.any():
            break

    return input_ids


def _smoke_test():
    config = ModelConfig(
        vocab_size=128,
        hidden_size=64,
        n_heads=4,
        n_kv_heads=4,
        n_kv_groups=None,
        head_dim=16,
        n_layers=2,
        attention_bias=False,
        intermediate_size=256,
        mlp_bias=False,
        eps=1e-5,
    )
    model = Transformer(config)
    model.eval()

    batch, seq_len = 2, 8
    input_ids = torch.randint(0, config.vocab_size, (batch, seq_len))
    logits, loss = model(input_ids, targets=input_ids)

    assert logits.shape == (batch, seq_len, config.vocab_size)
    assert loss.dim() == 0
    print("Smoke test passed: logits shape", logits.shape, "loss", loss.detach().item())


if __name__ == "__main__":
    _smoke_test()