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

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

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

# --- THE SECRET MAP (Completely hidden from the SLM) ---
anonymous_menu = {
    "A": {"label": "Strategy ALPHA", "func": "relu"},
    "B": {"label": "Strategy BETA", "func": "leaky_relu"},
    "C": {"label": "Strategy GAMMA", "func": "sigmoid"},
    "D": {"label": "Strategy DELTA", "func": "tanh"}
}

token_map = {letter: tokenizer.convert_tokens_to_ids(letter) for letter in anonymous_menu.keys()}
token_ids = list(token_map.values())
id_to_letter = {v: k for k, v in token_map.items()}

# Generate mock training data
np.random.seed(42)
from scipy.io.wavfile import read

# Load training and test data
X_train = read('../X_train.wav')[1].reshape(-1, 784)
y_train = (read('../y_train.wav')[1] * 9).astype(int)
X_test = read('../X_test.wav')[1].reshape(-1, 784)
y_test = (read('../y_test.wav')[1] * 9).astype(int)
X20 = X_test[:1000]
yt20 = y_test[:1000]

# --- Keeping the DynamicMLP logic the exact same ---
class DynamicMLP:
    def __init__(self, input_size, hidden_size, output_size, learning_rate=0.1):
        self.W1 = np.random.randn(input_size, hidden_size) * 0.01
        self.b1 = np.zeros((1, hidden_size))
        self.W2 = np.random.randn(hidden_size, output_size) * 0.01
        self.b2 = np.zeros((1, output_size))
        self.learning_rate = learning_rate
        self.current_activation = "tanh"
        
    def activate(self, x):
        if self.current_activation == "relu": return np.maximum(0, x)
        elif self.current_activation == "leaky_relu": return np.where(x > 0, x, x * 0.01)
        elif self.current_activation == "sigmoid": return 1 / (1 + np.exp(-np.clip(x, -50, 50)))
        elif self.current_activation == "tanh": return np.tanh(x)

    def activate_derivative(self, x, act_output):
        if self.current_activation == "relu": return np.where(x > 0, 1, 0)
        elif self.current_activation == "leaky_relu": return np.where(x > 0, 1, 0.01)
        elif self.current_activation == "sigmoid": return act_output * (1 - act_output)
        elif self.current_activation == "tanh": return 1 - act_output ** 2

    def softmax(self, x):
        exp_x = np.exp(x - np.max(x, axis=1, keepdims=True))
        return exp_x / np.sum(exp_x, axis=1, keepdims=True)

    def forward(self, X):
        self.z1 = np.dot(X, self.W1) + self.b1
        self.a1 = self.activate(self.z1)
        self.z2 = np.dot(self.a1, self.W2) + self.b2
        return self.softmax(self.z2)

    def update(self, X, y_true):
        m = y_true.shape[0]
        y_pred = self.forward(X)
        dz2 = y_pred - y_true
        dW2 = np.dot(self.a1.T, dz2) / m
        db2 = np.sum(dz2, axis=0, keepdims=True) / m
        da1 = np.dot(dz2, self.W2.T)
        dz1 = da1 * self.activate_derivative(self.z1, self.a1)
        dW1 = np.dot(X.T, dz1) / m
        db1 = np.sum(dz1, axis=0, keepdims=True) / m
        self.W1 -= self.learning_rate * dW1
        self.b1 -= self.learning_rate * db1
        self.W2 -= self.learning_rate * dW2
        self.b2 -= self.learning_rate * db2

    def score(self, X, y_true):
        preds = np.argmax(self.forward(X), axis=1)
        return np.mean(preds == y_true)


if __name__ == "__main__":
    mlp = DynamicMLP(input_size=784, hidden_size=100, output_size=10, learning_rate=0.1)
    performance_log = {key: [] for key in anonymous_menu.keys()}
    N = 10 

    for _ in range(100):
        idx = np.random.randint(0, 60000, 100)
        X = X_train[idx]
        yt = y_train[idx]
        print(_, mlp.score(X,yt))
        mlp.update(X, np.eye(10)[yt])
    
    
    print(f"Starting ANONYMOUS data-window optimization (N = {N})...")
    
    for epoch in range(1, 4):
        print(f"\n================ META CYCLE {epoch} ================")
        
        # --- PHASE 1: ANONYMOUS EXPLORATION ---
        print("Evaluating hidden mathematical performance...")
        for letter, config in anonymous_menu.items():
            mlp.current_activation = config["func"]
            start_acc = mlp.score(X_train[:500], y_train[:500])
            
            for _ in range(N):
                idx = np.random.randint(0, 5000, 64)
                mlp.update(X_train[idx], np.eye(10)[y_train[idx]])
                
            end_acc = mlp.score(X_train[:500], y_train[:500])
            delta_acc = end_acc - start_acc
            performance_log[letter].append(delta_acc)
            print(f" -> Tested {config['label']}: Accuracy Metric Performance = {delta_acc:+.4f}")
            
        # --- PHASE 2: ANONYMOUS EXPLOITATION ---
        # The prompt contains absolutely ZERO clues about what the functions actually are
        prompt = (
            f"<|im_start|>system\n"
            f"You are a pure numerical analysis engine tracking statistical trajectories.\n"
            f"Analyze the performance changes from the last matrix processing window:\n"
            f"- Strategy ALPHA performance score: {performance_log['A'][-1]:+.4f}\n"
            f"- Strategy BETA performance score: {performance_log['B'][-1]:+.4f}\n"
            f"- Strategy GAMMA performance score: {performance_log['C'][-1]:+.4f}\n"
            f"- Strategy DELTA performance score: {performance_log['D'][-1]:+.4f}\n\n"
            f"Select the single character token option linked to the highest positive numerical score value:\n"
            f"A = Strategy ALPHA\n"
            f"B = Strategy BETA\n"
            f"C = Strategy GAMMA\n"
            f"D = Strategy DELTA\n"
            f"Which option letter is optimal? Output letter:<|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, :]
            selection_logits = logits[token_ids]
            probs = torch.softmax(selection_logits, dim=-1).float().cpu().numpy()
            
            chosen_token = token_ids[np.argmax(probs)]
            chosen_letter = id_to_letter[chosen_token]
            
        # Behind the scenes, map their anonymous token choice back to real math
        selected_strategy = anonymous_menu[chosen_letter]
        mlp.current_activation = selected_strategy["func"]
        
        print(f"\n[Meta-Decision Logits]: A(ALPHA): {probs[0]:.2f} | B(BETA): {probs[1]:.2f} | C(GAMMA): {probs[2]:.2f} | D(DELTA): {probs[3]:.2f}")
        print(f"✨ SmolLM2 selected code {chosen_letter} ({selected_strategy['label']}).")
        print(f"   [System Execution mapping]: Operating on mathematical backend '{selected_strategy['func'].upper()}'")
