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]

# Normalize input data to prevent exploding activations from raw audio scales
X_train = X_train / (np.max(np.abs(X_train)) + 1e-9)
X_test = X_test / (np.max(np.abs(X_test)) + 1e-9)

class SmoothBioMetalClassifier:
    def __init__(self, input_size, shell_size, core_size, output_size):
        # Using a stable, calibrated learning rate instead of pure random numbers to prevent gradient explosion
        self.lr = 0.01 
        
        # 1. Surface to Outer Shell
        self.W_shell = np.random.randn(input_size, shell_size) * np.sqrt(2.0 / input_size)
        self.b_shell = np.zeros((1, shell_size))
        
        # 2. Outer Shell to Core 
        self.W_core = np.random.randn(shell_size, core_size) * np.sqrt(2.0 / shell_size)
        self.b_core = np.zeros((1, core_size))
        
        # 3. 10 Smooth Decimal Processing Streams
        self.W_streams = [np.random.randn(core_size, output_size) * 0.01 for _ in range(10)]
        self.b_streams = [np.zeros((1, output_size)) for _ in range(10)]
        
        # 4. Outward Radiation
        self.W_out = np.random.randn(output_size, output_size) * 0.01
        self.b_out = np.zeros((1, output_size))
    
    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 extract_smooth_stream(self, matrix, stream_idx):
        """
        Applies a continuous, differentiable scaling window.
        As stream_idx increases, it isolates finer decimal micro-variations.
        """
        # Centers the scale factor around the decimal threshold
        scale = 10 ** (stream_idx - 4) 
        return np.tanh(matrix * scale)

    def stream_derivative(self, matrix, stream_idx):
        """ The analytical derivative of tanh(matrix * 10^(idx-4)) """
        scale = 10 ** (stream_idx - 4)
        tanh_val = np.tanh(matrix * scale)
        return scale * (1.0 - tanh_val ** 2)

    def forward(self, X):
        # Step 1: Into the shell
        self.z_shell = np.dot(X, self.W_shell) + self.b_shell
        self.a_shell = self.relu(self.z_shell)
        
        # Step 2: Inward to Core
        self.z_core = np.dot(self.a_shell, self.W_core) + self.b_core
        self.a_core = self.relu(self.z_core)
        
        # Step 3: Diverge into 10 smooth decimal tolerance scales
        self.stream_outputs = []
        self.stream_activations = []
        combined_signal = np.zeros((X.shape[0], 10))
        
        for idx in range(10):
            a_str = self.extract_smooth_stream(self.a_core, idx)
            self.stream_activations.append(a_str)
            
            z_str = np.dot(a_str, self.W_streams[idx]) + self.b_streams[idx]
            self.stream_outputs.append(z_str)
            combined_signal += z_str
            
        # Step 4: Radiate back out
        self.z_out = np.dot(combined_signal, self.W_out) + self.b_out
        output = self.softmax(self.z_out)
        return output
    
    def backward(self, X, y_true, y_pred):
        m = y_true.shape[0]
        
        # Outward Layer Error
        dz_out = y_pred - y_true
        
        # Summed stream outputs for the outer layer weight update
        summed_streams = np.sum(self.stream_outputs, axis=0)
        dW_out = np.dot(summed_streams.T, dz_out) / m
        db_out = np.sum(dz_out, axis=0, keepdims=True) / m
        
        # Project error back through the Outer radiation matrix to the streams
        dz_streams = np.dot(dz_out, self.W_out.T)
        
        dW_streams_list = []
        db_streams_list = []
        da_core_accumulated = np.zeros_like(self.a_core)
        
        # Calculate gradients for each decimal scale stream
        for idx in range(10):
            dW_str = np.dot(self.stream_activations[idx].T, dz_streams) / m
            db_str = np.sum(dz_streams, axis=0, keepdims=True) / m
            dW_streams_list.append(dW_str)
            db_streams_list.append(db_str)
            
            # Key fix: Chain rule includes the smooth stream derivative
            da_stream_to_core = np.dot(dz_streams, self.W_streams[idx].T)
            da_core_accumulated += da_stream_to_core * self.stream_derivative(self.a_core, idx)
            
        # Backprop through Core
        dz_core = da_core_accumulated * self.relu_derivative(self.z_core)
        dW_core = np.dot(self.a_shell.T, dz_core) / m
        db_core = np.sum(dz_core, axis=0, keepdims=True) / m
        
        # Backprop through Shell
        da_shell = np.dot(dz_core, self.W_core.T)
        dz_shell = da_shell * self.relu_derivative(self.z_shell)
        dW_shell = np.dot(X.T, dz_shell) / m
        db_shell = np.sum(dz_shell, axis=0, keepdims=True) / m
        
        return dW_shell, db_shell, dW_core, db_core, dW_streams_list, db_streams_list, dW_out, db_out
    
    def update(self, X, y_true):
        y_pred = self.forward(X)
        dW_shell, db_shell, dW_core, db_core, dW_streams, db_streams, dW_out, db_out = self.backward(X, y_true, y_pred)
        
        self.W_shell -= self.lr * dW_shell
        self.b_shell -= self.lr * db_shell
        self.W_core  -= self.lr * dW_core
        self.b_core  -= self.lr * db_core
        
        for idx in range(10):
            self.W_streams[idx] -= self.lr * dW_streams[idx]
            self.b_streams[idx] -= self.lr * db_streams[idx]
            
        self.W_out -= self.lr * dW_out
        self.b_out -= self.lr * db_out

    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)

if __name__ == "__main__":
    f = SmoothBioMetalClassifier(input_size=784, shell_size=128, core_size=32, output_size=10)

    i = 0
    while True:
        idx = np.random.randint(0, X_train.shape[0], 100)
        X = X_train[idx]
        yt = y_train[idx]
        
        if i % 100 == 0:
            print(f"Iteration {i:5d} | Stable Collaborative Accuracy: {f.score(X, yt) * 100:.2f}%")
            
        f.update(X, np.eye(10)[yt])
        i += 1