import math
from dataclasses import dataclass

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


@dataclass
class GPTConfig:
    vocab_size: int = 50257
    block_size: int = 512

    n_layer: int = 24
    n_head: int = 20
    n_embd: int = 1280

    dropout: float = 0.1
    bias: bool = False


class CausalSelfAttention(nn.Module):

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

        assert config.n_embd % config.n_head == 0

        self.n_head = config.n_head
        self.n_embd = config.n_embd
        self.head_dim = config.n_embd // config.n_head

        self.c_attn = nn.Linear(
            config.n_embd,
            3 * config.n_embd,
            bias=config.bias
        )

        self.c_proj = nn.Linear(
            config.n_embd,
            config.n_embd,
            bias=config.bias
        )

        self.attn_dropout = nn.Dropout(config.dropout)
        self.resid_dropout = nn.Dropout(config.dropout)

    def forward(self, x):

        B, T, C = x.size()

        qkv = self.c_attn(x)

        q, k, v = qkv.split(self.n_embd, dim=2)

        q = q.view(
            B,
            T,
            self.n_head,
            self.head_dim
        ).transpose(1, 2)

        k = k.view(
            B,
            T,
            self.n_head,
            self.head_dim
        ).transpose(1, 2)

        v = v.view(
            B,
            T,
            self.n_head,
            self.head_dim
        ).transpose(1, 2)

        y = F.scaled_dot_product_attention(
            q,
            k,
            v,
            dropout_p=self.attn_dropout.p if self.training else 0.0,
            is_causal=True
        )

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

        y = y.view(B, T, C)

        y = self.c_proj(y)

        y = self.resid_dropout(y)

        return y


class MLP(nn.Module):

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

        self.c_fc = nn.Linear(
            config.n_embd,
            4 * config.n_embd,
            bias=config.bias
        )

        self.gelu = nn.GELU()

        self.c_proj = nn.Linear(
            4 * config.n_embd,
            config.n_embd,
            bias=config.bias
        )

        self.dropout = nn.Dropout(config.dropout)

    def forward(self, x):

        x = self.c_fc(x)
        x = self.gelu(x)
        x = self.c_proj(x)
        x = self.dropout(x)

        return x


class Block(nn.Module):

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

        self.ln_1 = nn.LayerNorm(config.n_embd)

        self.attn = CausalSelfAttention(config)

        self.ln_2 = nn.LayerNorm(config.n_embd)

        self.mlp = MLP(config)

    def forward(self, x):

        x = x + self.attn(
            self.ln_1(x)
        )

        x = x + self.mlp(
            self.ln_2(x)
        )

        return x


class GPT(nn.Module):

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

        self.config = config

        self.transformer = nn.ModuleDict(
            dict(
                wte=nn.Embedding(
                    config.vocab_size,
                    config.n_embd
                ),

                wpe=nn.Embedding(
                    config.block_size,
                    config.n_embd
                ),

                drop=nn.Dropout(
                    config.dropout
                ),

                h=nn.ModuleList(
                    [
                        Block(config)
                        for _ in range(config.n_layer)
                    ]
                ),

                ln_f=nn.LayerNorm(
                    config.n_embd
                ),
            )
        )

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

        self.transformer.wte.weight = self.lm_head.weight

        self.apply(self._init_weights)

        print(
            f"Model Parameters: "
            f"{self.get_num_params()/1e6:.2f}M"
        )

    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
            )

    def get_num_params(self):

        return sum(
            p.numel()
            for p in self.parameters()
        )

    def forward(
        self,
        idx,
        targets=None
    ):

        B, T = idx.size()

        assert (
            T <= self.config.block_size
        ), "Sequence too long"

        pos = torch.arange(
            0,
            T,
            device=idx.device
        )

        tok_emb = self.transformer.wte(idx)

        pos_emb = self.transformer.wpe(pos)

        x = tok_emb + pos_emb

        x = self.transformer.drop(x)

        for block in self.transformer.h:
            x = block(x)

        x = self.transformer.ln_f(x)

        logits = self.lm_head(x)

        loss = None

        if targets is not None:

            loss = F.cross_entropy(
                logits.reshape(
                    -1,
                    logits.size(-1)
                ),
                targets.reshape(-1)
            )

        return logits, loss

    @torch.no_grad()
    def generate(
        self,
        idx,
        max_new_tokens,
        temperature=1.0,
        top_k=50
    ):

        for _ in range(max_new_tokens):

            idx_cond = idx[
                :,
                -self.config.block_size:
            ]

            logits, _ = self(idx_cond)

            logits = logits[:, -1, :]

            logits = logits / temperature

            if top_k is not None:

                v, _ = torch.topk(
                    logits,
                    min(
                        top_k,
                        logits.size(-1)
                    )
                )

                logits[
                    logits < v[:, [-1]]
                ] = -float("inf")

            probs = F.softmax(
                logits,
                dim=-1
            )

            idx_next = torch.multinomial(
                probs,
                num_samples=1
            )

            idx = torch.cat(
                (idx, idx_next),
                dim=1
            )

        return idx


if __name__ == "__main__":

    config = GPTConfig()

    model = GPT(config)

    print(
        f"Total Parameters: "
        f"{model.get_num_params():,}"
    )