import numpy as np
from tensorflow.keras.datasets import mnist

class LinearGraphicalClassifier:
    """
    Linear Graphical Classifier for MNIST.
    
    Architecture:
        scores = X @ P.T   (P acts as class prototypes)
        probs = softmax(scores)
    
    Dual‑gradient update:
        total_grad = (label_error.T @ X) + λ * [ (label_error.T @ X) - P * sum(label_error) ]
        where label_error = probs - y_onehot
    """
    def __init__(self, num_classes=10, input_size=784, lam=0.1, lr=0.01):
        self.num_classes = num_classes
        self.input_size = input_size
        self.lam = lam          # weight of the graphical pixel error term
        self.lr = lr            # learning rate

        # Prototypes initialised uniformly in [0,1] – same image space as data
        self.P = np.random.uniform(low=0.0, high=1.0,
                                   size=(num_classes, input_size))
        
    def softmax(self, x):
        """Stable softmax along axis=1."""
        e_x = np.exp(x - np.max(x, axis=1, keepdims=True))
        return e_x / e_x.sum(axis=1, keepdims=True)

    def forward(self, X):
        """Compute class scores and probabilities."""
        scores = X @ self.P.T               # (N, C)
        probs = self.softmax(scores)
        return scores, probs

    def predict(self, X):
        """Return predicted class indices."""
        scores = X @ self.P.T
        return np.argmax(scores, axis=1)

    def cross_entropy_loss(self, probs, y_onehot):
        """Average cross‑entropy loss (with small epsilon for safety)."""
        eps = 1e-12
        return -np.mean(np.sum(y_onehot * np.log(probs + eps), axis=1))

    def update(self, X, y):
        """
        Perform one batch update using the dual‑gradient rule.
        
        Parameters:
            X: batch of images, shape (N, D)
            y: true labels, shape (N,) with integers 0..9
        Returns:
            loss: average cross‑entropy loss for the batch
        """
        N = X.shape[0]
        # One‑hot encode labels
        y_onehot = np.zeros((N, self.num_classes))
        y_onehot[np.arange(N), y] = 1.0

        # Forward pass
        scores, probs = self.forward(X)
        loss = self.cross_entropy_loss(probs, y_onehot)

        # Label error: (probs - y_true)
        label_error = probs - y_onehot          # (N, C)

        # 1. Standard linear gradient: d(cross‑entropy)/dP = label_error.T @ X
        grad_linear = label_error.T @ X         # (C, D)

        # 2. Graphical pixel error term
        #    For class j: Σ_i label_error[i,j] * (X_i - P_j)
        #    = (label_error.T @ X)[j] - P_j * Σ_i label_error[i,j]
        sum_error_per_class = label_error.sum(axis=0)   # (C,)
        #grad_pixel = grad_linear - label_error.T @ self.P[y]
        grad_pixel = -label_error.T @ self.P[y]
        #grad_pixel = grad_linear - self.P * sum_error_per_class[:, np.newaxis]

        # Combined update with λ weighting
        total_grad = grad_linear + self.lam * grad_pixel

        # SGD step
        self.P -= self.lr * total_grad

        return loss


# ----------------------------------------------------------------------
# Training and evaluation
# ----------------------------------------------------------------------
def main():
    # 1. Load and preprocess MNIST
    (x_train, y_train), (x_test, y_test) = mnist.load_data()
    x_train = x_train.reshape(-1, 784).astype(np.float32) / 255.0   # [0,1]
    x_test  = x_test.reshape(-1, 784).astype(np.float32) / 255.0

    # 2. Hyperparameters
    num_classes = 10
    input_size = 784
    lam = 0.01          # weight for graphical pixel error term
    lr = 0.01
    batch_size = 100
    epochs = 2000

    # 3. Instantiate classifier
    model = LinearGraphicalClassifier(num_classes=num_classes,
                                      input_size=input_size,
                                      lam=lam, lr=lr)

    # 4. Training loop
    n_train = x_train.shape[0]
    for epoch in range(epochs):
        # Shuffle training data each epoch
        perm = np.random.permutation(n_train)
        x_train_shuf = x_train[perm]
        y_train_shuf = y_train[perm]

        total_loss = 0.0
        n_batches = 0
        for i in range(0, n_train, batch_size):
            X_batch = x_train_shuf[i:i+batch_size]
            y_batch = y_train_shuf[i:i+batch_size]
            loss = model.update(X_batch, y_batch)
            total_loss += loss
            n_batches += 1

        # Compute accuracies
        train_pred = model.predict(x_train)
        test_pred  = model.predict(x_test)
        train_acc = np.mean(train_pred == y_train)
        test_acc  = np.mean(test_pred == y_test)

        avg_loss = total_loss / n_batches
        print(f"Epoch {epoch+1:2d}/{epochs} | Loss: {avg_loss:.4f} | "
              f"Train Acc: {train_acc:.4f} | Test Acc: {test_acc:.4f}")

    print("\nTraining finished.")

if __name__ == "__main__":
    main()