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)

# Map action letters to specific activation functions
activation_actions = {
    "A": "relu",
    "B": "leaky_relu",
    "C": "sigmoid",
    "D": "tanh"
}
token_map = {letter: tokenizer.convert_tokens_to_ids(letter) for letter in activation_actions.keys()}
token_ids = list(token_map.values())
id_to_letter = {v: k for k, v in token_map.items()}

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]

# 2. Flexible MLP Classifier with Swappable Activation Logic
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 = "relu" # Default state
        
    # --- Activation Functions & Derivatives ---
    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)
        
        # Backward path utilizing current dynamic derivative rules
        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
        
        # Apply step
        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)


# 3. Execution Loop
if __name__ == "__main__":
    mlp = DynamicMLP(input_size=784, hidden_size=100, output_size=10, learning_rate=0.2)
    
    print("Starting Activation Function Swapping Loop...\n")
    
    for step in range(1, 11):
        idx = np.random.randint(0, 2000, 128)
        X = X_train[idx]
        yt = y_train[idx]
        
        # Record accuracy prior to this step's update
        accuracy = mlp.score(X, yt)
        
        print(f"--- Step {step} ---")
        print(f"Current Activation: {mlp.current_activation.upper()}")
        print(f"Batch Accuracy: {accuracy:.4f}")
        
        # Build prompt telling SmolLM2 how things are currently performing
        prompt = (
            f"<|im_start|>system\n"
            f"Select the best hidden layer activation function token to improve training accuracy.\n"
            f"Current batch accuracy: {accuracy:.4f}. Current activation: {mlp.current_activation}\n\n"
            f"Menu options:\n"
            f"A = Switch to ReLU\n"
            f"B = Switch to LeakyReLU\n"
            f"C = Switch to Sigmoid\n"
            f"D = Switch to Tanh\n"
            f"Which single letter token is best? 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, :]
            
            # Extract raw neural probabilities for choices A, B, C, D
            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]
            
        # Swap the activation function state
        mlp.current_activation = activation_actions[chosen_letter]
        
        print(f"SLM Intuition Probs -> A(ReLU): {probs[0]:.2f} | B(Leaky): {probs[1]:.2f} | C(Sigmoid): {probs[2]:.2f} | D(Tanh): {probs[3]:.2f}")
        print(f"Result -> Meta-optimizer selected: {mlp.current_activation.upper()}")
        
        # Perform training step with the selected function
        mlp.update(X, np.eye(10)[yt])
        print("\n")