import numpy as np

# Linear component
x = np.linspace(0, 1, 100)

# 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]


class MLPClassifier:
    def __init__(self, layer_dims, learning_rate=None):
        """
        layer_dims: list of layer sizes [input, hidden1, hidden2, ..., output]
        10 layers total: alternating linear and ReLU
        Even indices (0,2,4,6,8): linear layers
        Odd indices (1,3,5,7): ReLU layers
        Last layer (9): softmax output
        """
        self.layer_dims = layer_dims
        self.num_layers = len(layer_dims)
        self.x = x
        
        # Initialize weights and biases
        self.weights = []
        self.biases = []
        
        for i in range(self.num_layers - 1):
            # Linear layer weights
            W = np.random.randn(layer_dims[i], layer_dims[i+1]) * 0.01
            b = np.zeros((1, layer_dims[i+1]))
            self.weights.append(W)
            self.biases.append(b)
        
        # Learning rates for each layer (W and b for each)
        if learning_rate is None:
            self.learning_rate = np.random.rand(self.num_layers - 1) * 0.01
        else:
            self.learning_rate = learning_rate
    
    def linear(self, Z):
        """Linear activation (identity)"""
        return Z
    
    def relu(self, Z):
        """ReLU activation"""
        return np.maximum(0, Z)
    
    def relu_derivative(self, Z):
        """ReLU derivative"""
        return np.where(Z > 0, 1, 0)
    
    def softmax(self, Z):
        """Softmax activation"""
        exp_Z = np.exp(Z - np.max(Z, axis=1, keepdims=True))
        return exp_Z / np.sum(exp_Z, axis=1, keepdims=True)
    
    def forward_layer(self, A, layer_idx):
        """Forward pass through single layer"""
        Z = np.dot(A, self.weights[layer_idx]) + self.biases[layer_idx]
        
        # Apply activation based on layer type
        if layer_idx == self.num_layers - 2:
            # Last layer: softmax
            A_next = self.softmax(Z)
        elif layer_idx % 2 == 0:
            # Even layer (0,2,4,6,8): linear
            A_next = self.linear(Z)
        else:
            # Odd layer (1,3,5,7): ReLU
            A_next = self.relu(Z)
        
        return Z, A_next
    
    def forward(self, X):
        """Full forward pass through all layers"""
        self.cache_Z = []  # Store Z values for backprop
        self.cache_A = [X]  # Store A values (activations)
        
        A = X
        for i in range(self.num_layers - 1):
            Z, A = self.forward_layer(A, i)
            self.cache_Z.append(Z)
            self.cache_A.append(A)
        
        return A  # Output predictions
    
    def backward_layer(self, dA, layer_idx):
        """Backward pass through single layer"""
        Z = self.cache_Z[layer_idx]
        A_prev = self.cache_A[layer_idx]
        W = self.weights[layer_idx]
        
        # Compute dZ based on activation type
        if layer_idx == self.num_layers - 2:
            # Last layer: dZ = dA (softmax gradient handled in output layer)
            dZ = dA
        elif layer_idx % 2 == 0:
            # Even layer: linear (dZ = dA)
            dZ = dA
        else:
            # Odd layer: ReLU
            dZ = dA * self.relu_derivative(Z)
        
        # Gradients
        m = A_prev.shape[0]
        dW = np.dot(A_prev.T, dZ) / m
        db = np.sum(dZ, axis=0, keepdims=True) / m
        
        # Pass gradient to previous layer
        dA_prev = np.dot(dZ, W.T)
        
        return dA_prev, dW, db
    
    def backward(self, y_true, y_pred):
        """Full backward pass"""
        self.grad_weights = []
        self.grad_biases = []
        
        # Output layer gradient (softmax cross-entropy)
        dA = y_pred - y_true
        
        for i in reversed(range(self.num_layers - 1)):
            dA, dW, db = self.backward_layer(dA, i)
            self.grad_weights.insert(0, dW)
            self.grad_biases.insert(0, db)
        
        return self.grad_weights, self.grad_biases
    
    def update(self, X, y_true):
        """Update weights and biases"""
        y_pred = self.forward(X)
        dW, db = self.backward(y_true, y_pred)
        
        for i in range(self.num_layers - 1):
            self.weights[i] -= self.learning_rate[i] * dW[i]
            self.biases[i] -= self.learning_rate[i] * db[i]
    
    def predict(self, X):
        """Make predictions"""
        probabilities = self.forward(X)
        return np.argmax(probabilities, axis=1)
    
    def score(self, X, y_true):
        """Calculate accuracy"""
        return np.mean(self.predict(X) == y_true)


# Initialize model with 10 layers
# Architecture: 784 -> 100 -> 100 -> 100 -> 100 -> 100 -> 100 -> 100 -> 100 -> 10
# Layer indices: 0     1     2     3     4     5     6     7     8     9
# Types:          L    R    L    R    L    R    L    R    L    S
# L = Linear, R = ReLU, S = Softmax

layer_dims = [784, 100, 100, 100, 100, 100, 100, 100, 100, 10]
learning_rate = np.random.rand(len(layer_dims) - 1) * 0.01

f = MLPClassifier(layer_dims=layer_dims, learning_rate=learning_rate)

# Training loop
i = 0
while True:
    idx = np.random.randint(0, 60000, 100)
    X = X_train[idx]
    yt = y_train[idx]
    f.update(X, np.eye(10)[yt])
    print(i, f.score(X20, yt20))
    i += 1