import numpy as np
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)

# Cheat subset for clean evaluation
cheat_test_idx = np.random.choice(X_train.shape[0], 1000, replace=False)
X_cheat_test_gray = X_train[cheat_test_idx]
y_cheat_test = y_train[cheat_test_idx]

X20 = X_test[:1000]
yt20 = y_test[:1000]

# Deterministic color profiles mapping class -> [R, G, B] weights
COLOR_MAP = {
    0: [1.0, 0.0, 0.0], 1: [0.0, 1.0, 0.0], 2: [0.0, 0.0, 1.0],
    3: [1.0, 1.0, 0.0], 4: [0.0, 1.0, 1.0], 5: [1.0, 0.0, 1.0],
    6: [1.0, 0.5, 0.0], 7: [0.5, 0.0, 0.5], 8: [0.5, 0.5, 0.5],
    9: [0.2, 0.8, 0.4]
}

def prepare_gray_as_rgb(X_gray):
    """Converts standard gray images to uniform 3 channels (R=G=B)."""
    N = X_gray.shape[0]
    images = X_gray.reshape(N, 28, 28, 1)
    rgb_images = np.repeat(images, 3, axis=-1)
    return rgb_images.reshape(N, -1)

def apply_functional_color(X_gray, y_labels):
    """Transforms grayscale images into functional RGB based on the class label."""
    N = X_gray.shape[0]
    images = X_gray.reshape(N, 28, 28, 1)
    rgb_images = np.repeat(images, 3, axis=-1).astype(float)
    for i in range(N):
        rgb_images[i] = rgb_images[i] * np.array(COLOR_MAP[y_labels[i]])
    return rgb_images.reshape(N, -1)

class MLPClassifier:
    def __init__(self, input_size, hidden_size, output_size, learning_rate=np.random.rand(4)):
        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
    
    def relu(self, x): return np.maximum(0, x)
    def relu_derivative(self, x): return np.where(x > 0, 1, 0)
    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):
        # We store activations dynamically per forward pass to avoid mixups in dual updates
        state = {}
        state['z1'] = np.dot(X, self.W1) + self.b1
        state['a1'] = self.relu(state['z1'])
        state['z2'] = np.dot(state['a1'], self.W2) + self.b2
        state['output'] = self.softmax(state['z2'])
        return state
    
    def get_gradients(self, X, y_true, state):
        """Computes and returns gradients without changing the weights yet."""
        m = y_true.shape[0]
        dz2 = state['output'] - y_true
        dW2 = np.dot(state['a1'].T, dz2) / m
        db2 = np.sum(dz2, axis=0, keepdims=True) / m
        da1 = np.dot(dz2, self.W2.T)
        dz1 = da1 * self.relu_derivative(state['z1'])
        dW1 = np.dot(X.T, dz1) / m
        db1 = np.sum(dz1, axis=0, keepdims=True) / m
        return dW1, db1, dW2, db2
    
    def apply_gradients(self, dW1, db1, dW2, db2):
        self.W1 -= self.learning_rate[0] * dW1
        self.b1 -= self.learning_rate[1] * db1
        self.W2 -= self.learning_rate[2] * dW2
        self.b2 -= self.learning_rate[3] * db2

    def predict(self, X):
        state = self.forward(X)
        return np.argmax(state['output'], axis=1)

    def score(self, X, y_true):
        y = self.predict(X)
        return np.mean(y == y_true)

if __name__ == "__main__":
    learning_rate = np.random.rand(4)
    f = MLPClassifier(input_size=2352, hidden_size=500, output_size=10, learning_rate=learning_rate)

    X_cheat_test_rgb = prepare_gray_as_rgb(X_cheat_test_gray)

    i = 0
    while True:
        idx = np.random.randint(0, 60000, 100)
        X_batch_raw = X_train[idx]
        yt = y_train[idx]
        y_hot = np.eye(10)[yt]

        idx = np.random.randint(0, 60000, 100)
        X_batch_raw2 = X_train[idx]
        yt2 = y_train[idx]
        y_hot2 = np.eye(10)[yt2]

        # 1. Generate both variations for the exact same batch
        X_batch_color = apply_functional_color(X_batch_raw, yt)
        X_batch_gray  = prepare_gray_as_rgb(X_batch_raw)
        X_batch_gray2  = prepare_gray_as_rgb(X_batch_raw2)
        
        # 2. Forward pass for both representations
        state_color = f.forward(X_batch_color)
        state_gray  = f.forward(X_batch_gray)
        state_gray2  = f.forward(X_batch_gray2)
        
        # 3. Calculate gradients for both streams
        dW1_c, db1_c, dW2_c, db2_c = f.get_gradients(X_batch_color, y_hot, state_color)
        dW1_g, db1_g, dW2_g, db2_g = f.get_gradients(X_batch_gray, y_hot, state_gray)
        dW1_g2, db1_g2, dW2_g2, db2_g2 = f.get_gradients(X_batch_gray2, y_hot2, state_gray2)
        
        # 4. Blend the error signals (50% Color Task, 50% Grayscale structural baseline)
        dW1 = 0.1 * dW1_c + 0.45 * dW1_g
        db1 = 0.1 * db1_c + 0.45 * db1_g
        dW2 = 0.1 * dW2_c + 0.45 * dW2_g
        db2 = 0.1 * db2_c + 0.45 * db2_g

        dW1 += 0.45 * dW1_g2
        db1 += 0.45 * db1_g2
        dW2 += 0.45 * dW2_g2
        db2 += 0.45 * db2_g2
        
        # 5. Step weights
        f.apply_gradients(dW1, db1, dW2, db2)
    
        if i % 100 == 0:
            X20_rgb = prepare_gray_as_rgb(X20)
            test_score = f.score(X20_rgb, yt20)
            print(f"Iteration {i:5d} | Color Accuracy: {np.mean(np.argmax(state_color['output'], axis=1) == yt):.4f} | Gray Cheat-Test Accuracy: {test_score:.4f}")
        
        i += 1