import torch
import torch.nn as nn
import torch.nn.functional as F
from data_utils import *

from attention import SelfAttentionHead, MultiHeadAttention, FeedForwardNet, DecoderBlock


class BigramLanguageModel(nn.Module):
    def __init__(self, vocab_size, n_embed, block_size, num_heads, n_layers) -> None:
        super().__init__()
        self.token_embedding_table = nn.Embedding(vocab_size, n_embed)
        self.position_embedding_table = nn.Embedding(block_size, n_embed)
        self.decoder_blocks = nn.Sequential(*[DecoderBlock(n_embed, num_heads, block_size=block_size) for _ in range(n_layers)] )
        self.ln_final = nn.LayerNorm(n_embed)

        ## self.sa_head = SelfAttentionHead(vocab_size, n_embed, block_size)
        # self.sa_heads = MultiHeadAttention(num_heads=4, head_size=n_embed//4, n_embed=n_embed, block_size=block_size)
        # self.ffn = FeedForwardNet(n_embed, dropout=0.2)

        self.lm_head = nn.Linear(n_embed, vocab_size)


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

        # idx and targets both are tensors of shape (B, T) -> B = batch_sz, T = seq_len ("time steps", here 8)
        B, T = idx.shape
        tok_embed = self.token_embedding_table(idx) # (B, T, C)  C = "channels", here vocab_size or embedding dim for each token
        pos_embed = self.position_embedding_table(torch.arange(T, device=idx.device)) # (T, C)  C = "channels", here vocab_size or embedding dim for each token
        x_in = tok_embed + pos_embed
        # x_in = self.sa_heads(x_in)
        # x_in = self.ffn(x_in)
        x_in = self.ln_final(self.decoder_blocks(x_in))
        logits = self.lm_head(x_in) # (B, T, C)  C = "channels", here vocab_size or embedding dim for each token

        if targets is None:
            loss = None
        else:
            B, T, C = logits.shape
            # Cross entropy requires the 2nd param to be C "channels"
            loss = F.cross_entropy(logits.view(B*T, C), targets.view(B*T), ignore_index=0)

        return logits, loss
    
    def generate(self, idx, max_new_tokens):
        # idx is (B, T) shaped array of indices in current context
        for _ in range(max_new_tokens):
            #limit input idx to last "block size" tokens
            idx_cond = idx[:, -BLOCK_SIZE:]
            logits, loss = self(idx_cond)
            #focus only on the last time step
            logits = logits[:, -1, :] # becomes (B, C)
            # apply softmax for probs
            probs = F.softmax(logits, dim=-1) # (B, C)
            #sample from distribudion
            idx_next = torch.multinomial(probs, num_samples=1) # (B, 1)
            #append sampled index to running sequence idx
            idx = torch.cat([idx, idx_next], dim=1) # (B, T+1)

        return idx
    
    def get_num_params(self, non_embedding=True):
        """
        Return the number of parameters in the model.
        For non-embedding count (default), the position embeddings get subtracted.
        The token embeddings would too, except due to the parameter sharing these
        params are actually used as weights in the final layer, so we include them.
        """
        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)



if __name__ == "__main__":
    from data_utils import *
    xb, yb = get_random_batch('train')
    xb = xb.to(device)
    yb = yb.to(device)

    m = BigramLanguageModel(vocab_size=65, n_embed=n_embed, block_size=BLOCK_SIZE, num_heads=n_head, n_layers=n_layer).to(device)
    logits, loss = m(xb, yb)
    print(logits.shape)
    print(loss)