import torch
import torch.nn.functional as F
from model.transformer import Transformer
from tokenizer.tokenizer import BPETokenizer

class InferenceEngine:
    def __init__(self, model: Transformer, tokenizer: BPETokenizer, device: str = "cpu"):
        self.model = model
        self.tokenizer = tokenizer
        self.device = device
        self.model.to(self.device)
        self.model.eval()

    def _get_kv_caches(self, batch_size: int, max_seq_len: int):
        params = self.model.params
        head_dim = params.dim // params.n_heads
        n_kv_heads = params.n_kv_heads if params.n_kv_heads else params.n_heads
        
        kv_caches = []
        for _ in range(params.n_layers):
            k_cache = torch.zeros(batch_size, max_seq_len, n_kv_heads, head_dim).to(self.device)
            v_cache = torch.zeros(batch_size, max_seq_len, n_kv_heads, head_dim).to(self.device)
            kv_caches.append((k_cache, v_cache))
        return kv_caches

    @torch.no_grad()
    def generate(
        self, 
        prompt: str, 
        max_new_tokens: int = 50, 
        temperature: float = 0.7, 
        top_p: float = 0.9,
        top_k: int = 50,
        stream: bool = False
    ):
        if not stream:
            return self._generate_no_stream(prompt, max_new_tokens, temperature, top_p, top_k)
        return self._generate_stream(prompt, max_new_tokens, temperature, top_p, top_k)

    def _generate_no_stream(self, prompt, max_new_tokens, temperature, top_p, top_k):
        tokens = self.tokenizer.encode(prompt, bos=True, eos=False)
        x = torch.tensor(tokens).unsqueeze(0).to(self.device)
        
        bsz = x.shape[0]
        kv_caches = self._get_kv_caches(bsz, self.model.params.max_seq_len)
        
        # Initial forward to fill cache
        logits = self.model(x, start_pos=0, kv_caches=kv_caches)
        
        generated = tokens
        start_pos = x.shape[1]
        
        for _ in range(max_new_tokens):
            logits = self.model(x[:, -1:], start_pos=start_pos, kv_caches=kv_caches)
            logits = logits[:, -1, :] / max(temperature, 1e-5)
            
            # Top-k filtering
            if top_k > 0:
                v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
                logits[logits < v[:, [-1]]] = float('-inf')
                
            # Top-p (nucleus) filtering
            if top_p < 1.0:
                sorted_logits, sorted_indices = torch.sort(logits, descending=True)
                cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1)
                sorted_indices_to_remove = cumulative_probs > top_p
                sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()
                sorted_indices_to_remove[..., 0] = False
                indices_to_remove = sorted_indices[sorted_indices_to_remove]
                logits[:, indices_to_remove] = float('-inf')

            probs = F.softmax(logits, dim=-1)
            next_token = torch.multinomial(probs, num_samples=1)
            
            token_id = next_token.item()
            generated.append(token_id)
            x = next_token
            start_pos += 1
            
            if token_id == self.tokenizer.special_tokens["<eos>"]:
                break
                
        return self.tokenizer.decode(generated)

    def _generate_stream(self, prompt, max_new_tokens, temperature, top_p, top_k):
        tokens = self.tokenizer.encode(prompt, bos=True, eos=False)
        x = torch.tensor(tokens).unsqueeze(0).to(self.device)
        bsz = x.shape[0]
        kv_caches = self._get_kv_caches(bsz, self.model.params.max_seq_len)
        self.model(x, start_pos=0, kv_caches=kv_caches)
        start_pos = x.shape[1]
        
        for _ in range(max_new_tokens):
            logits = self.model(x[:, -1:], start_pos=start_pos, kv_caches=kv_caches)
            logits = logits[:, -1, :] / max(temperature, 1e-5)
            
            if top_k > 0:
                v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
                logits[logits < v[:, [-1]]] = float('-inf')
            if top_p < 1.0:
                sorted_logits, sorted_indices = torch.sort(logits, descending=True)
                cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1)
                sorted_indices_to_remove = cumulative_probs > top_p
                sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()
                sorted_indices_to_remove[..., 0] = False
                indices_to_remove = sorted_indices[sorted_indices_to_remove]
                logits[:, indices_to_remove] = float('-inf')

            probs = F.softmax(logits, dim=-1)
            next_token = torch.multinomial(probs, num_samples=1)
            token_id = next_token.item()
            yield self.tokenizer.decode([token_id])
            x = next_token
            start_pos += 1
            if token_id == self.tokenizer.special_tokens["<eos>"]:
                break

if __name__ == "__main__":
    # Test KV-cache inference
    from model.transformer import ModelArgs
    params = ModelArgs(dim=256, n_layers=2, n_heads=4, vocab_size=100284)
    model = Transformer(params)
    tokenizer = BPETokenizer()
    engine = InferenceEngine(model, tokenizer)
    
    print("Testing KV-cache generation...")
    res = engine.generate("Once upon a time", max_new_tokens=10, stream=False)
    print(f"Result: {res}")