import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
import numpy as np

model_id = "HuggingFaceTB/SmolLM2-135M-Instruct"
device = "cuda" if torch.cuda.is_available() else "cpu"

tokenizer = AutoTokenizer.from_pretrained(model_id)
llm_model = AutoModelForCausalLM.from_pretrained(model_id).to(device)

choice_strings = ["1", "2", "3", "4"]
choice_tokens = [tokenizer.convert_tokens_to_ids(c) for c in choice_strings]
token_to_digit = {v: int(k) for k, v in zip(choice_strings, choice_tokens)}

# Start with your repeating test array
sequence_history = [3, 3, 3, 3, 3]

print("Starting Anti-Repetition Sequence Loop...\n")

for epoch in range(1, 10):
    sequence_str = ", ".join(map(str, sequence_history))
    
    prompt = (
        f"<|im_start|>system\n"
        f"Predict the next token integer to balance the sequence exploration path.\n"
        f"Sequence: [{sequence_str}, \n"
        f"Next integer token is:<|im_end|>\n"
        f"<|im_start|>assistant\n"
    )
    
    inputs = tokenizer(prompt, return_tensors="pt").to(device)
    
    with torch.no_grad():
        outputs = llm_model(**inputs)
        logits = outputs.logits[0, -1, choice_tokens]
        
        # 🔥 THE ANTI-TRAP HACK:
        # Check if the last 3 moves were identical. If they were, penalize that specific token logit!
        if len(sequence_history) >= 3 and len(set(sequence_history[-3:])) == 1:
            stuck_digit = sequence_history[-1]
            stuck_idx = choice_strings.index(str(stuck_digit))
            logits[stuck_idx] -= 8.0  # Apply a heavy penalty to force a state switch
            
        probs = torch.softmax(logits, dim=-1).float().cpu().numpy()
        chosen_token_id = choice_tokens[np.argmax(probs)]
        predicted_digit = token_to_digit[chosen_token_id]

    print(f"Sequence Given: [{sequence_str}, ?]")
    print(f"Calibrated Probs -> 1: {probs[0]:.2f} | 2: {probs[1]:.2f} | 3: {probs[2]:.2f} | 4: {probs[3]:.2f}")
    print(f"✨ Model Output: {predicted_digit}\n")
    
    sequence_history.append(predicted_digit)