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

# Load training and test data
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]

# Define the MLP Classifier
class MLPClassifier:
    def __init__(self, input_size, hidden_size, output_size, learning_rate=np.random.rand(4)):
        self.W1 = np.random.randn(input_size, hidden_size) * 0.01
        self.b1 = np.zeros((1, hidden_size))
        self.W2 = np.random.randn(hidden_size, output_size) * 0.01
        self.b2 = np.zeros((1, output_size))
        self.learning_rate = learning_rate
    
    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):
        self.z1 = np.dot(X, self.W1) + self.b1
        self.a1 = self.relu(self.z1)
        self.z2 = np.dot(self.a1, self.W2) + self.b2
        output = self.softmax(self.z2)
        return output
    
    def compute_loss(self, y_true, y_pred):
        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, y_pred_conditional=None, alpha=0.5):
        """
        Backward pass supporting combined ordinary and conditional probabilities.
        
        Args:
            X: Input data
            y_true: True labels (one-hot encoded)
            y_pred: Ordinary probability P(y|x)
            y_pred_conditional: Conditional probability P(y|x,t) or None
            alpha: Weight for conditional probability (0=only ordinary, 1=only conditional)
                   When y_pred_conditional is None, alpha is ignored
        
        Returns:
            Gradients dW1, db1, dW2, db2
        """
        m = y_true.shape[0]
        
        if y_pred_conditional is not None:
            # Combine probabilities: P_combined = (1-alpha) * P(y|x) + alpha * P(y|x,t)
            # The gradient becomes: dz2 = (1-alpha) * (y_pred - y_true) + alpha * (y_pred_conditional - y_true)
            dz2_ordinary = y_pred - y_true
            dz2_conditional = y_pred_conditional - y_true
            dz2 = (1 - alpha) * dz2_ordinary + alpha * dz2_conditional
        else:
            # Original behavior: only ordinary probability
            dz2 = y_pred - y_true
        
        dW2 = np.dot(self.a1.T, dz2) / m
        db2 = np.sum(dz2, axis=0, keepdims=True) / m
        da1 = np.dot(dz2, self.W2.T)
        dz1 = da1 * self.relu_derivative(self.z1)
        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, thinking_time=None, alpha=0.5):
        """
        Update weights using combined ordinary and conditional probabilities.
        
        Args:
            X: Input data
            y_true: True labels
            thinking_time: Thinking time for conditional probability (None to skip)
            alpha: Weight for conditional probability (only used if thinking_time provided)
        """
        y_pred = self.forward(X)
        
        if thinking_time is not None:
            y_pred_conditional = self.forward_conditional(X, thinking_time)
            dW1, db1, dW2, db2 = self.backward(X, y_true, y_pred, y_pred_conditional, alpha)
        else:
            dW1, db1, dW2, db2 = self.backward(X, y_true, y_pred)
        
        self.W1 -= self.learning_rate[0] * dW1
        self.b1 -= self.learning_rate[1] * db1
        self.W2 -= self.learning_rate[2] * dW2
        self.b2 -= self.learning_rate[3] * db2
    
    def predict(self, X):
        probabilities = self.forward(X)
        return np.argmax(probabilities, axis=1)

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

    def forward_conditional(self, X, thinking_time):
        """
        Calculate forward pass with conditional probability given thinking time.
        P(y|x, t) - probability of class y given input x and thinking time t
        
        Different samples can take different paths based on their thinking time.
        thinking_time can be a scalar (same for all) or array (per-sample)
        """
        if isinstance(thinking_time, (int, float)):
            thinking_time = np.full(X.shape[0], thinking_time)
        
        # Store for potential backward pass
        self.z1 = np.dot(X, self.W1) + self.b1
        self.a1 = self.relu(self.z1)
        
        # Apply thinking time modulation to hidden layer
        # Longer thinking time = more refined representations
        # Modulate activations based on thinking time
        t_normalized = thinking_time.reshape(-1, 1) / 100.0  # Normalize by expected max iterations
        modulation = np.minimum(t_normalized, 1.0)  # Cap at 1.0
        
        # Apply modulation to hidden activations
        a1_modulated = self.a1 * modulation
        
        # Second layer with modulated activations
        self.z2 = np.dot(a1_modulated, self.W2) + self.b2
        
        # Temperature scaling based on thinking time
        # Higher thinking time = lower temperature (more confident)
        temperature = 1.0 / (0.5 + 0.5 * modulation)
        z2_scaled = self.z2 * temperature
        
        output = self.softmax(z2_scaled)
        return output

    def conditional_probability(self, X, thinking_time, y_true=None):
        """
        Calculate P(y|x, t) - conditional probability given input and thinking time.
        Returns probabilities and optionally the loss.
        """
        probabilities = self.forward_conditional(X, thinking_time)
        
        if y_true is not None:
            loss = self.compute_loss(y_true, probabilities)
            return probabilities, loss
        
        return probabilities

# Initialize and train the MLP Classifier
learning_rate = np.random.rand(4)
f = MLPClassifier(input_size=784, hidden_size=100, output_size=10, learning_rate=learning_rate)
g = MLPClassifier(input_size=784, hidden_size=100, output_size=10, learning_rate=learning_rate)
h = MLPClassifier(input_size=784, hidden_size=100, output_size=10, learning_rate=learning_rate)
max_acc = 0
time = 0


alpha = 0.1  # Weight for conditional probability (0.3 = 30% conditional, 70% ordinary)
i = 0
while True:
    if i%300==0:
        idx1 = np.random.randint(0,60000,1000)
        idx2 = np.random.randint(0,60000,1000)
    X = X_train[idx1]
    yt = y_train[idx1]
    g.update(X, np.eye(10)[yt], thinking_time=50, alpha=alpha)
    X = X_train[idx2]
    yt = y_train[idx2]
    h.update(X, np.eye(10)[yt], thinking_time=50, alpha=alpha)
    
    f.W1 = 0.5 * (g.W1 + h.W1)
    
    idx = np.random.randint(0, 60000, 100)
    X = X_train[idx]
    yt = y_train[idx]        
    f.update(X, np.eye(10)[yt], thinking_time=50, alpha=alpha)
    g.W1 = f.W1.copy()
    h.W1 = f.W1.copy()
    max_acc = np.maximum(max_acc, f.score(X20, yt20))
    if i % 100 == 0:
        print(f"Iteration {i}: Score={max_acc:.4f}, alpha={alpha}", g.score(X_train[idx1],y_train[idx1]))
    i += 1