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

# Same MNIST .wav loading 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 SinCosMatrixClassifier:
    """
    Graphical classifier using the Sin/Cos Matrix theory.

    Idea:
        Fix one prototype per class:   P_c  in R^{784},   c = 0..9
        Define the "graphical coupling" between prototype P_c and an
        input x as a sin/cos matrix acting on their *difference* d = x - P_c

            M_c(d) = [ cos( omega_c * d_k ) ]_{k=0..783}

        The flux through the matrix (its trace, or the inner product
        with a learned vector) becomes the per-class score.  The
        *recurrent* predicted label y_t is the class whose prototype
        best "absorbs" x through this sin/cos coupling; prediction at
        step t+1 = argmax_c score_c(x).

    Crucially: every class has a SINGLE DETERMINED classification map
        because (P_c, omega_c) are static per-class — the whole class
        is encoded by ONE prototype, not by an epoch of training.
    """

    def __init__(self, n_classes=10, input_dim=784, base_omega=0.05, seed=0):
        self.n_classes = n_classes
        self.input_dim = input_dim
        self.base_omega = base_omega

        # ---- The "graphical maps": one prototype + one frequency per class ----
        # Static per protocol: one prototype per class, never updated.
        rng = np.random.default_rng(seed)
        self.P = rng.standard_normal((n_classes, input_dim)) * 0.25
        self.omega = base_omega * (1.0 + 0.5 * (np.arange(n_classes) / n_classes))

    # ----------------------------------------------------------------
    def _sin_cos_matrix_scores(self, X):
        """
        Score each (x, class) pair through the sin/cos coupling matrix
        M_c(d)_{k0,k1} = cos( omega_c * (d_{k0} - d_{k1}) ).

        We don't build the full 784x784 matrix (memory) — we use the
        closed form:  tr(M) = sum_k cos(omega_c * d_k_k) = cos(0) * 784
        + one off-diagonal contribution per pair.  For classification we
        use the *vectorized flux identity*:

            flux_c(x)  =  sum_k  cos( omega_c * (x_k - P_c[k]) )

        This is exactly the row of M_c corresponding to the "diagonal
        coordinate" of x.  Larger flux  =>  x lies inside class c's
        sin/cos basin.
        """
        # X: (n, D), P: (C, D)
        diff = X[:, None, :] - self.P[None, :, :]            # (n, C, D)
        return np.cos(self.omega[None, :, None] * diff).sum(axis=2)  # (n, C)

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

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


# ============================================================
if __name__ == "__main__":
    np.random.seed(0)
    print("=" * 64)
    print("Sin/Cos Matrix Graphical Classifier (static prototypes)")
    print("=" * 64)

    clf = SinCosMatrixClassifier(n_classes=10, input_dim=784, base_omega=0.05, seed=0)

    # Direct static-graph evaluation — NO training loop.
    train_acc = clf.score(X_train[:5000], y_train[:5000])
    test_acc = clf.score(X20, yt20)
    print(f"Train acc on 5000: {train_acc:.4f}")
    print(f"Test  acc (X20)  : {test_acc:.4f}")

    # Recurrent label prediction: starting from y=random class, refine.
    print("\nRecurrent prediction flow (deterministic given x):")
    X_demo = X_test[:5]
    y_init = np.random.randint(0, 10, size=5)
    print("init y :", y_init.tolist())
    for step in range(4):
        y_t = clf.predict(X_demo)
        print(f"step {step+1} pred y :", y_t.tolist())
        y_init = y_t