from dataclasses import asdict, dataclass

import torch
from torch import nn
from torch.nn import functional as functional
from torch.utils.checkpoint import checkpoint


@dataclass
class ModelConfig:
    vocab_size: int = 8192
    context_length: int = 1024
    layers: int = 8
    hidden_size: int = 512
    heads: int = 8
    intermediate_size: int = 1408

    def __post_init__(self):
        if any(value <= 0 for value in asdict(self).values()):
            raise ValueError("All model dimensions must be positive")
        if self.hidden_size % self.heads or (self.hidden_size // self.heads) % 2:
            raise ValueError("Head dimension must be an even integer")


class Attention(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.heads = config.heads
        self.head_size = config.hidden_size // config.heads
        self.qkv = nn.Linear(config.hidden_size, 3 * config.hidden_size, bias=False)
        self.output = nn.Linear(config.hidden_size, config.hidden_size, bias=False)
        frequencies = 1.0 / (10000 ** (torch.arange(0, self.head_size, 2).float() / self.head_size))
        angles = torch.outer(torch.arange(config.context_length).float(), frequencies)
        self.register_buffer("cosines", angles.cos(), persistent=False)
        self.register_buffer("sines", angles.sin(), persistent=False)

    def rotate(self, tensor):
        length = tensor.size(-2)
        cosines = self.cosines[:length].to(tensor.dtype)
        sines = self.sines[:length].to(tensor.dtype)
        even, odd = tensor[..., 0::2], tensor[..., 1::2]
        return torch.stack((even * cosines - odd * sines, even * sines + odd * cosines), dim=-1).flatten(-2)

    def forward(self, hidden):
        batch_size, length, width = hidden.shape
        query, key, value = self.qkv(hidden).reshape(batch_size, length, 3, self.heads, self.head_size).unbind(2)
        attended = functional.scaled_dot_product_attention(
            self.rotate(query.transpose(1, 2)),
            self.rotate(key.transpose(1, 2)),
            value.transpose(1, 2),
            is_causal=True,
        )
        return self.output(attended.transpose(1, 2).contiguous().view(batch_size, length, width))


class Block(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.attention_norm = nn.RMSNorm(config.hidden_size, eps=1e-5)
        self.attention = Attention(config)
        self.ffn_norm = nn.RMSNorm(config.hidden_size, eps=1e-5)
        self.gate = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
        self.up = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
        self.down = nn.Linear(config.intermediate_size, config.hidden_size, bias=False)

    def forward(self, hidden):
        hidden = hidden + self.attention(self.attention_norm(hidden))
        normalized = self.ffn_norm(hidden)
        return hidden + self.down(functional.silu(self.gate(normalized)) * self.up(normalized))


class LanguageModel(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.config = config
        self.embedding = nn.Embedding(config.vocab_size, config.hidden_size)
        self.blocks = nn.ModuleList([Block(config) for _ in range(config.layers)])
        self.norm = nn.RMSNorm(config.hidden_size, eps=1e-5)
        self.gradient_checkpointing = False
        self.apply(self.initialize)

    def set_gradient_checkpointing(self, enabled):
        self.gradient_checkpointing = enabled

    @staticmethod
    def initialize(module):
        if isinstance(module, (nn.Linear, nn.Embedding)):
            nn.init.normal_(module.weight, mean=0.0, std=0.02)

    def forward(self, tokens, targets=None, return_logits=False, logit_positions=None):
        if tokens.size(1) > self.config.context_length:
            raise ValueError("Input exceeds context length")
        hidden = self.embedding(tokens)
        for block in self.blocks:
            if self.gradient_checkpointing and self.training and torch.is_grad_enabled():
                hidden = checkpoint(block, hidden, use_reentrant=False)
            else:
                hidden = block(hidden)
        hidden = self.norm(hidden)
        if logit_positions is not None:
            hidden = hidden[:, logit_positions, :]
        if targets is None and not return_logits:
            return functional.linear(hidden[:, -1:, :], self.embedding.weight)
        logits = functional.linear(hidden, self.embedding.weight)
        if targets is None:
            return logits
        return functional.cross_entropy(logits.reshape(-1, self.config.vocab_size), targets.reshape(-1))