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 our internal numbers (1-4) back to activation functions
strategy_map = {
    1: "relu",
    2: "leaky_relu",
    3: "sigmoid",
    4: "tanh"
}

# Explicitly find the unique Token IDs for characters '1', '2', '3', '4' in SmolLM2
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)}

# Generate mock data
from scipy.io.wavfile import read

np.random.seed(42)
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]

# --- Simplified MLP Class ---
class DynamicMLP:
    def __init__(self, input_size=784, hidden_size=100, output_size=10, learning_rate=0.1):
        self.W1 = np.random.randn(input_size, hidden_size) * np.sqrt(2.0 / input_size)
        self.b1 = np.zeros((1, hidden_size))
        self.W2 = np.random.randn(hidden_size, output_size) * np.sqrt(2.0 / hidden_size)
        self.b2 = np.zeros((1, output_size))
        self.learning_rate = learning_rate
        self.current_activation = "relu"
        
    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 forward(self, X):
        return np.exp(np.dot(self.activate(np.dot(X, self.W1) + self.b1), self.W2) + self.b2)

    def update(self, X, y_true):
        # Perform basic forward/backward adjustments...
        pass

    def score(self, X, y_true):
        return 0.85 # Mock output metric tracker


if __name__ == "__main__":
    mlp = DynamicMLP()
    
    # Initialize your sequence history tracker (e.g., seed it with previous run observations)
    sequence_history = [3, 3, 3, 3, 3, 3]
    
    print(f"Starting Sequence Progression Optimization...")
    print(f"Initial Sequence History: {sequence_history}\n")
    
    for epoch in range(1, 6):
        # 1. Transform sequence history list directly into a compact string token pattern
        sequence_str = ",".join(map(str, sequence_history))
        
        # Build a prompt that acts purely as a sequence continuation challenge
        prompt = (
            f"<|im_start|>system\n"
            f"Predict the next single digit pattern completion integer in the sequence.\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, :]
            
            # Slice the probabilities for exactly tokens '1', '2', '3', '4'
            target_logits = logits[choice_tokens]
            probs = torch.softmax(target_logits, dim=-1).float().cpu().numpy()
            
            # Extract chosen token integer digit
            chosen_token_id = choice_tokens[np.argmax(probs)]
            predicted_digit = token_to_digit[chosen_token_id]
            
        print(f"--- Meta Step {epoch} ---")
        print(f"Prompt Sequence Given: [{sequence_str}, ?]")
        print(f"SLM Continuation Probs -> 1: {probs[0]:.2f} | 2: {probs[1]:.2f} | 3: {probs[2]:.2f} | 4: {probs[3]:.2f}")
        print(f"✨ Intuition Choice -> Appending {predicted_digit} to sequence history.")
        
        # 2. Dynamically assign the activation based on the pattern prediction sequence
        mlp.current_activation = strategy_map[predicted_digit]
        print(f"   [Execution Step]: Running backend via '{mlp.current_activation.upper()}' configuration.")
        
        # 3. Append the choice to the history list for the next iteration loop sequence
        sequence_history.append(predicted_digit)
        
        # Run training loop step
        idx = np.random.randint(0, 2000, 64)
        mlp.update(X_train[idx], np.eye(10)[y_train[idx]])
        print("\n")