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]

class SemanticMeshClassifier:
    def __init__(self, input_size, hidden_size, output_clusters, mesh_dim=10, learning_rate=0.01):
        """
        input_size: 784 (pixels)
        output_clusters: 10 (digits 0-9)
        mesh_dim: The size of the 2D Semantic Mesh for each class (e.g., 10x10)
        """
        self.mesh_dim = mesh_dim
        self.output_clusters = output_clusters
        
        # Standard Weights
        self.W1 = np.random.randn(input_size, hidden_size) * 0.01
        self.b1 = np.zeros((1, hidden_size))
        
        # The Mesh Head: Instead of outputting (1, 10), 
        # we output (1, 10 * mesh_dim * mesh_dim) to form a tensor field.
        self.W2 = np.random.randn(hidden_size, output_clusters * mesh_dim**2) * 0.01
        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):
        """
        Transforms the raw output into a set of 2D probability meshes.
        Each mesh must sum to 1.0 (conservation of truth).
        """
        # Reshape to (batch, clusters, mesh_dim, mesh_dim)
        batch_size = x.shape[0]
        meshes = x.reshape(batch_size, self.output_clusters, self.mesh_dim, self.mesh_dim)
        
        # Softmax across the mesh dimensions (i, j) for each cluster independently
        # This creates the 'Basins of Attraction' mentioned in the Codex
        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)
        
        return exp_meshes / sum_meshes

    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 is now a 4D Tensor: (batch, cluster, mesh_i, mesh_j)
        self.meshes = self.mesh_softmax(self.z2)
        return self.meshes

    def get_mesh_metrics(self, mesh):
        """
        Implements TELEPASM-M operators: Gradient and Curl.
        Simplified as discrete differences over the mesh.
        """
        # Discrete Gradient (Approximate)
        grad_i = np.diff(mesh, axis=-2)
        grad_j = np.diff(mesh, axis=-1)
        
        # Discrete Curl (Approximate vorticity: dPj/di - dPi/dj)
        # This detects "Paradoxes" in the digit's semantic representation
        curl = np.zeros_like(mesh)
        curl[..., :-1, :-1] = grad_j[..., :-1] - grad_i[..., :-1]
        
        return np.mean(np.abs(grad_i)), np.mean(np.abs(curl))

    def predict(self, X):
        """
        The 'Collapse' operator (FOLD).
        Collapses the 2D mesh into a 1D scalar by summing the energy of the mesh.
        """
        meshes = self.forward(X)
        # Fold: Sum the total probability mass for each cluster
        # In a real TELEPASM-M system, this is where the 'Attractor Basin' is measured
        collapsed = np.sum(meshes, axis=(2, 3)) 
        return np.argmax(collapsed, axis=1)

    def update(self, X, y_true):
        """
        Standard backprop adjusted for the mesh-tensor output.
        """
        m = X.shape[0]
        # Convert y_true to a target mesh (One-hot cluster, flat mesh distribution)
        # For simplicity, we target a 'peak' in the center of the mesh for the correct class
        y_pred_meshes = self.forward(X)
        
        # Flatten the meshes back to 2D for gradient calculation
        y_pred_flat = y_pred_meshes.reshape(m, -1)
        
        # Target: The correct cluster gets the 'truth mass', others get 0
        target_flat = np.zeros_like(y_pred_flat)
        for i, label in enumerate(y_true):
            start = label * self.mesh_dim**2
            # We distribute the truth across the mesh for the correct label
            target_flat[i, start : start + self.mesh_dim**2] = 1.0 / (self.mesh_dim**2)

        # Gradient of cross-entropy for the mesh
        dz2 = y_pred_flat - target_flat
        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) # ReLU derivative
        dW1 = np.dot(X.T, dz1) / m
        db1 = np.sum(dz1, axis=0, keepdims=True) / m

        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)

# Usage Example
if __name__ == "__main__":
    # Initialize the Mesh Classifier
    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]
        print(i, f.score(X,yt))
        f.update(X, yt)
        i += 1