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)

# Use a cheat subset of the training data as a proxy test set (kept strictly in gray format)
cheat_test_idx = np.random.choice(X_test.shape[0], 1000, replace=False)
X_cheat_test_gray = X_test[cheat_test_idx]
y_cheat_test = y_test[cheat_test_idx]

# Helper to convert grayscale images to 3-channel RGB
def to_rgb(X_gray):
    # Reshape to (N, 28, 28), expand dims to channel axis, and repeat 3 times
    N = X_gray.shape[0]
    images = X_gray.reshape(N, 28, 28, 1)
    return np.repeat(images, 3, axis=-1)

def apply_color_retrofit(X_gray):
    """
    Transforms grayscale images into RGB format and applies arbitrary colorations
    to teach the network color invariance.
    """
    N = X_gray.shape[0]
    rgb_images = to_rgb(X_gray).astype(float)
    
    for i in range(N):
        # Example retrofit color injection: Choose a random channel to boost or tint
        # Alternatively, you can tint the background, add random shifts, etc.
        color_tint = np.random.rand(3) * 255.0  
        mask = rgb_images[i] > 30 # Apply color mostly to where the digit is drawn
        
        # Blend random color with existing pixel intensity
        rgb_images[i][mask[:, :, 0]] *= (color_tint / 255.0) 
        
    # Flatten back to vector format for the MLP (28 * 28 * 3 = 2352)
    return rgb_images.reshape(N, -1)

def prepare_gray_as_rgb(X_gray):
    """
    Converts pure grayscale images to 3-channel representation without changing color balance,
    ensuring it fits the model's new input dimensions.
    """
    N = X_gray.shape[0]
    return to_rgb(X_gray).reshape(N, -1)


# Define the MLP Classifier
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):
        self.z1 = np.dot(X, self.W1) + self.b1
        self.a1 = self.relu(self.z1)
        self.z2 = np.dot(self.a1, self.W2) + self.b2
        output = self.softmax(self.z2)
        return output
    
    def compute_loss(self, y_true, y_pred):
        m = y_true.shape[0]
        loss = -np.sum(y_true * np.log(y_pred + 1e-9)) / m
        return loss
    
    def backward(self, X, y_true, y_pred):
        m = y_true.shape[0]
        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.relu_derivative(self.z1)
        dW1 = np.dot(X.T, dz1) / m
        db1 = np.sum(dz1, axis=0, keepdims=True) / m
        return dW1, db1, dW2, db2
    
    def update(self, X, y_true):
        y_pred = self.forward(X)
        dW1, db1, dW2, db2 = self.backward(X, y_true, y_pred)
        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):
        probabilities = self.forward(X)
        return np.argmax(probabilities, axis=1)

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

# Initialize and train the MLP Classifier
if __name__ == "__main__":
    learning_rate = np.random.rand(4)
    # The input dimension scales up to 784 * 3 channels = 2352
    f = MLPClassifier(input_size=2352, hidden_size=100, output_size=10, learning_rate=learning_rate)

    # Convert our cheat test set into standard gray-RGB format for continuous evaluation
    X_cheat_test_rgb = prepare_gray_as_rgb(X_cheat_test_gray)

    i = 0
    while True:
        idx = np.random.randint(0, 60000, 100)
        
        # 1. Grab raw training batches
        X_batch_gray = X_train[idx]
        yt = y_train[idx]
        
        # 2. Inject retrofitted colors into the training batch
        X_batch_rgb = apply_color_retrofit(X_batch_gray)
        
        # 3. Evaluate progress using the pure grayscale cheat-test substitute
        if i % 10 == 0:
            test_score = f.score(X_cheat_test_rgb, y_cheat_test)
            print(f"Iteration {i:4d} | Train Batch Accuracy: {f.score(X_batch_rgb, yt):.4f} | Cheat Test Accuracy (Gray): {test_score:.4f}")
        
        # 4. Backprop on the colorized samples
        f.update(X_batch_rgb, np.eye(10)[yt])
        i += 1