import numpy as np
from .model import PotatoGPT
from .tokenizer import CharTokenizer

def generate(model: PotatoGPT, idx: np.ndarray, max_new_tokens: int, temperature: float = 1.0, top_k: int = None):
    """
    Generates tokens based on the prompt.
    """
    for _ in range(max_new_tokens):
        # Crop context if it's too long
        idx_cond = idx[:, -model.block_size:]
        
        # Forward pass
        logits = model(idx_cond)
        
        # Focus only on the last time step
        logits = logits[:, -1, :] / temperature
        
        if top_k is not None:
            v = np.sort(logits, axis=-1)[:, -top_k]
            logits[logits < v[:, None]] = -np.inf
        
        # Softmax to get probabilities
        exp_logits = np.exp(logits - np.max(logits, axis=-1, keepdims=True))
        probs = exp_logits / np.sum(exp_logits, axis=-1, keepdims=True)
        
        # Sample or greedy
        if temperature == 0:
            idx_next = np.argmax(probs, axis=-1, keepdims=True)
        else:
            # Simple multinomial sampling
            idx_next = np.array([[np.random.choice(len(p), p=p) for p in probs]])
            
        idx = np.concatenate((idx, idx_next), axis=1)
        
    return idx
