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

# Add parent directory to import flux_matrix
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from flux_matrix import FluxMatrix

# Load training and test data
X_train = read('../X_train.wav')[1].reshape(-1, 784).astype(np.float64)
y_train = (read('../y_train.wav')[1] * 9).astype(int)
X_test = read('../X_test.wav')[1].reshape(-1, 784).astype(np.float64)
y_test = (read('../y_test.wav')[1] * 9).astype(int)
X20 = X_test[:1000]
yt20 = y_test[:1000]


class FluxMLPClassifier:
    """
    MLP Classifier using Flux Algebra for uncertainty-aware learning.
    
    All weights and biases are FluxMatrices that track:
    - V: Current weight values
    - S: Uncertainty in each weight
    - T: Rate of change (flux/velocity)
    
    Learning = Entropy Collapse through Work payment
    """
    
    def __init__(self, input_size, hidden_size, output_size, 
                 learning_rate=0.01, init_entropy=0.01, work_budget=0.001):
        # Initialize weights as FluxMatrices
        # He initialization for values, small entropy, tiny flux
        scale1 = np.sqrt(2.0 / input_size)
        scale2 = np.sqrt(2.0 / hidden_size)
        
        self.W1 = FluxMatrix.from_random(
            (input_size, hidden_size),
            v_scale=scale1,
            s_init=init_entropy,
            t_scale=scale1 * 0.01
        )
        self.b1 = FluxMatrix.from_random(
            (1, hidden_size),
            v_scale=0.0,
            s_init=init_entropy,
            t_scale=0.001
        )
        self.W2 = FluxMatrix.from_random(
            (hidden_size, output_size),
            v_scale=scale2,
            s_init=init_entropy,
            t_scale=scale2 * 0.01
        )
        self.b2 = FluxMatrix.from_random(
            (1, output_size),
            v_scale=0.0,
            s_init=init_entropy,
            t_scale=0.001
        )
        
        self.learning_rate = learning_rate
        self.work_budget = work_budget
        
        # Track learning statistics
        self.total_work_paid = 0.0
        self.entropy_history = []
    
    def relu(self, x):
        return np.maximum(0, x)

    def relu_derivative(self, x):
        return np.where(x > 0, 1, 0)

    def softmax(self, x):
        exp_x = np.exp(x - np.max(x, axis=1, keepdims=True))
        return exp_x / np.sum(exp_x, axis=1, keepdims=True)

    def forward(self, X):
        """
        Forward pass with uncertainty propagation.
        
        Returns:
            output: Class probabilities (shape: batch x output_size)
            uncertainty: Uncertainty in predictions (shape: batch x output_size)
        """
        # Layer 1: X @ W1 + b1
        # Flux matrix multiplication propagates uncertainty
        z1_V = X @ self.W1.V + self.b1.V
        z1_S = np.sqrt((X**2 @ self.W1.S**2) + self.b1.S**2)
        z1_T = X @ self.W1.T + self.b1.T  # Flux propagates through linear transform
        
        self.z1 = FluxMatrix(z1_V, z1_S, z1_T)
        self.a1 = FluxMatrix(
            self.relu(z1_V),
            z1_S,  # Entropy passes through relu
            self.relu(z1_T)  # Flux passes through relu (gated by activation)
        )
        
        # Layer 2: a1 @ W2 + b2
        z2_V = self.a1.V @ self.W2.V + self.b2.V
        # Entropy propagation: σ² = a1² @ W2.S² + a1.S² @ W2² + b2.S²
        z2_S = np.sqrt((self.a1.V**2 @ self.W2.S**2) + 
                       (self.a1.S**2 @ self.W2.V**2) + 
                       self.b2.S**2)
        z2_T = self.a1.V @ self.W2.T + self.a1.T @ self.W2.V + self.b2.T
        
        self.z2 = FluxMatrix(z2_V, z2_S, z2_T)
        output = self.softmax(z2_V)
        
        # Uncertainty in predictions
        uncertainty = self.softmax(z2_V + z2_S) - self.softmax(z2_V - z2_S)
        
        return output, uncertainty

    def compute_loss(self, y_true, y_pred):
        """Cross-entropy loss."""
        m = y_true.shape[0]
        loss = -np.sum(y_true * np.log(y_pred + 1e-9)) / m
        return loss

    def backward(self, X, y_true, y_pred):
        """
        Backward pass - computes gradients (standard approach).
        
        Returns:
            Gradients for all weights and biases as standard matrices
        """
        m = y_true.shape[0]
        
        # Output layer gradients
        dz2 = y_pred - y_true
        dW2 = np.dot(self.a1.V.T, dz2) / m
        db2 = np.sum(dz2, axis=0, keepdims=True) / m
        
        # Hidden layer gradients
        da1 = np.dot(dz2, self.W2.V.T)
        dz1 = da1 * self.relu_derivative(self.z1.V)
        dW1 = np.dot(X.T, dz1) / m
        db1 = np.sum(dz1, axis=0, keepdims=True) / m
        
        return dW1, db1, dW2, db2

    def update(self, X, y_true):
        """
        Flux Algebra update: gradient step + entropy collapse.
        
        This is the key difference from standard MLP:
        1. Compute gradients (backward pass)
        2. Apply gradient step using Flux algebra
        3. Collapse entropy by paying Work
        """
        # Forward pass
        y_pred, uncertainty = self.forward(X)
        
        # Backward pass - get gradients
        dW1, db1, dW2, db2 = self.backward(X, y_true, y_pred)
        
        # Flux gradient steps
        # Each weight update: θ_{t+1} = θ_t ⊕ (-lr * gradient)
        # Then collapse: C(θ, W) to reduce uncertainty
        self.W1 = self.W1.gradient_step(dW1, self.learning_rate, self.work_budget)
        self.b1 = self.b1.gradient_step(db1, self.learning_rate, self.work_budget)
        self.W2 = self.W2.gradient_step(dW2, self.learning_rate, self.work_budget)
        self.b2 = self.b2.gradient_step(db2, self.learning_rate, self.work_budget)
        
        # Track work paid
        self.total_work_paid += self.work_budget * 4  # 4 weight matrices
        
        # Record entropy
        total_entropy = self.W1.total_entropy() + self.W2.total_entropy()
        self.entropy_history.append(total_entropy)
        
        return y_pred, uncertainty

    def predict(self, X):
        """Predict classes."""
        probabilities, uncertainty = self.forward(X)
        return np.argmax(probabilities, axis=1)

    def predict_with_confidence(self, X):
        """
        Predict classes with confidence scores.
        
        Returns:
            predictions: Class predictions
            confidences: Confidence in each prediction [0, 1]
        """
        probabilities, uncertainty = self.forward(X)
        predictions = np.argmax(probabilities, axis=1)
        
        # Confidence = 1 - normalized uncertainty
        max_prob = np.max(probabilities, axis=1, keepdims=True)
        confidence = np.exp(-uncertainty.sum(axis=1) / probabilities.shape[1])
        
        return predictions, confidence

    def score(self, X, y_true):
        """Compute accuracy."""
        y = self.predict(X)
        return np.mean(y == y_true)
    
    def score_with_uncertainty(self, X, y_true):
        """
        Compute accuracy with uncertainty breakdown.
        
        Returns:
            accuracy: Overall accuracy
            high_conf_acc: Accuracy on high-confidence predictions
            low_conf_acc: Accuracy on low-confidence predictions
        """
        predictions, confidences = self.predict_with_confidence(X)
        correct = predictions == y_true
        
        accuracy = np.mean(correct)
        
        # Split by confidence
        median_conf = np.median(confidences)
        high_conf_mask = confidences >= median_conf
        low_conf_mask = confidences < median_conf
        
        high_conf_acc = np.mean(correct[high_conf_mask]) if high_conf_mask.any() else 0.0
        low_conf_acc = np.mean(correct[low_conf_mask]) if low_conf_mask.any() else 0.0
        
        return accuracy, high_conf_acc, low_conf_acc
    
    def get_total_entropy(self):
        """Get total model uncertainty."""
        return (self.W1.total_entropy() + self.b1.total_entropy() + 
                self.W2.total_entropy() + self.b2.total_entropy())
    
    def get_stability_score(self):
        """Get overall model stability [0, 1]."""
        entropy = self.get_total_entropy()
        return np.exp(-entropy)


# ──────────────────────────────────────────────────────────────
# TRAINING
# ──────────────────────────────────────────────────────────────

if __name__ == "__main__":
    print("=" * 60)
    print("FLUX ALGEBRA MLP CLASSIFIER")
    print("Uncertainty-Aware Neural Network Training")
    print("=" * 60)
    
    # Initialize Flux MLP
    f = FluxMLPClassifier(
        input_size=784,
        hidden_size=100,
        output_size=10,
        learning_rate=0.01,
        init_entropy=0.01,
        work_budget=0.0005
    )
    
    print(f"\nInitial state:")
    print(f"  Total entropy: {f.get_total_entropy():.4f}")
    print(f"  Stability: {f.get_stability_score():.4f}")
    
    i = 0
    batch_size = 100
    print_freq = 100
    
    print(f"\nTraining (printing every {print_freq} iterations)...")
    print(f"{'Iter':<8} {'Accuracy':<12} {'Entropy':<12} {'Stability':<12} {'Work Paid':<12}")
    print("-" * 60)
    
    while True:
        idx = np.random.randint(0, 60000, batch_size)
        X = X_train[idx]
        yt = y_train[idx]
        
        # Flux update
        y_pred, uncertainty = f.update(X, np.eye(10)[yt])
        y_pred, uncertainty = f.update(X, np.eye(10)[yt])
        y_pred, uncertainty = f.update(X, np.eye(10)[yt])
        y_pred, uncertainty = f.update(X, np.eye(10)[yt])
        
        if i % print_freq == 0:
            acc = f.score(X, yt)
            entropy = f.get_total_entropy()
            stability = f.get_stability_score()
            
            print(f"{i:<8} {acc:<12.4f} {entropy:<12.4f} {stability:<12.4f} {f.total_work_paid:<12.4f}")
            
            # Check if model has collapsed to stable state
            if entropy < 0.1 and i > 1000:
                print(f"\nModel stabilized after {i} iterations!")
                break
        
        i += 1
        
        # Safety break
        if i > 100000:
            print(f"\nReached maximum iterations ({i})")
            break
    
    # Final evaluation
    print("\n" + "=" * 60)
    print("FINAL EVALUATION")
    print("=" * 60)
    
    test_acc = f.score(X20, yt20)
    print(f"Test accuracy: {test_acc:.4f}")
    
    # Accuracy with confidence breakdown
    acc, high_acc, low_acc = f.score_with_uncertainty(X20, yt20)
    print(f"High-confidence accuracy: {high_acc:.4f}")
    print(f"Low-confidence accuracy: {low_acc:.4f}")
    print(f"Final entropy: {f.get_total_entropy():.4f}")
    print(f"Final stability: {f.get_stability_score():.4f}")
    
    print("\n" + "=" * 60)
    print("Flux Algebra enables uncertainty-aware learning!")
    print("=" * 60)