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

def precompute(head_dim : int, seq_len : int, device : str, base : float = 10000.0):
    assert head_dim % 2 == 0

    theta_pair = torch.arange(0,head_dim,2).float()
    theta = 1.0 / (base ** (theta_pair / head_dim)).to(device)

    m = torch.arange(seq_len,device=device)
    freqs = torch.outer(m,theta).float()
    
    freqs_complex = torch.polar(torch.ones_like(freqs),freqs)

    return freqs_complex

def apply_rope(x : torch.tensor, freqs_complex : torch.tensor, device:str):
    x_complex = torch.view_as_complex(x.float().reshape(*x.shape[:-1],-1,2))
    freqs_complex = freqs_complex.unsqueeze(0).unsqueeze(0)

    x_rotated = x_complex * freqs_complex
    x_out = torch.view_as_real(x_rotated)
    x_out = x_out.reshape(*x.shape)

    return x_out.type_as(x)


class CasualSelfAttention(nn.Module):
    def __init__(self,config):
        super().__init__()

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

        self.n_head = config.n_head
        self.n_embd = config.n_embd


    def forward(self,x,freqs):
        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, C // self.n_head).transpose(1,2)
        k = k.view(B,T, self.n_head, C // self.n_head).transpose(1,2)
        v = v.view(B,T , self.n_head, C // self.n_head).transpose(1,2)

        q = apply_rope(q, freqs, x.device)
        k = apply_rope(k, freqs, x.device)

        y = F.scaled_dot_product_attention(q,k,v,is_causal=True)
        y = y.transpose(1,2).contiguous().view(B,T,C)

        y = self.c_proj(y)


        return y
    


class FeedForward(nn.Module):
    def __init__(self,config):
        super().__init__()
        hidden_dim = int(2 * config.n_embd / 3)
        hidden_dim = config.multiple_of * ((hidden_dim + config.multiple_of - 1) // config.multiple_of)

        self.w1 = nn.Linear(config.n_embd,hidden_dim, bias=False)
        self.w2 = nn.Linear(hidden_dim,config.n_embd, bias = False)
        self.w3  = nn.Linear(config.n_embd, hidden_dim, bias= False)

    def forward(self,x : torch.tensor):
        return self.w2(F.silu(self.w1(x)) * self.w3(x))
    

class Block(nn.Module):
    def __init__(self,config):
        super().__init__()

        self.sa = CasualSelfAttention(config)
        self.mlp = FeedForward(config)
        self.ln1 = nn.RMSNorm(config.n_embd,eps=1e-05)
        self.ln2 = nn.RMSNorm(config.n_embd,eps =1e-05)


    def forward(self,x,freqs):
        x = x + self.sa(self.ln1(x),freqs)
        x = x + self.mlp(self.ln2(x))

        return x
    

class Model(nn.Module):
    def __init__(self,config):
        super().__init__()

        self.register_buffer(
            'freqs_complex',precompute(
                config.n_embd // config.n_head,
                config.block_size,
                device=config.device
            )
        )

        self.config = config
        self.transformer = nn.ModuleDict(dict(
            wte = nn.Embedding(config.vocab_size,config.n_embd),
            h = nn.ModuleList([Block(config) for _ in range(config.n_layer)]),
            ln_f = nn.RMSNorm(config.n_embd,eps = 1e-05)
        ))
        self.lm_head = nn.Linear(config.n_embd, config.vocab_size,bias=False)


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

    def forward(self,idx,targets=None):
        B,T = idx.size()

        x = self.transformer.wte(idx)

        for block in self.transformer.h:
            x = block(x,self.freqs_complex[:T])

        x = self.transformer.ln_f(x)
        logits = self.lm_head(x)


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

        return logits,loss


    def sample_top_p(self,probs,p):
        probs_sort, probs_idx = torch.sort(probs,dim=-1,descending=True)
        probs_sum = torch.cumsum(probs_sort,dim=-1)
        mask = probs_sum - probs_sort > p
        probs_sort[mask] = 0.0
        probs_sort.div_(probs_sort.sum(dim=-1,keepdim=True))

        next_token = torch.multinomial(probs_sort,num_samples=1)
        next_token = torch.gather(probs_idx,-1,next_token)

        return next_token

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

        for _ in range(max_new_tokens):

            idx_cond = idx if idx.size(1) <= self.config.block_size else idx[:, -self.config.block_size:]

            logits, _ = self(idx_cond)

            logits = logits[:, -1, :] / 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)

            
            if top_p is not None:
                idx_next = self.sample_top_p(probs,top_p)
            else:
                idx_next = torch.multinomial(probs, num_samples=1)

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

        return idx