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

class LMBase(nn.Module):
    def __init__(self):
        super().__init__()
        self.name=self.full_name="LMBase"
    
    def calculate_loss(self, logits, target_tokens, l1_loss_lambda=None):
        loss = F.cross_entropy(
            logits.reshape(-1, logits.size(-1)),
            target_tokens.reshape(-1),
            reduction='mean'
            )
        return loss
    
    def init_weights(self, module, num_layers=None):
        if isinstance(module, nn.Linear):
            std = 0.02 if num_layers is None else 0.02 / math.sqrt(2 * num_layers)
            nn.init.normal_(module.weight, mean=0.0, std=std)
            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)

        elif isinstance(module, nn.LayerNorm):
            nn.init.ones_(module.weight)
            nn.init.zeros_(module.bias)
    
    def count_parameters(self):
        total_params = sum(p.numel() for p in self.parameters())
        embed_params = sum(p.numel() for name, p in self.named_parameters() if "embed" in name.lower())
        non_embed_params = total_params - embed_params
        return total_params, embed_params, non_embed_params

    @torch.no_grad()
    def generate(
        self, 
        input_ids, 
        max_generation_length, 
        tokenizer, 
        temperature=1.0, 
        top_p=0.9, 
        return_generation_only=False
    ):

        self.eval()

        batch_size = input_ids.size(0)
        device = input_ids.device

        generated = input_ids.clone()
        finished = torch.zeros(batch_size, dtype=torch.bool, device=device)

        for _ in range(max_generation_length):

            logits = self(generated)[:, -1, :] / temperature
            probs = F.softmax(logits, dim=-1)

            sorted_probs, sorted_indices = torch.sort(probs, dim=-1, descending=True)
            cumulative_probs = torch.cumsum(sorted_probs, dim=-1)

            cutoff_mask = cumulative_probs > top_p
            cutoff_mask[:, 1:] = cutoff_mask[:, :-1].clone()
            cutoff_mask[:, 0] = False

            sorted_probs = sorted_probs.masked_fill(cutoff_mask, 0.0)
            normalized_probs = sorted_probs / sorted_probs.sum(dim=-1, keepdim=True)

            probs = torch.zeros_like(normalized_probs).scatter(-1, sorted_indices, normalized_probs)

            next_token = torch.multinomial(probs, num_samples=1).squeeze(-1)
            next_token = torch.where(finished, torch.full_like(next_token, tokenizer.pad_token_id), next_token)

            generated = torch.cat([generated, next_token.unsqueeze(1)], dim=1)

            finished |= next_token == tokenizer.eos_token_id

            if finished.all():
                break

        if return_generation_only:
            return generated[:, input_ids.size(1):]
        else:
            return generated