import math

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


class RMSNorm(nn.Module):
    def __init__(self, dim, eps=1e-6):
        super().__init__()
        self.weight = nn.Parameter(torch.ones(dim))
        self.eps = eps

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


def rotate_half(x):
    x1 = x[..., : x.shape[-1] // 2]
    x2 = x[..., x.shape[-1] // 2 :]
    return torch.cat((-x2, x1), dim=-1)


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


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

        inv_freq = 1.0 / (
            base ** (
                torch.arange(0, head_dim, 2).float() / head_dim
            )
        )

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

        positions = torch.arange(max_seq_len).float()

        freqs = torch.outer(positions, inv_freq)

        emb = torch.cat((freqs, freqs), dim=-1)

        self.register_buffer(
            "cos_cached",
            emb.cos()[None, None, :, :],
            persistent=False,
        )

        self.register_buffer(
            "sin_cached",
            emb.sin()[None, None, :, :],
            persistent=False,
        )

    def forward(self, seq_len):
        return (
            self.cos_cached[:, :, :seq_len, :],
            self.sin_cached[:, :, :seq_len, :],
        )


class SwiGLU(nn.Module):
    def __init__(self, hidden_size, intermediate_size):
        super().__init__()

        self.gate_proj = nn.Linear(
            hidden_size,
            intermediate_size,
            bias=False,
        )

        self.up_proj = nn.Linear(
            hidden_size,
            intermediate_size,
            bias=False,
        )

        self.down_proj = nn.Linear(
            intermediate_size,
            hidden_size,
            bias=False,
        )

    def forward(self, x):
        return self.down_proj(
            F.silu(self.gate_proj(x)) * self.up_proj(x)
        )


class GQAttention(nn.Module):
    def __init__(
        self,
        hidden_size,
        num_heads,
        num_kv_heads,
        max_seq_len,
    ):
        super().__init__()

        assert hidden_size % num_heads == 0

        self.num_heads = num_heads
        self.num_kv_heads = num_kv_heads
        self.head_dim = hidden_size // num_heads

        assert num_heads % num_kv_heads == 0

        self.q_proj = nn.Linear(
            hidden_size,
            num_heads * self.head_dim,
            bias=False,
        )

        self.k_proj = nn.Linear(
            hidden_size,
            num_kv_heads * self.head_dim,
            bias=False,
        )

        self.v_proj = nn.Linear(
            hidden_size,
            num_kv_heads * self.head_dim,
            bias=False,
        )

        self.o_proj = nn.Linear(
            hidden_size,
            hidden_size,
            bias=False,
        )

        self.rope = RotaryEmbedding(
            self.head_dim,
            max_seq_len,
        )

    def forward(self, x):
        batch, seq_len, hidden = x.shape

        q = self.q_proj(x)
        k = self.k_proj(x)
        v = self.v_proj(x)

        q = q.view(
            batch,
            seq_len,
            self.num_heads,
            self.head_dim,
        ).transpose(1, 2)

        k = k.view(
            batch,
            seq_len,
            self.num_kv_heads,
            self.head_dim,
        ).transpose(1, 2)

        v = v.view(
            batch,
            seq_len,
            self.num_kv_heads,
            self.head_dim,
        ).transpose(1, 2)

        cos, sin = self.rope(seq_len)

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

        # Repeat KV heads for grouped-query attention.
        repeat_factor = self.num_heads // self.num_kv_heads

        k = k.repeat_interleave(
            repeat_factor,
            dim=1,
        )

        v = v.repeat_interleave(
            repeat_factor,
            dim=1,
        )

        y = F.scaled_dot_product_attention(
            q,
            k,
            v,
            is_causal=True,
        )

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

        y = y.view(
            batch,
            seq_len,
            hidden,
        )

        return self.o_proj(y)


class TransformerBlock(nn.Module):
    def __init__(
        self,
        hidden_size,
        num_heads,
        num_kv_heads,
        intermediate_size,
        max_seq_len,
    ):
        super().__init__()

        self.attn_norm = RMSNorm(hidden_size)
        self.attn = GQAttention(
            hidden_size,
            num_heads,
            num_kv_heads,
            max_seq_len,
        )

        self.ffn_norm = RMSNorm(hidden_size)

        self.ffn = SwiGLU(
            hidden_size,
            intermediate_size,
        )

    def forward(self, x):
        x = x + self.attn(
            self.attn_norm(x)
        )

        x = x + self.ffn(
            self.ffn_norm(x)
        )

        return x


class VegaLM(nn.Module):
    def __init__(
        self,
        vocab_size=32000,
        hidden_size=512,
        num_layers=12,
        num_heads=8,
        num_kv_heads=4,
        intermediate_size=880,
        max_seq_len=2048,
    ):
        super().__init__()

        self.vocab_size = vocab_size
        self.hidden_size = hidden_size

        self.token_embedding = nn.Embedding(
            vocab_size,
            hidden_size,
        )

        self.layers = nn.ModuleList(
            [
                TransformerBlock(
                    hidden_size=hidden_size,
                    num_heads=num_heads,
                    num_kv_heads=num_kv_heads,
                    intermediate_size=intermediate_size,
                    max_seq_len=max_seq_len,
                )
                for _ in range(num_layers)
            ]
        )

        self.final_norm = RMSNorm(hidden_size)

        self.lm_head = nn.Linear(
            hidden_size,
            vocab_size,
            bias=False,
        )

        # Tie input embeddings and output projection.
        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,
            )

        elif isinstance(module, nn.Embedding):
            nn.init.normal_(
                module.weight,
                mean=0.0,
                std=0.02,
            )

    def forward(self, input_ids, targets=None):
        x = self.token_embedding(input_ids)

        for layer in self.layers:
            x = layer(x)

        x = self.final_norm(x)

        logits = self.lm_head(x)

        loss = None

        if targets is not None:
            loss = F.cross_entropy(
                logits.view(-1, self.vocab_size),
                targets.view(-1),
            )

        return logits, loss


def count_parameters(model):
    return sum(
        p.numel()
        for p in model.parameters()
        if p.requires_grad
    )


if __name__ == "__main__":
    model = VegaLM()

    params = count_parameters(model)

    print(f"Parameters: {params:,}")
    print(f"Parameters: {params / 1e6:.2f}M")

    x = torch.randint(
        0,
        32000,
        (2, 128),
    )

    logits, _ = model(x)

    print(f"Input:  {x.shape}")
    print(f"Output: {logits.shape}")