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

import tiktoken
enc = tiktoken.get_encoding("gpt2")
print("tiktoken cached ✅")

@dataclass
class GPTConfig:
    block_size:int=512
    vocab_size:int=50304
    n_layer:int=10
    n_head:int=8
    n_embd:int=256
    dropout:float=0.0
    bias:bool=False

class LayerNorm(nn.Module):
    def __init__(self,ndim,bias):
        super().__init__()
        self.weight=nn.Parameter(torch.ones(ndim))
        self.bias=nn.Parameter(torch.zeros(ndim)) if bias else None

    def forward(self,input):
        return F.layer_norm(input,self.weight.shape, self.weight,self.bias,1e-5)

class CausalSelfAttention(nn.Module):
    def __init__(self,config):
        super().__init__()
        assert config.n_embd % config.n_head==0
        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)
        self.n_head = config.n_head
        self.n_embd = config.n_embd
        self.dropout = config.dropout
        self.flash=hasattr(torch.nn.functional,'scaled_dot_product_attention')
        if not self.flash:
            print("WARNING using slow attention")
            self.register_buffer("bias",torch.tril(torch.ones(config.block_size,config.block_size)).view(1,1,config.block_size,config.block_size))

    def forward(self,x):
        B,T,C=x.size()
        
        q, k, v=self.c_attn(x).split(self.n_embd,dim=2)
        k=k.view(B,T,self.n_head, C //self.n_head).transpose(1,2)
        q=q.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)
        
        if self.flash:
            y=torch.nn.functional.scaled_dot_product_attention(q, k, v,attn_mask=None, dropout_p=self.dropout if self.training else 0,is_causal=True)

        else:
            att=(q@k.transpose(-2,-1))*(1.0/math.sqrt(k.size(-1)))
            att=att.masked_fill(self.bias[:,:,:T,:T]==0,float('-inf'))
            att=F.softmax(att,dim=-1)
            att=self.attn_dropout(att)
            y=att@v
        y=y.transpose(1,2).contiguous().view(B, T, C)
        y=self.resid_dropout(self.c_proj(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=LayerNorm(config.n_embd,bias=config.bias)
        self.attn=CausalSelfAttention(config)
        self.ln_2=LayerNorm(config.n_embd,bias=config.bias)
        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__()
        assert config.vocab_size is not None
        assert config.block_size is not None
        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=LayerNorm(config.n_embd, bias=config.bias),
        ))
        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)
        for pn,p in self.named_parameters():
            if pn.endswith('c_proj.weight'):
                torch.nn.init.normal_(p,mean=0.0,std=0.02/math.sqrt(2*config.n_layer))
        print("Number of Parameters: %.2fM" % (self.get_num_params()/1e6))
    
    def get_num_params(self,non_embedding=True):
        n_params=sum(p.numel() for p in self.parameters())
        if non_embedding:
            n_params-=self.transformer.wpe.weight.numel()
        return n_params

    def _init_weights(self,module):
        if isinstance(module,nn.Linear):
            torch.nn.init.normal_(module.weight, mean=0.0,std=0.02)
            if module.bias is not None:
                torch.nn.init.zeros_(module.bias)
        elif isinstance(module,nn.Embedding):
            torch.nn.init.normal_(module.weight,mean=0.0,std=0.02)

    def forward(self,idx,targets=None):
        device =idx.device
        b,t=idx.size()

        assert t<=self.config.block_size,f"Cannot forward Sequence of Length {t}  block size is only {self.config.block_size}"
        pos=torch.arange(0,t,dtype=torch.long, device=device)
        tok_emb=self.transformer.wte(idx)
        pos_emb=self.transformer.wpe(pos)
        x=self.transformer.drop(tok_emb + pos_emb)
        for block in self.transformer.h:
            x=block(x)
        x=self.transformer.ln_f(x)
        if targets is not None:
            logits=self.lm_head(x)
            loss=F.cross_entropy(logits.view(-1,logits.size(-1)),targets.view(-1),ignore_index=-1)
        else:
            logits=self.lm_head(x[:,[-1],:])
            loss=None
        return logits,loss
    def configure_optimizers(self, weight_decay,learning_rate,betas,device_type):
        param_dict={pn:p for pn, p in self.named_parameters()}
        param_dict={pn:p for pn, p in param_dict.items() if p.requires_grad}
        
    
        decay_params=[p for n,p in param_dict.items() if p.dim()>=2]
        ndecay_params=[p for n, p in param_dict.items() if p.dim()<2]
        optim_grps=[
            {'params':decay_params,'weight_decay':weight_decay},
            {'params':ndecay_params,'weight_decay':0.0}
        ]
        num_decay_params=sum(p.numel() for p in decay_params)
        num_nodecay_params=sum(p.numel() for p in ndecay_params)
        print(f"num decayed parameter tensors: {len(decay_params)}, with {num_decay_params:,} parameters")
        print(f"num non-decayed parameter tensors: {len(ndecay_params)}, with {num_nodecay_params:,} parameters")
        fused_available='fused' in inspect.signature(torch.optim.AdamW).parameters
        use_fused = fused_available and device_type=='cuda'
        extra_args=dict(fused=True) if use_fused else dict()
        optimizer = torch.optim.AdamW(optim_grps,lr=learning_rate,betas=betas, **extra_args)
        print(f"Using Fused AdamW: {use_fused}")
        return optimizer
        
    @torch.no_grad()
    def generate(self,idx,max_new_tokens,temperature=1.0,top_k=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)
            idx_next=torch.multinomial(probs,num_samples=1)
            if idx_next.item()==50256:
                break
            idx =torch.cat((idx,idx_next),dim=1)
        return idx