import numpy as np
from scipy.io.wavfile import read

# ---------- Product Quantization utilities ----------
def pq_dot_loop(X, W_idx, codebooks):
    """
    Compute X @ W_factorized where W is represented by codebooks and indices.
    X: (batch_size, K) float32
    W_idx: (d_out, M) int8 – each row selects centroids per subspace
    codebooks: list of M arrays, each (C, d) float32, where d = K // M
    Returns: (batch_size, d_out) float32
    """
    M = len(codebooks)          # number of subspaces
    d = X.shape[1] // M         # subspace dimension
    batch_size = X.shape[0]
    d_out = W_idx.shape[0]

    # Precompute projections: list of (batch_size, C) per subspace
    X_proj_list = []
    for m, cb in enumerate(codebooks):
        X_sub = X[:, m*d:(m+1)*d]          # (batch_size, d)
        X_proj_list.append(X_sub @ cb.T)   # (batch_size, C)

    O = np.zeros((batch_size, d_out), dtype=np.float32)
    for i in range(batch_size):
        for j in range(d_out):
            acc = 0.0
            for m in range(M):
                centroid_idx = W_idx[j, m]      # integer
                acc += X_proj_list[m][i, centroid_idx]
            O[i, j] = acc
    return O

class PQLinear:
    """
    A linear layer where the weight matrix is factorised via Product Quantization.
    The weight is represented by:
      - W_idx: (out_features, M) int8, indices into codebooks
      - codebooks: list of M arrays, each (n_centroids, subspace_dim)
    The forward pass computes X @ W_approx via pq_dot_loop.
    The backward pass computes gradients for each centroid (indices are fixed).
    """
    def __init__(self, in_features, out_features, M, n_centroids=256, init='random'):
        self.M = M
        self.d = in_features // M
        self.out_features = out_features
        self.n_centroids = n_centroids

        # Initialize codebooks (centroids) with random normal
        self.codebooks = [np.random.randn(n_centroids, self.d).astype(np.float32) * 0.01
                          for _ in range(M)]

        # Initialize indices randomly
        self.W_idx = np.random.randint(0, n_centroids, size=(out_features, M)).astype(np.int8)

        # Gradient accumulators for codebooks (same shape as codebooks)
        self.grad_codebooks = [np.zeros_like(cb) for cb in self.codebooks]

    def forward(self, X):
        """Forward pass: compute X @ W_approx."""
        self.X = X  # store for backward
        self.out = pq_dot_loop(X, self.W_idx, self.codebooks)
        return self.out

    def backward(self, grad_output):
        """
        Backward pass: given gradient of loss w.r.t. output (grad_output),
        compute gradients w.r.t. codebooks (and also w.r.t. X if needed).
        We accumulate gradients in self.grad_codebooks.
        Returns gradient w.r.t. X (for upstream layers).
        """
        batch_size, d_out = grad_output.shape
        M = self.M
        d = self.d
        X = self.X
        W_idx = self.W_idx

        # Zero out gradients (or accumulate if we want to accumulate across batches)
        for g in self.grad_codebooks:
            g.fill(0.0)

        # For each subspace, compute gradient w.r.t. its codebook.
        # The output O[i,j] = sum_m X_proj_list[m][i, W_idx[j,m]]
        # where X_proj_list[m] = X_sub @ codebooks[m].T
        # So grad wrt codebook[m][c, :] accumulates over i,j where W_idx[j,m]==c
        # and the contribution is grad_output[i,j] * X_sub[i, :].
        for m in range(M):
            cb = self.codebooks[m]          # (C, d)
            X_sub = X[:, m*d:(m+1)*d]       # (batch, d)
            # For each centroid c, we need to sum over j where W_idx[j,m]==c and over i
            for c in range(self.n_centroids):
                # Find which output neurons use centroid c for subspace m
                mask = (W_idx[:, m] == c)   # (out_features,)
                if not np.any(mask):
                    continue
                # grad_output[:, mask] shape (batch, n_active)
                # sum over active output neurons and over batch dimension
                # grad_codebook[m][c, :] = sum_i sum_{j in active} grad_output[i,j] * X_sub[i,:]
                # This is: (grad_output[:, mask].T @ X_sub).sum(axis=0)? Let's compute carefully.
                # For each i, we have sum_j grad_output[i,j] * X_sub[i,:] summed over j in active.
                # That is equivalent to: sum_i X_sub[i,:] * (sum_{j in active} grad_output[i,j])
                # So we can compute: temp = grad_output[:, mask].sum(axis=1)  # (batch,)
                # then grad = temp.T @ X_sub  # (d,)
                temp = grad_output[:, mask].sum(axis=1)   # (batch,)
                self.grad_codebooks[m][c, :] = temp @ X_sub   # (d,)

        # Gradient w.r.t. X: for each subspace, we have
        # dO/dX_sub = sum_{j} grad_output[i,j] * codebooks[m][W_idx[j,m], :]
        # so we compute for each i.
        grad_X = np.zeros_like(X)
        for m in range(M):
            cb = self.codebooks[m]
            # For each i, we need to sum over j: grad_output[i,j] * cb[W_idx[j,m], :]
            # This is like a matrix multiplication: grad_output (batch, out) times
            # a matrix where each row j is cb[W_idx[j,m], :] (d,)
            # We can compute for each i: sum_j grad_output[i,j] * cb[W_idx[j,m], :]
            # That is: grad_output @ cb[W_idx[:, m], :]  (but indexing cb with W_idx[:,m] gives (out, d))
            # So grad_X_sub = grad_output @ cb[W_idx[:, m], :]  # (batch, d)
            selected = cb[W_idx[:, m], :]   # (out, d)
            grad_X_sub = grad_output @ selected   # (batch, d)
            grad_X[:, m*d:(m+1)*d] += grad_X_sub

        # Normalize gradients by batch size (optional, done in update step)
        # We'll scale in the update method.

        return grad_X

    def update(self, learning_rate):
        """Update codebooks using accumulated gradients."""
        for m in range(self.M):
            self.codebooks[m] -= learning_rate * self.grad_codebooks[m] / self.X.shape[0]

    def score(self, X, y_true):  # for compatibility
        # Dummy – we use the outer MLP score.
        pass

# ---------- MLP Classifier with PQ layers ----------
class MLPClassifier:
    def __init__(self, input_size, hidden_size, output_size,
                 use_pq=False, M=8, n_centroids=256, learning_rate=np.random.rand(4)):
        """
        If use_pq is True, replace dense weight matrices with PQLinear layers.
        """
        self.use_pq = use_pq
        self.learning_rate = learning_rate  # [lr_W1, lr_b1, lr_W2, lr_b2]

        if use_pq:
            # W1: (input_size, hidden_size) factorized
            self.W1 = PQLinear(input_size, hidden_size, M, n_centroids)
            # W2: (hidden_size, output_size) factorized
            self.W2 = PQLinear(hidden_size, output_size, M, n_centroids)
            # We still need bias vectors as dense parameters
            self.b1 = np.zeros((1, hidden_size))
            self.b2 = np.zeros((1, output_size))
        else:
            # Dense version (original)
            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))

    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):
        if self.use_pq:
            # W1 is PQLinear
            self.z1 = self.W1.forward(X) + self.b1
        else:
            self.z1 = np.dot(X, self.W1) + self.b1
        self.a1 = self.relu(self.z1)

        if self.use_pq:
            self.z2 = self.W2.forward(self.a1) + self.b2
        else:
            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
        # gradients for W2 and b2
        if self.use_pq:
            # W2 is PQLinear
            # compute gradient w.r.t. input of W2 (i.e., a1) via backward
            grad_a1 = self.W2.backward(dz2)   # returns gradient w.r.t. a1
            # gradients for codebooks will be stored inside self.W2.grad_codebooks
            # we don't compute dW2 directly; we'll update later via W2.update
            # for b2
            db2 = np.sum(dz2, axis=0, keepdims=True) / m
        else:
            dW2 = np.dot(self.a1.T, dz2) / m
            db2 = np.sum(dz2, axis=0, keepdims=True) / m
            # gradient w.r.t. a1 for backprop to W1
            grad_a1 = np.dot(dz2, self.W2.T)

        # backprop through ReLU
        dz1 = grad_a1 * self.relu_derivative(self.z1)

        # gradients for W1 and b1
        if self.use_pq:
            # W1 is PQLinear
            grad_X = self.W1.backward(dz1)   # gradient w.r.t. X (input)
            # We don't use grad_X for updating input, but it's computed.
            db1 = np.sum(dz1, axis=0, keepdims=True) / m
        else:
            dW1 = np.dot(X.T, dz1) / m
            db1 = np.sum(dz1, axis=0, keepdims=True) / m

        # Store gradients for update
        if not self.use_pq:
            self.dW1 = dW1
            self.dW2 = dW2
        self.db1 = db1
        self.db2 = db2

    def update(self, X, y_true):
        y_pred = self.forward(X)
        self.backward(X, y_true, y_pred)

        if self.use_pq:
            # Update codebooks with learning rates
            self.W1.update(self.learning_rate[0])
            self.W2.update(self.learning_rate[2])
            # Update biases
            self.b1 -= self.learning_rate[1] * self.db1
            self.b2 -= self.learning_rate[3] * self.db2
        else:
            self.W1 -= self.learning_rate[0] * self.dW1
            self.b1 -= self.learning_rate[1] * self.db1
            self.W2 -= self.learning_rate[2] * self.dW2
            self.b2 -= self.learning_rate[3] * self.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)

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

# ---------- Training ----------
if __name__ == "__main__":
    # Choose whether to use PQ
    use_pq = True   # set to False to use dense version
    learning_rate = np.random.rand(4)
    f = MLPClassifier(input_size=784, hidden_size=100, output_size=10,
                      use_pq=use_pq, M=8, n_centroids=256,
                      learning_rate=learning_rate)

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