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

# --- Data Loading ---
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:
    print("Files not found, using mock data...")
    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)

X_train /= 255.0
X_test /= 255.0
X20, yt20 = X_test[:1000], y_test[:1000]

class LinearGraphicalClassifier:
    """
    Graphical classifier using linear projections (inner products).
    
    Instead of Sin/Cos flux, it uses a linear score: score = X @ P.T
    It maintains the graphical error signals:
    1. Label Error: (probs - y_true) -> directs which prototype to move.
    2. Pixel Error: (X - P) -> directs the direction of movement in image space.
    """
    def __init__(self, n_classes=10, input_dim=784, lr_p=0.01, lambda_pixel=0.1):
        self.n_classes = n_classes
        self.input_dim = input_dim
        self.lr_p = lr_p
        self.lambda_pixel = lambda_pixel
        
        rng = np.random.default_rng(0)
        # Prototypes initialized as random "images"
        self.P = rng.uniform(0, 1, (n_classes, input_dim))

    @staticmethod
    def _softmax(s):
        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):
        # Linear Layer: X (m, D) dot P.T (D, C) -> (m, C)
        return X @ self.P.T

    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
        scores = self._scores(X)
        probs = self._softmax(scores)
        
        # 2. Label Error (m, C)
        y_onehot = np.eye(self.n_classes)[y_true]
        label_err = (probs - y_onehot) / m 
        
        # 3. Graphical Error Signal
        # The 'Linear' gradient is simply the input X.
        # But to keep it "Graphical", we combine the linear gradient 
        # with the "Pixel Error" (X - P).
        
        # Standard Linear Gradient: sum(label_err * X)
        # (m, C).T @ (m, D) -> (C, D)
        grad_linear = label_err.T @ X
        
        # Graphical Pixel Correction:
        # For each class, we look at the samples attributed to it 
        # and pull the prototype toward the actual pixels.
        diff = X[:, np.newaxis, :] - self.P[np.newaxis, :, :] # (m, C, D)
        # We weight the pixel difference by the label error
        grad_pixel = (label_err[:, :, np.newaxis] * diff).sum(axis=0) 
        
        # Total weight update
        # We subtract the linear gradient and add the weighted pixel error
        total_dP = grad_linear + self.lambda_pixel * grad_pixel
        
        self.P -= self.lr_p * total_dP
        
        return scores, probs

    def train(self, X, y_true, iters=1000, batch=128, eval_every=100):
        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)
            
            if t % eval_every == 0 or t == 1:
                tr = self.score(X[:5000], y_true[:5000])
                te = self.score(X20, yt20)
                print(f"Step {t:4d} | Train Acc: {tr:.4f} | Test Acc: {te:.4f}")

if __name__ == "__main__":
    np.random.seed(0)
    print("=" * 64)
    print("Linear Graphical Classifier (Linear layers + Graphical Error)")
    print("=" * 64)

    clf = LinearGraphicalClassifier(lr_p=0.01, lambda_pixel=0.5)
    clf.train(X_train, y_train, iters=1000, batch=128, eval_every=100)
    print(f"\nFinal test acc: {clf.score(X20, yt20):.4f}")