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

# Load training and test data
try:
    X_train = read('../../X_train.wav')[1].reshape(-1, 784)
    y_train = (read('../../y_train.wav')[1] * 9).astype(int)
except:
    X_train = np.random.randn(60000, 784)
    y_train = np.random.randint(0, 10, 60000)

class SemanticMeshClassifier:
    def __init__(self, input_size, hidden_size, output_clusters, mesh_dim=10, learning_rate=0.01):
        self.mesh_dim = mesh_dim
        self.output_clusters = output_clusters
        
        # He Initialization
        self.W1 = np.random.randn(input_size, hidden_size) * np.sqrt(2. / input_size)
        self.b1 = np.zeros((1, hidden_size))
        self.W2 = np.random.randn(hidden_size, output_clusters * mesh_dim**2) * np.sqrt(2. / hidden_size)
        self.b2 = np.zeros((1, output_clusters * mesh_dim**2))
        
        self.lr = learning_rate

    def relu(self, x):
        return np.maximum(0, x)

    def mesh_softmax(self, x):
        batch_size = x.shape[0]
        meshes = x.reshape(batch_size, self.output_clusters, self.mesh_dim, self.mesh_dim)
        shift_meshes = meshes - np.max(meshes, axis=(2, 3), keepdims=True)
        exp_meshes = np.exp(shift_meshes)
        sum_meshes = np.sum(exp_meshes, axis=(2, 3), keepdims=True)
        res = exp_meshes / (sum_meshes + 1e-9)
        return res

    def forward(self, X):
        self.z1 = np.dot(X, self.W1) + self.b1
        # LOG 1: Check if input signal survives W1
        print(f"  [LOG] z1 mean: {np.mean(self.z1):.6f} | z1 std: {np.std(self.z1):.6f}")
        
        self.a1 = self.relu(self.z1)
        # LOG 2: Check for dead neurons (ReLU)
        print(f"  [LOG] a1 active%: {100 * np.mean(self.a1 > 0):.2f}%")
        
        self.z2 = np.dot(self.a1, self.W2) + self.b2
        # LOG 3: Check raw energy distribution before softmax
        print(f"  [LOG] z2 mean: {np.mean(self.z2):.6f} | z2 std: {np.std(self.z2):.6f}")
        
        self.meshes = self.mesh_softmax(self.z2)
        return self.meshes

    def predict(self, X):
        meshes = self.forward(X)
        collapsed = np.max(meshes, axis=(2, 3)) 
        preds = np.argmax(collapsed, axis=1)
        # LOG 4: Check if model is predicting the same class for everyone
        unique, counts = np.unique(preds, return_counts=True)
        print(f"  [LOG] Pred Distrib: {dict(zip(unique, counts))}")
        return preds

    def update(self, X, y_true):
        m = X.shape[0]
        y_pred_meshes = self.forward(X)
        y_pred_flat = y_pred_meshes.reshape(m, -1)
        
        target_flat = np.zeros_like(y_pred_flat)
        c = self.mesh_dim // 2
        for i, label in enumerate(y_true):
            idx_center = label * (self.mesh_dim**2) + (c * self.mesh_dim + c)
            target_flat[i, idx_center] = 1.0

        dz2 = y_pred_flat - target_flat
        # LOG 5: Check if gradient is zero
        print(f"  [LOG] dz2 mean: {np.mean(np.abs(dz2)):.6f}")
        
        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.z1 > 0) 
        dW1 = np.dot(X.T, dz1) / m
        db1 = np.sum(dz1, axis=0, keepdims=True) / m
        
        # LOG 6: Check if weight updates are actually happening
        print(f"  [LOG] dW1 mean: {np.mean(np.abs(dW1)):.8f} | dW2 mean: {np.mean(np.abs(dW2)):.8f}")

        self.W1 -= self.lr[0] * dW1
        self.b1 -= self.lr[1] * db1
        self.W2 -= self.lr[2] * dW2
        self.b2 -= self.lr[3] * db2

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

if __name__ == "__main__":
    # KEEPING ORIGINAL LEARNING RATE LOGIC
    learning_rate = np.random.rand(4) * 0.1
    f = SemanticMeshClassifier(input_size=784, hidden_size=100, output_clusters=10, learning_rate=learning_rate)
    
    i = 0
    while True:
        idx = np.random.randint(0, 60000, 100)
        X = X_train[idx]
        yt = y_train[idx]
        
        # Normalization to prevent saturation
        x_max = np.max(X)
        X_norm = X / x_max if x_max > 1 else X
        
        print(f"Iteration {i}:")
        acc = f.score(X_norm, yt)
        print(f"Accuracy: {acc:.4f}")
        
        f.update(X_norm, yt)
        print("-" * 30)
        i += 1