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

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

print(f"Loading {model_id} on {device}...")
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(model_id).to(device)

# Map our target text choices to their exact token IDs in SmolLM2's vocabulary
index_tokens = [tokenizer.convert_tokens_to_ids(str(i)) for i in range(5)] # Tokens for '0', '1', '2', '3', '4'
plus_token = tokenizer.convert_tokens_to_ids("+")
minus_token = tokenizer.convert_tokens_to_ids("-")

# 2. Data generation
np.random.seed(42)
x_data = np.linspace(-2, 2, 20)
y_true = 0.5*x_data**4 - 1.2*x_data**3 + 0.2*x_data**2 + 2.0*x_data - 0.5 + np.random.normal(0, 0.1, 20)

coeffs = np.array([0.0, 0.0, 0.0, 0.0, 0.0])
learning_rate = 0.1

# Initialize tracking variables before the loop
prev_mse = None
last_index = None
last_direction = None

print("Starting True Hybrid Probabilistic Optimizer...\n")

for iteration in range(1, 15): # Expanded to 14 steps to see the auto-correction
    y_pred = np.polyval(coeffs, x_data)
    mse = np.mean((y_true - y_pred) ** 2)
    
    print(f"--- Iteration {iteration} ---")
    print(f"Current Coeffs: {np.round(coeffs, 3)}")
    print(f"MSE: {mse:.4f}")
    
    # --- PROBABILISTIC OVERRIDE HACK ---
    # If the error increased, we know the LAST move was wrong. 
    # Forcefully reverse the last action to undo the damage, and block that index for this turn.
    if prev_mse is not None and mse > prev_mse:
        print(f"⚠️ [System Override]: MSE increased! Forcefully undoing last move and reversing direction.")
        # Undo the bad move
        coeffs[last_index] -= last_direction * learning_rate
        # Calculate corrected MSE
        y_pred = np.polyval(coeffs, x_data)
        mse = np.mean((y_true - y_pred) ** 2)
        
        feedback_str = "Your last move was terrible and was forcefully reverted. Pick a different index!"
    else:
        feedback_str = "Your last move successfully lowered the error. Keep going."
        
    prev_mse = mse
    
    prompt = (
        f"<|im_start|>system\n"
        f"Optimize the 4th-degree polynomial. Current error: {mse:.2f}.\n"
        f"Feedback: {feedback_str}\n"
        f"Coefficients: {list(np.round(coeffs,2))}.\n"
        f"Output the best coefficient index to change (0-4):<|im_end|>\n"
        f"<|im_start|>assistant\n"
    )
    
    inputs = tokenizer(prompt, return_tensors="pt").to(device)
    
    with torch.no_grad():
        outputs = model(**inputs)
        next_token_logits = outputs.logits[0, -1, :] 
        index_logits = next_token_logits[index_tokens]
        
        # If we just had a failure, drastically penalize the index that caused it
        if feedback_str.startswith("Your last move was terrible"):
            index_logits[last_index] -= 10.0 # Artificially suppress the stuck token
            
        index_probs = torch.softmax(index_logits, dim=-1).float().cpu().numpy()
        chosen_index = np.argmax(index_probs)
        
    direction_prompt = prompt + f"{chosen_index}\nShould we change it with + or -?<|im_end|>\n<|im_start|>assistant\n"
    dir_inputs = tokenizer(direction_prompt, return_tensors="pt").to(device)
    
    with torch.no_grad():
        dir_outputs = model(**dir_inputs)
        dir_logits = dir_outputs.logits[0, -1, :]
        sign_logits = torch.tensor([dir_logits[plus_token], dir_logits[minus_token]])
        sign_probs = torch.softmax(sign_logits, dim=-1).float().cpu().numpy()
        chosen_direction = 1 if np.argmax(sign_probs) == 0 else -1

    # Save actions for the next iteration's evaluation
    last_index = chosen_index
    last_direction = chosen_direction

    # Apply the step
    coeffs[chosen_index] += chosen_direction * learning_rate
    print(f"AI Probabilities (0-4): {np.round(index_probs, 2)}")
    print(f"Applied Action -> Index {chosen_index} modified by {chosen_direction * learning_rate}")
    print("\n")