import numpy as np
from scipy.io.wavfile import read
from sklearn.cluster import KMeans

# 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):
        m = y_true.shape[0]
        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):
        y_pred = self.forward(X)
        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)

# Initialize and train the MLP Classifier
from time import sleep
learning_rate = np.random.rand(4)
import numpy as np
from time import sleep

# Initialize models
f = MLPClassifier(input_size=784, hidden_size=100, output_size=10, learning_rate=learning_rate)
g = KMeans(n_clusters=20, random_state=42, n_init=10)

# Fit KMeans on a subset first
print("Fitting KMeans on training data...")
g.fit(X_train[:5000])  # Use more samples for better centroids

# Pre-compute cluster-to-class mappings
print("Computing cluster-class distributions...")
cluster_class_dist = np.zeros((20, 10))
for i in range(len(X_train)):
    cluster_id = g.predict(X_train[i:i+1])[0]
    cluster_class_dist[cluster_id, y_train[i]] += 1

# Normalize to get probability distributions
cluster_probs = cluster_class_dist / (cluster_class_dist.sum(axis=1, keepdims=True) + 1e-8)

# Find pairs of clusters with different dominant classes
def find_distinct_pairs(probs, threshold=0.3):
    """Find cluster pairs where dominant classes differ significantly."""
    pairs = []
    for i in range(len(probs)):
        for j in range(i+1, len(probs)):
            # Check if classes with highest probability differ
            dominant_i = np.argmax(probs[i])
            dominant_j = np.argmax(probs[j])
            if dominant_i != dominant_j:
                # Check separation quality (both clusters are somewhat pure)
                purity = max(probs[i].max(), probs[j].max())
                if purity > threshold:
                    pairs.append((i, j, purity))
    return sorted(pairs, key=lambda x: -x[2])  # Sort by purity descending

distinct_pairs = find_distinct_pairs(cluster_probs)
print(f"Found {len(distinct_pairs)} highly separable cluster pairs")

# Training loop
best_score = 0
patience = 0
max_patience = 50  # Stop if no improvement for 50 iterations

print("\nStarting training...")
for i in range(10000):
    # Random batch without replacement cycling
    start_idx = (i * 1000) % 60000
    idx = np.arange(start_idx, start_idx + 1000) % 60000
    X = X_train[idx]
    yt = y_train[idx]
    
    # Predict cluster assignments for batch
    ids = g.predict(X)
    
    # Update on all distinct pairs found in this batch
    updates_made = 0
    for c1, c2, _ in distinct_pairs:
        mask1 = ids == c1
        mask2 = ids == c2
        n1, n2 = mask1.sum(), mask2.sum()
        
        if n1 >= 10 and n2 >= 10:
            # Combine samples from both clusters
            combined_mask = mask1 | mask2  # Fixed: use bitwise OR
            X_batch = X[combined_mask]
            y_batch = yt[combined_mask]
            
            # One-hot encode targets
            y_onehot = np.eye(10)[y_batch]
            
            # Update MLP
            f.update(X_batch, y_onehot)
            updates_made += 1
    
    # Progress output every 100 iterations
    if i % 100 == 0:
        current_score = f.score(X_train[:5000], y_train[:5000])
        print(f"Iter {i:5d} | Score: {current_score:.4f} | Updates: {updates_made}")
        
        # Early stopping check
        if current_score > best_score + 0.001:
            best_score = current_score
            patience = 0
        else:
            patience += 1
            if patience >= max_patience:
                print(f"\nConverged at iteration {i}")
                break

print(f"\nFinal score: {f.score(X_test, y_test):.4f}")