import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
import numpy as np
from scipy.io.wavfile import read

# ========================================================
# 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} as the Meta-Optimizer on {device}...")
tokenizer = AutoTokenizer.from_pretrained(model_id)
# Load in float32 or float16 depending on availability to prevent numpy alignment bottlenecks
llm_model = AutoModelForCausalLM.from_pretrained(model_id).to(device)

# The Secret Mapping (Fully hidden from the SLM to completely bypass semantic biases)
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"}
}

# Map clean multiple-choice string integers to vocabulary IDs
choice_strings = ["1", "2", "3", "4"]
choice_tokens = [tokenizer.convert_tokens_to_ids(c) for c in choice_strings]
index_to_letter = {1: "A", 2: "B", 3: "C", 4: "D"}

# ========================================================
# 2. LOAD YOUR WAVE DATASET
# ========================================================
print("Loading audio wave dataset layers...")
try:
    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)
except FileNotFoundError:
    print("⚠️ Local dataset files not found. Creating random simulation matrices for fallback execution.")
    X_train = np.random.randn(5000, 784)
    y_train = np.random.randint(0, 10, 5000)

# ========================================================
# 3. DYNAMIC MLP ARCHITECTURE
# ========================================================
class DynamicMLP:
    def __init__(self, input_size, hidden_size, output_size, learning_rate=0.1):
        # Xavier initialization method to stabilize early activation changes fairly
        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 = "tanh"  # Base activation to initiate warm-up
        
    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)


# ========================================================
# 4. EXECUTION DRIVER
# ========================================================
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  # Window size per exploration segment
    
    # --- PHASE 0: WARM-UP SLOPE PHASE ---
    print("\n🚀 Initiating Classifier Model Warm-Up Phase...")
    print("Iterating all strategies evenly across base arrays to settle erratic weight scales...")
    
    for warm_step in range(100):
        idx = np.random.randint(0, X_train.shape[0], 100)
        X = X_train[idx]
        yt = y_train[idx]
        
        # Cycle activation function parameters dynamically during warm-up
        if warm_step % 25 == 0:
            act_pool = ["relu", "leaky_relu", "sigmoid", "tanh"]
            mlp.current_activation = act_pool[warm_step // 25]
            
        if warm_step % 20 == 0:
            print(f"  Warm-up Iteration {warm_step}/100 | Current Batch Score: {mlp.score(X, yt):.4f}")
            
        mlp.update(X, np.eye(10)[yt])
        
    print(f"Warm-Up Phase Completed. Current Network Stable Baseline Score: {mlp.score(X_train[:500], y_train[:500]):.4f}")
    print("\nStarting Meta-Optimizer Closed Loop Evaluation...")
    
    # --- MAIN META LEARNING CYCLES ---
    for epoch in range(1, 4):
        print(f"\n================ META CYCLE {epoch} ================")
        
        # --- PHASE 1: ANONYMOUS EXPLORATION (Data Collection) ---
        print("Gathering hidden performance metrics across strategy variations...")
        for letter, config in anonymous_menu.items():
            mlp.current_activation = config["func"]
            start_acc = mlp.score(X_train[:500], y_train[:500])
            
            # Allow strategy to progress over explicit test window interval N
            for _ in range(N):
                idx = np.random.randint(0, X_train.shape[0], 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" -> {config['label']} window run completed: Delta Performance = {delta_acc:+.4f}")
            
        # Pull current interval's metric deltas
        deltas = [
            performance_log['A'][-1],
            performance_log['B'][-1],
            performance_log['C'][-1],
            performance_log['D'][-1]
        ]
        
# --- PHASE 2: RANDOMIZED POSITION MATCHING ---
    import random
    
    # Raw strategy keys and their corresponding deltas for this loop
    raw_strategies = ["A", "B", "C", "D"]
    deltas = [
        performance_log['A'][-1],
        performance_log['B'][-1],
        performance_log['C'][-1],
        performance_log['D'][-1]
    ]
    
    # Zip them up and shuffle their positions completely randomly
    combined = list(zip(raw_strategies, deltas))
    random.shuffle(combined)
    
    # Create a dynamic dynamic menu based on the shuffle
    # mapping 1,2,3,4 to whatever landed there
    index_to_letter = {i+1: combined[i][0] for i in range(4)}
    shuffled_deltas = [combined[i][1] for i in range(4)]
    
    # A completely flattened comparison prompt that breaks vertical stack bias
    prompt = (
        f"<|im_start|>system\n"
        f"You are a basic math calculator comparing numbers.\n"
        f"Which option has the greatest value: Option 1 is {shuffled_deltas[0]:+.4f}, Option 2 is {shuffled_deltas[1]:+.4f}, Option 3 is {shuffled_deltas[2]:+.4f}, Option 4 is {shuffled_deltas[3]:+.4f}.\n"
        f"Output only the best option digit (1, 2, 3, or 4):<|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[choice_tokens]
        probs = torch.softmax(selection_logits, dim=-1).float().cpu().numpy()
        chosen_index = int(np.argmax(probs)) + 1 
        
    # Unpack what strategy actually occupied that shuffled number slot!
    chosen_letter = index_to_letter[chosen_index]
    selected_strategy = anonymous_menu[chosen_letter]
    mlp.current_activation = selected_strategy["func"]
    
    print("\n--- Meta Decision Metrics (Shuffled!) ---")
    for idx, prob in enumerate(probs):
        true_strat = anonymous_menu[index_to_letter[idx+1]]['label']
        print(f"Option [{idx+1}] containing {true_strat} ({shuffled_deltas[idx]:+.4f}): Probability = {prob:.2f}")
        
    print(f"✨ SmolLM2 picked Option Index {chosen_index}")
    print(f"   [System Execution mapping]: Revealed to be {selected_strategy['label']} -> Operating on backend '{selected_strategy['func'].upper()}'")