#!/usr/bin/env python3
"""
Speculative 'Look-Ahead' Generation for HuggingFaceTB/SmolLM2-135M-Instruct.
Prioritizes sequence intelligence over raw speed by branching 4 tokens into 
the future and picking the path with the highest overall log-probability.
"""

import sys
import torch
import torch.nn.functional as F
from transformers import AutoModelForCausalLM, AutoTokenizer

MODEL_NAME = "HuggingFaceTB/SmolLM2-135M-Instruct"

def load_model_and_tokenizer():
    print(f"🔄 Loading {MODEL_NAME}...")
    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token
        
    device = "cuda" if torch.cuda.is_available() else ("mps" if torch.backends.mps.is_available() else "cpu")
    print(f"💻 Using device: {device}")
    
    model = AutoModelForCausalLM.from_pretrained(
        MODEL_NAME,
        torch_dtype=torch.float16 if device in ["cuda", "mps"] else torch.float32,
        low_cpu_mem_usage=True,
    ).to(device)
    
    return model, tokenizer, device

def lookahead_generate(model, tokenizer, device, prompt, max_new_tokens=600, num_branches=1, lookahead_depth=1):
    """
    Generates text by looking ahead 4 tokens across 4 different branches,
    evaluating which branch makes the most 'sense' globally, and committing to it.
    """
    # Apply Chat Template
    messages = [{"role": "user", "content": prompt}]
    input_text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
    input_ids = tokenizer(input_text, return_tensors="pt")["input_ids"].to(device)
    
    print("🧠 Thinking with look-ahead branching... \n")
    
    # Generate until max tokens or EOS
    for _ in range(max_new_tokens):
        best_branch_first_token = None
        best_branch_score = float('-inf')
        
        # 1. Roll out alternative futures (Branches)
        for _ in range(num_branches):
            branch_ids = input_ids.clone()
            branch_log_prob = 0.0
            first_token_of_branch = None
            
            # Lookahead Depth (e.g., 4 tokens deep into this specific future)
            for depth in range(lookahead_depth):
                with torch.no_grad():
                    outputs = model(branch_ids)
                    next_token_logits = outputs.logits[0, -1, :]
                    
                    # Apply temperature to keep branches diverse
                    probs = F.softmax(next_token_logits / 0.7, dim=-1)
                    
                    # Sample a token from the distribution
                    next_token = torch.multinomial(probs, num_samples=1)
                    token_log_prob = torch.log(probs[next_token]).item()
                    
                    if depth == 0:
                        first_token_of_branch = next_token
                    
                    # Track the cumulative health/intelligence of this branch
                    branch_log_prob += token_log_prob
                    branch_ids = torch.cat([branch_ids, next_token.unsqueeze(0)], dim=-1)
                    
                    if next_token.item() == tokenizer.eos_token_id:
                        break
            
            # Normalized score by depth to avoid penalizing early EOS
            branch_score = branch_log_prob / (depth + 1)
            
            # Keep the branch that holds the highest total confidence over the 4-step window
            if branch_score > best_branch_score:
                best_branch_score = branch_score
                best_branch_first_token = first_token_of_branch
        
        # 2. Commit the first token of the single smartest branch found
        input_ids = torch.cat([input_ids, best_branch_first_token.unsqueeze(0)], dim=-1)
        
        # Live-print the committed token
        print(tokenizer.decode(best_branch_first_token, skip_special_tokens=True), end="", flush=True)
        
        if best_branch_first_token.item() == tokenizer.eos_token_id:
            break
            
    print("\n")

def main():
    model, tokenizer, device = load_model_and_tokenizer()
    
    # Example logic puzzle prompt that traditional small models struggle with due to greediness
    default_prompt = "If Sally is taller than Tom, and Tom is taller than Sam, who is the shortest? Explain step by step."
    
    prompt = " ".join(sys.argv[1:]) if len(sys.argv) > 1 else default_prompt
    print(f"📝 Prompt: {prompt}")
    
    lookahead_generate(model, tokenizer, device, prompt)

if __name__ == "__main__":
    main()