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

# 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)

# 2. Generate Ground Truth Data (A hidden 4th-degree curve with some noise)
np.random.seed(42)
x_data = np.linspace(-2, 2, 20)
# Target coefficients: [0.5, -1.2, 0.2, 2.0, -0.5]
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)

# Initial guess for coefficients (c4, c3, c2, c1, c0)
coeffs = np.array([0.0, 0.0, 0.0, 0.0, 0.0])
learning_rate = 0.05

print("Starting iteration loop with SmolLM2...\n")

# ... [Keep your previous imports and setup the same] ...
    
for iteration in range(1, 6):
    y_pred = np.polyval(coeffs, x_data)
    mse = np.mean((y_true - y_pred) ** 2)
    
    print(f"--- Iteration {iteration} ---")
    print(f"Current Coeffs [c4, c3, c2, c1, c0]: {np.round(coeffs, 3)}")
    print(f"Mean Squared Error: {mse:.4f}")
    
    # 1. We switch to an explicit few-shot format that small models can copy.
    prompt = (
        "<|im_start|>system\n"
        "You are an optimization assistant. You must ONLY output a single JSON block. "
        "No explanation. No introduction. Follow the examples exactly.\n"
        "Example 1:\n"
        "Input: Coeffs are [0.0, 0.0, 0.0, 0.0, 0.0]. MSE error is 15.80.\n"
        "Output: {\"index\": 0, \"direction\": 1}\n"
        "Example 2:\n"
        "Input: Coeffs are [0.5, -0.2, 0.0, 1.0, -0.1]. MSE error is 4.23.\n"
        "Output: {\"index\": 1, \"direction\": -1}\n"
        "<|im_end|>\n"
        f"<|im_start|>user\n"
        f"Input: Coeffs are {list(np.round(coeffs,3))}. MSE error is {mse:.2f}.\n"
        "Output:<|im_end|>\n"
        "<|im_start|>assistant\n"
    )
    
    inputs = tokenizer(prompt, return_tensors="pt").to(device)
    
    with torch.no_grad():
        # Lowering temperature drastically forces it to pick the most confident tokens
        outputs = model.generate(
            **inputs, 
            max_new_tokens=25, 
            temperature=0.01, 
            do_sample=False
        )
    
    response = tokenizer.decode(outputs[0][inputs.input_ids.shape[1]:], skip_special_tokens=True).strip()
    print(f"SmolLM2 Decision: {response}")
    
    # Parse decisions safely
    try:
        idx_match = re.search(r'"index":\s*(\d)', response)
        dir_match = re.search(r'"direction":\s*(-?\d)', response)
        
        if idx_match and dir_match:
            target_idx = int(idx_match.group(1))
            direction = int(dir_match.group(1))
            
            if 0 <= target_idx <= 4 and direction in [-1, 1]:
                coeffs[target_idx] += direction * learning_rate
                print(f"Applied update: Shifted coefficient {target_idx} by {direction * learning_rate}")
            else:
                print("Parsed values out of bounds. Skipping update.")
        else:
            print("Could not parse JSON. Skipping update.")
    except Exception as e:
        print(f"Parsing error: {e}")
    print("\n")