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

# --- Data Loading ---
# Assuming .wav files contain MNIST data as described
try:
    X_train = read('../X_train.wav')[1].reshape(-1, 784).astype(np.float32)
    y_train = (read('../y_train.wav')[1] * 9).astype(int)
    X_test = read('../X_test.wav')[1].reshape(-1, 784).astype(np.float32)
    y_test = (read('../y_test.wav')[1] * 9).astype(int)
except FileNotFoundError:
    # Mock data for demonstration if files aren't present
    print("Files not found, using mock data for demonstration...")
    X_train = np.random.randint(0, 256, (60000, 784)).astype(np.float32)
    y_train = np.random.randint(0, 10, 60000)
    X_test = np.random.randint(0, 256, (10000, 784)).astype(np.float32)
    y_test = np.random.randint(0, 10, 10000)

# Normalization is CRITICAL for trigonometric activations
X_train /= 255.0
X_test /= 255.0

X20 = X_test[:1000]
yt20 = y_test[:1000]

class SinCosMatrixClassifier:
    """
    Optimized Graphical classifier using the Sin/Cos Matrix theory.
    """
    def __init__(self, n_classes=10, input_dim=784, base_omega=np.pi, 
                 lr_p=0.01, lr_w=0.001, seed=0):
        self.n_classes = n_classes
        self.input_dim = input_dim
        self.lr_p = lr_p
        self.lr_w = lr_w
        
        rng = np.random.default_rng(seed)
        # Initialize P in the range of normalized data [0, 1]
        self.P = rng.uniform(0, 1, (n_classes, input_dim))
        self.omega = np.full(n_classes, base_omega)

    @staticmethod
    def _softmax(s):
        # Stability trick: subtract max
        exp_s = np.exp(s - np.max(s, axis=1, keepdims=True))
        return exp_s / np.sum(exp_s, axis=1, keepdims=True)

    def _scores(self, X):
        # X: (m, D), P: (C, D), omega: (C,)
        # result: (m, C)
        # Broadcast X to (m, 1, D) and P to (1, C, D)
        diff = X[:, np.newaxis, :] - self.P[np.newaxis, :, :]
        # omega is (C,), broadcast to (1, C, 1)
        return np.cos(self.omega[np.newaxis, :, np.newaxis] * diff).sum(axis=2)

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

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

    def update(self, X, y_true):
        m = X.shape[0]
        
        # 1. Forward pass
        diff = X[:, np.newaxis, :] - self.P[np.newaxis, :, :] # (m, C, D)
        # omega_inv = self.omega[np.newaxis, :, np.newaxis]
        scores = np.cos(self.omega[np.newaxis, :, np.newaxis] * diff).sum(axis=2)
        probs = self._softmax(scores)
        
        # 2. Compute Label Error (Softmax gradient)
        y_onehot = np.eye(self.n_classes)[y_true]
        label_err = (probs - y_onehot) / m # (m, C)
        
        # 3. Prototype Gradient (dP)
        # d_score/dP = sin(omega * (X-P)) * omega
        # Because we are minimizing loss, we use the chain rule: 
        # dLoss/dP = label_err * d_score/dP
        sin_term = np.sin(self.omega[np.newaxis, :, np.newaxis] * diff)
        dP = (label_err[:, :, np.newaxis] * sin_term * self.omega[np.newaxis, :, np.newaxis]).sum(axis=0)
        
        # 4. Omega Gradient (dw)
        # d_score/dw = -(X-P) * sin(omega * (X-P))
        dw = (label_err[:, :, np.newaxis] * (-diff * sin_term)).sum(axis=(0, 2)) / m
        
        # Updates
        self.P -= self.lr_p * dP
        self.omega -= self.lr_w * dw
        
        # Clipping omega to prevent vanishing/exploding cycles
        self.omega = np.clip(self.omega, 0.1, 10.0)
        
        return scores, probs

    def train(self, X, y_true, iters=1000, batch=128, eval_every=100, eval_train_n=5000):
        history = []
        for t in range(1, iters + 1):
            idx = np.random.randint(0, X.shape[0], batch)
            Xb, yb = X[idx], y_true[idx]
            
            scores, probs = self.update(Xb, yb)
            
            loss = -np.log(probs[np.arange(batch), yb] + 1e-12).mean()
            
            if t % eval_every == 0 or t == 1:
                tr = self.score(X[:eval_train_n], y_true[:eval_train_n])
                te = self.score(X20, yt20)
                print(f"Step {t:4d} | Loss: {loss:.4f} | Train Acc: {tr:.4f} | Test Acc: {te:.4f}")
                
        return history

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

    clf = SinCosMatrixClassifier(lr_p=0.01, lr_w=0.001)

    print(f"Initial test acc: {clf.score(X20, yt20):.4f}")

    clf.train(X_train, y_train, iters=1000, batch=128, eval_every=100)

    print(f"\nFinal test acc: {clf.score(X20, yt20):.4f}")
    print(f"Final omega: {np.round(clf.omega, 4)}")