import numpy as np
from scipy.io.wavfile import read
from vgpu_cache import VGPUCache   # <-- import the generalized cache

cache = VGPUCache()   # global instance

# Load 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)

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).astype(np.float32) * 0.01
        self.b1 = np.zeros((1, hidden_size), dtype=np.float32)
        self.W2 = np.random.randn(hidden_size, output_size).astype(np.float32) * 0.01
        self.b2 = np.zeros((1, output_size), dtype=np.float32)
        self.lr = 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):
        # z1 = X @ W1 + b1
        self.z1 = cache.matmul(X, self.W1) + self.b1
        self.a1 = self.relu(self.z1)
        # z2 = a1 @ W2 + b2
        self.z2 = cache.matmul(self.a1, self.W2) + self.b2
        return self.softmax(self.z2)

    def backward(self, X, y_true, y_pred):
        m = y_true.shape[0]
        dz2 = y_pred - y_true
        # dW2 = (a1.T @ dz2) / m
        dW2 = cache.matmul(self.a1.T, dz2) / m
        db2 = np.sum(dz2, axis=0, keepdims=True) / m
        # da1 = dz2 @ W2.T
        da1 = cache.matmul(dz2, self.W2.T)
        dz1 = da1 * self.relu_derivative(self.z1)
        # dW1 = (X.T @ dz1) / m
        dW1 = cache.matmul(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.lr[0] * dW1
        self.b1 -= self.lr[1] * db1
        self.W2 -= self.lr[2] * dW2
        self.b2 -= self.lr[3] * db2

    def predict(self, X):
        return np.argmax(self.forward(X), axis=1)

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

# Training loop
if __name__ == "__main__":
    model = MLPClassifier(784, 100, 10)
    for i in range(1000):   # you can change the loop condition
        idx = np.random.randint(0, 60000, 100)
        X = X_train[idx]
        y = y_train[idx]
        y_onehot = np.eye(10, dtype=np.float32)[y]
        model.update(X, y_onehot)
        if i % 100 == 0:
            acc = model.score(X_test[:1000], y_test[:1000])
            print(f"Iter {i}: train acc {model.score(X, y):.3f}, test acc {acc:.3f}")

    print("\nCache stats:", cache._stats)