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)
X20 = X_test[:1000]
yt20 = y_test[:1000]


class ModMLPClassifier:
    """
    MLP with modulo-based weight sharing.

    Instead of storing full weight matrices, we store small "base" matrices.
    The full weight matrix is reconstructed on-the-fly using modulo indexing:

        W_full[i, j] = W_base[i % mod, j]

    This dramatically reduces parameter count while introducing a useful
    inductive bias: input features with the same modulo index share weights.

    For MNIST (28x28 images):
        - w1_mod=28  → W1 goes from (784, 100)=78,400  to  (28, 100)=2,800
        - w2_mod=10  → W2 goes from (100, 10)  =1,000   to  (10, 10) =100
        Total: 79,400 → 2,900  (~27x reduction)
    """

    def __init__(self, input_size, hidden_size, output_size,
                 w1_mod=28, w2_mod=10, learning_rate=None):
        if learning_rate is None:
            learning_rate = np.full(4, 0.005, dtype=np.float32)

        # ── Small base weight matrices ──
        self.W1_base = np.random.randn(w1_mod, hidden_size) * np.sqrt(2.0 / w1_mod)
        self.b1 = np.zeros((1, hidden_size))
        self.W2_base = np.random.randn(w2_mod, output_size) * np.sqrt(2.0 / hidden_size)
        self.b2 = np.zeros((1, output_size))

        self.learning_rate = np.asarray(learning_rate, dtype=np.float32)
        self.w1_mod = w1_mod
        self.w2_mod = w2_mod
        self.input_size = input_size
        self.hidden_size = hidden_size
        self.output_size = output_size

        # Precompute index lookups (constant, never change)
        self._idx_in = np.arange(input_size) % w1_mod   # (784,) → [0..27, 0..27, ...]
        self._idx_h  = np.arange(hidden_size) % w2_mod  # (100,) → [0..9, 0..9, ...]

    # ── Activation functions (unchanged) ──
    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)

    # ── Forward pass: expand base weights with modulo ──
    def forward(self, X):
        # Expand: W1_eff[i, :] = W1_base[i % w1_mod, :]
        W1_eff = self.W1_base[self._idx_in]    # (784, 100)
        self.z1 = np.dot(X, W1_eff) + self.b1
        self.a1 = self.relu(self.z1)

        # Expand: W2_eff[i, :] = W2_base[i % w2_mod, :]
        W2_eff = self.W2_base[self._idx_h]     # (100, 10)
        self.z2 = np.dot(self.a1, W2_eff) + self.b2
        output = self.softmax(self.z2)
        return output

    def compute_loss(self, y_true, y_pred):
        m = y_true.shape[0]
        return -np.sum(y_true * np.log(y_pred + 1e-9)) / m

    # ── Backward pass: compute full gradients, then fold back with modulo ──
    def backward(self, X, y_true, y_pred):
        m = y_true.shape[0]

        # Reconstruct effective weights (needed for da1)
        W2_eff = self.W2_base[self._idx_h]     # (100, 10)

        dz2 = (y_pred - y_true) / m            # (m, 10)

        # --- W2 gradient ---
        dW2_full = np.dot(self.a1.T, dz2)      # (100, 10) — full gradient
        dW2_base = np.zeros_like(self.W2_base)
        np.add.at(dW2_base, self._idx_h, dW2_full)

        db2 = np.sum(dz2, axis=0, keepdims=True)

        # --- W1 gradient ---
        da1 = np.dot(dz2, W2_eff.T)            # (m, 100)
        dz1 = da1 * self.relu_derivative(self.z1)

        dW1_full = np.dot(X.T, dz1)            # (784, 100) — full gradient
        dW1_base = np.zeros_like(self.W1_base)
        np.add.at(dW1_base, self._idx_in, dW1_full)

        db1 = np.sum(dz1, axis=0, keepdims=True)

        return dW1_base, db1, dW2_base, db2

    # ── Update (same structure, now operates on base weights) ──
    def update(self, X, y_true):
        y_pred = self.forward(X)
        dW1_base, db1, dW2_base, db2 = self.backward(X, y_true, y_pred)
        self.W1_base -= self.learning_rate[0] * dW1_base
        self.b1      -= self.learning_rate[1] * db1
        self.W2_base -= self.learning_rate[2] * dW2_base
        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)


# ── Training loop ──
np.random.seed(0)
#learning_rate = np.full(4, 0.005, dtype=np.float32)
learning_rate = np.random.rand(4) *0.01
f = ModMLPClassifier(
    input_size=784,
    hidden_size=100,
    output_size=10,
    w1_mod=28,       # 784 → 28 base weights per hidden unit
    w2_mod=10,       # 100 → 10 base weights per output unit
    learning_rate=learning_rate,
)

# Quick parameter-count sanity check
n_params = (f.W1_base.size + f.b1.size + f.W2_base.size + f.b2.size)
print(f"Learnable parameters: {n_params}  (was 79,400 in the original)")
# →  Learnable parameters: 2900  (was 79,400 in the original)

i = 0
while True:
    idx = np.random.randint(0, len(X_train), 128)
    X = X_train[idx]
    yt = y_train[idx]
    f.update(X, np.eye(10)[yt])
    if i % 20 == 0:
        train_acc = f.score(X, yt)
        test_acc = f.score(X20, yt20)
        print(i, f"train={train_acc:.3f}", f"test={test_acc:.3f}")
    i += 1