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

# Load data (same as before)
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 LearnableSinCosClassifier:
    """
    Learnable classifier based on sin/cos matrix theory.

    Architecture (multi-pool Fourier features):
      1. Linear projection:          x -> z = W*x + b   (776 -> 100)
      2. Many independent phases:    theta_k = atan2(z[:, 2k+1], z[:, 2k])
      3. Fourier features per phase: [cos(theta_k), sin(theta_k), ..., cos(K*theta_k), sin(K*theta_k)]
      4. Linear softmax classifier on the Fourier feature bank.
    """

    def __init__(self, input_dim=784, hidden_dim=100, n_frequencies=16, n_classes=10,
                 lr=(0.01, 0.01, 0.005, 0.01)):
        self.input_dim = input_dim
        self.hidden_dim = hidden_dim
        self.n_frequencies = n_frequencies
        self.n_classes = n_classes

        if hidden_dim % 2 != 0:
            raise ValueError("hidden_dim must be even so it splits into (z0, z1) pairs.")
        self.n_pools = hidden_dim // 2
        self.feat_dim = self.n_pools * 2 * n_frequencies  # raw features before pooling
        self.pooled_dim = self.n_pools * 2 * n_frequencies

        # ---- FIX 2: proper init, not * 0.01 ----
        self.W = np.random.randn(input_dim, hidden_dim) * np.sqrt(2.0 / input_dim)
        self.b = np.zeros((1, hidden_dim))

        self.W_cls = np.random.randn(self.pooled_dim, n_classes) * np.sqrt(2.0 / self.pooled_dim)
        self.b_cls = np.zeros((1, n_classes))

        # Either a tuple/list or a length-4 array; both work with indexing.
        self.lr = np.array(lr, dtype=float)

        # Cache for backprop
        self.z = None
        self.thetas = None   # shape (n, n_pools)
        self.features = None

    # ------------------------------------------------------------------
    def _fourier_per_pool(self, thetas):
        """
        Per-pool Fourier bank:
            for k = 1..n_frequencies:
                [cos(k*theta_p), sin(k*theta_p)]
        then concatenate across pools.
        """
        n = thetas.shape[0]
        K = self.n_frequencies
        P = self.n_pools

        # thetas[:, p] (shape n) -> k * thetas[:, p] for each (k, p)
        # Resulting shape (K, n, P)
        ks = np.arange(1, K + 1).reshape(K, 1, 1)
        angles = ks * thetas.reshape(1, n, P)

        cos_block = np.cos(angles)   # (K, n, P)
        sin_block = np.sin(angles)   # (K, n, P)

        # interleave cos and sin along the frequency axis: (K, n, P, 2)
        per_freq = np.stack([cos_block, sin_block], axis=-1)
        # flatten K and the cos/sin dim together: (n, P, 2K)
        per_pool = per_freq.transpose(1, 2, 0, 3).reshape(n, P, 2 * K)
        # finally flatten pools: (n, P*2K)
        return per_pool.reshape(n, P * 2 * K)

    # ------------------------------------------------------------------
    def forward(self, X):
        # Projection
        self.z = np.dot(X, self.W) + self.b  # (n, hidden_dim)

        # ---- FIX 4 + FIX 3: many phases, each from a (z0, z1) pair, no epsilon ----
        z0 = self.z[:, 0::2]   # even columns (n, n_pools)
        z1 = self.z[:, 1::2]   # odd columns  (n, n_pools)
        self.thetas = np.arctan2(z1, z0)     # (n, n_pools)

        # Fourier bank
        self.features = self._fourier_per_pool(self.thetas)   # (n, pooled_dim)
        logits = np.dot(self.features, self.W_cls) + self.b_cls
        return logits

    # ------------------------------------------------------------------
    @staticmethod
    def softmax(logits):
        e = np.exp(logits - np.max(logits, axis=1, keepdims=True))
        return e / np.sum(e, axis=1, keepdims=True)

    def compute_loss(self, logits, y_true):
        probs = self.softmax(logits)
        m = y_true.shape[0]
        return -np.sum(np.log(probs[np.arange(m), y_true] + 1e-12)) / m

    # ------------------------------------------------------------------
    def backward(self, X, y_true, logits):
        m = y_true.shape[0]
        y_onehot = np.eye(self.n_classes)[y_true]
        probs = self.softmax(logits)
        dlogits = (probs - y_onehot) / m                          # (n, n_classes)

        # Classification layer
        dW_cls = np.dot(self.features.T, dlogits)                 # (pooled_dim, n_classes)
        db_cls = np.sum(dlogits, axis=0, keepdims=True)           # (1, n_classes)

        # Back into features
        dfeatures = np.dot(dlogits, self.W_cls.T)                 # (n, pooled_dim)

        # Reshape to (n, P, 2K) so we can compute dθ per pool
        n, P = m, self.n_pools
        K = self.n_frequencies
        dfeat_per_pool = dfeatures.reshape(n, P, 2 * K)

        # split cos / sin columns; remember ordering: cos(k*), sin(k*),
        # k = 1..K, all per pool.
        # d/cost = [-k sin(kθ)], d/sint = [k cos(kθ)]
        # The cos block is at index 2*(k-1), sin at 2*(k-1)+1 within each pool.
        dcos = dfeat_per_pool[:, :, 0::2]   # (n, P, K)
        dsin = dfeat_per_pool[:, :, 1::2]   # (n, P, K)

        ks = np.arange(1, K + 1).reshape(1, 1, K)
        # cos(kθ) derivative wrt θ = -k sin(kθ)
        # sin(kθ) derivative wrt θ =  k cos(kθ)
        dtheta_per_pool = np.sum(
            dcos * (-ks * np.sin(ks * self.thetas.reshape(n, P, 1))) +
            dsin * ( ks * np.cos(ks * self.thetas.reshape(n, P, 1))),
            axis=2
        )                                                           # (n, P)

        # dθ/dz0 = -z1 / (z0^2 + z1^2), dθ/dz1 = z0 / (z0^2 + z1^2)
        z0 = self.z[:, 0::2]
        z1 = self.z[:, 1::2]
        norm2 = z0 ** 2 + z1 ** 2                  # (n, P)

        dz = np.zeros_like(self.z)                 # (n, hidden_dim)
        dz_even = dtheta_per_pool * (-z1 / norm2)   # (n, P), lives at columns 0::2
        dz_odd = dtheta_per_pool * ( z0 / norm2)   # (n, P), lives at columns 1::2
        dz[:, 0::2] = dz_even
        dz[:, 1::2] = dz_odd

        # Projection layer
        dW = np.dot(X.T, dz)                       # (input_dim, hidden_dim)
        db = np.sum(dz, axis=0, keepdims=True)     # (1, hidden_dim)
        return dW, db, dW_cls, db_cls

    # ------------------------------------------------------------------
    def update(self, X, y_true):
        logits = self.forward(X)
        loss = self.compute_loss(logits, y_true)
        dW, db, dW_cls, db_cls = self.backward(X, y_true, logits)

        # ---- FIX 1: fixed lr, no longer np.random.rand ----
        self.W     -= self.lr[0] * dW
        self.b     -= self.lr[1] * db
        self.W_cls -= self.lr[2] * dW_cls
        self.b_cls -= self.lr[3] * db_cls
        return loss

    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)


# ============================================================
if __name__ == "__main__":
    print("=" * 60)
    print("Learnable Sin/Cos Matrix Classifier for MNIST (fixed)")
    print("=" * 60)

    model = LearnableSinCosClassifier(
        input_dim=784,
        hidden_dim=100,            # 50 independent (z0, z1) pools
        n_frequencies=16,          # -> 50 * 32 = 1600 features
        n_classes=10,
        lr=(0.01, 0.01, 0.005, 0.01),
    )

    print(f"hidden_dim={model.hidden_dim}, n_pools={model.n_pools}, "
          f"n_frequencies={model.n_frequencies}, "
          f"pooled_dim={model.pooled_dim}")

    epochs = 20000
    batch = 1000
    print_every = 10

    for epoch in range(1, epochs + 1):
        idx = np.random.randint(0, 60000, batch)
        X = X_train[idx]
        yt = y_train[idx]
        loss = model.update(X, yt)

        if epoch % print_every == 0 or epoch == 1:
            test_acc = model.score(X20, yt20)
            train_acc = model.score(X, yt)
            print(f"step {epoch:4d}: loss={loss:.4f}  "
                  f"train_acc={train_acc:.4f}  test_acc={test_acc:.4f}")

    print(f"\nFinal test accuracy (X20, first 1000): "
          f"{model.score(X20, yt20):.4f}")