import random
import math
from collections import defaultdict, Counter
from typing import List, Dict, Tuple, Optional
class SimpleLLM:
    """A simple character-level language model using n-grams"""
    
    def __init__(self, n: int = 3):
        """
        Initialize the language model with n-gram size
        
        Args:
            n: Size of n-grams (context window)
        """
        self.n = n
        self.ngrams: Dict[str, Dict[str, int]] = defaultdict(Counter)
        self.start_tokens = []
        self.vocab = set()
        
    def train(self, text: str) -> None:
        """
        Train the model on a text corpus
        
        Args:
            text: Training text
        """
        if len(text) < self.n:
            raise ValueError(f"Text too short for n={self.n}")
        
        # Add special tokens for start and end
        padded_text = "▁" * (self.n - 1) + text + "□"
        
        # Build n-grams
        for i in range(len(padded_text) - self.n):
            context = padded_text[i:i + self.n - 1]
            next_char = padded_text[i + self.n - 1]
            
            self.ngrams[context][next_char] += 1
            self.vocab.add(next_char)
            
            # Track starting contexts
            if i < self.n - 1:
                self.start_tokens.append(context)
    
    def _get_next_char(self, context: str) -> str:
        """
        Get next character based on context using weighted random choice
        
        Args:
            context: Current context string
            
        Returns:
            Predicted next character
        """
        if context not in self.ngrams:
            return random.choice(list(self.vocab)) if self.vocab else " "
        
        # Get frequency distribution for this context
        char_counts = self.ngrams[context]
        total = sum(char_counts.values())
        
        # Weighted random choice
        rand_val = random.random() * total
        cumulative = 0
        
        for char, count in char_counts.items():
            cumulative += count
            if rand_val <= cumulative:
                return char
        
        # Fallback
        return random.choice(list(self.vocab)) if self.vocab else " "
    
    def generate(self, 
                 max_length: int = 100, 
                 temperature: float = 1.0,
                 seed_text: Optional[str] = None) -> str:
        """
        Generate text from the model
        
        Args:
            max_length: Maximum length of generated text
            temperature: Controls randomness (0.0 to 1.0)
            seed_text: Starting text (optional)
            
        Returns:
            Generated text
        """
        if not self.ngrams:
            return "Model not trained yet!"
        
        # Set initial context
        if seed_text:
            # Use seed text if provided
            context = seed_text[-(self.n - 1):]
            if len(context) < self.n - 1:
                context = "▁" * (self.n - 1 - len(context)) + context
            result = seed_text
        else:
            # Start with random context from training
            context = random.choice(self.start_tokens)
            result = context.replace("▁", "")
        
        # Generate characters
        for _ in range(max_length - len(result)):
            # Get prediction with temperature sampling
            if temperature < 1.0:
                # Greedy or tempered sampling
                char = self._get_next_char_with_temperature(context, temperature)
            else:
                # Regular weighted random
                char = self._get_next_char(context)
            
            # Stop if we hit end token
            if char == "□":
                break
                
            result += char
            context = (context + char)[-(self.n - 1):]
        
        return result
    
    def _get_next_char_with_temperature(self, 
                                       context: str, 
                                       temperature: float) -> str:
        """
        Get next character with temperature sampling
        
        Args:
            context: Current context
            temperature: Sampling temperature
            
        Returns:
            Selected character
        """
        if context not in self.ngrams:
            return random.choice(list(self.vocab)) if self.vocab else " "
        
        char_counts = self.ngrams[context]
        
        # Apply temperature to probabilities
        probs = {}
        total = 0
        
        for char, count in char_counts.items():
            # Apply temperature scaling
            scaled = count ** (1.0 / temperature)
            probs[char] = scaled
            total += scaled
        
        # Sample from temperature-adjusted distribution
        rand_val = random.random() * total
        cumulative = 0
        
        for char, prob in probs.items():
            cumulative += prob
            if rand_val <= cumulative:
                return char
        
        return random.choice(list(self.vocab))
    
    def predict_next(self, context: str, top_k: int = 5) -> List[Tuple[str, float]]:
        """
        Predict the most likely next characters
        
        Args:
            context: Context string
            top_k: Number of top predictions to return
            
        Returns:
            List of (character, probability) tuples
        """
        if context not in self.ngrams:
            return []
        
        char_counts = self.ngrams[context]
        total = sum(char_counts.values())
        
        # Calculate probabilities
        predictions = [(char, count / total) 
                      for char, count in char_counts.items()]
        
        # Sort by probability (descending)
        predictions.sort(key=lambda x: x[1], reverse=True)
        
        return predictions[:top_k]
    
    def perplexity(self, text: str) -> float:
        """
        Calculate perplexity of the model on given text
        
        Args:
            text: Text to evaluate
            
        Returns:
            Perplexity score (lower is better)
        """
        if not self.ngrams:
            return float('inf')
        
        padded_text = "▁" * (self.n - 1) + text + "□"
        log_prob_sum = 0
        total_chars = 0
        
        for i in range(len(padded_text) - self.n):
            context = padded_text[i:i + self.n - 1]
            next_char = padded_text[i + self.n - 1]
            
            if context in self.ngrams and next_char in self.ngrams[context]:
                total = sum(self.ngrams[context].values())
                prob = self.ngrams[context][next_char] / total
                log_prob_sum += math.log(prob) if prob > 0 else -float('inf')
            else:
                # Add penalty for unseen n-grams
                log_prob_sum += math.log(1e-10)
            
            total_chars += 1
        
        if total_chars == 0:
            return float('inf')
        
        avg_log_prob = log_prob_sum / total_chars
        return math.exp(-avg_log_prob)
# Example usage and demonstration
def main():
    # Sample training text
    training_text = """
    The quick brown fox jumps over the lazy dog.
    Now is the time for all good men to come to the aid of their party.
    To be or not to be, that is the question.
    All that glitters is not gold.
    A journey of a thousand miles begins with a single step.
    """
    
    # Create and train model
    print("Creating and training model...")
    model = SimpleLLM(n=4)
    model.train(training_text)
    
    # Generate some text
    print("\n--- Generated Text ---")
    for i in range(3):
        text = model.generate(max_length=50, temperature=0.7)
        print(f"Sample {i+1}: {text}")
    
    # Generate with seed text
    print("\n--- Generated with Seed ---")
    seed_text = "The quick"
    generated = model.generate(max_length=60, temperature=0.8, seed_text=seed_text)
    print(f"Seed: '{seed_text}'")
    print(f"Generated: {generated}")
    
    # Show predictions
    print("\n--- Next Character Predictions ---")
    context = "the "
    predictions = model.predict_next(context, top_k=5)
    print(f"After '{context}', most likely next characters:")
    for char, prob in predictions:
        print(f"  '{char}': {prob:.2%}")
    
    # Calculate perplexity
    test_text = "The quick brown"
    perplexity = model.perplexity(test_text)
    print(f"\n--- Model Evaluation ---")
    print(f"Perplexity on '{test_text}': {perplexity:.2f}")
    
    # Interactive generation
    print("\n--- Interactive Mode ---")
    print("Type a starting phrase (or 'quit' to exit):")
    
    while True:
        user_input = input("\nYour text: ").strip()
        if user_input.lower() in ['quit', 'exit', 'q']:
            break
        
        if len(user_input) < model.n - 1:
            print(f"Please enter at least {model.n - 1} characters")
            continue
        
        generated = model.generate(
            max_length=100,
            temperature=0.7,
            seed_text=user_input
        )
        print(f"Generated: {generated}")
if __name__ == "__main__":
    main()