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)
    X_test = read('../X_test.wav')[1].reshape(-1, 784)
    y_test = (read('../y_test.wav')[1] * 9).astype(int)
except Exception:
    # Fallback to simulated data if files are not present
    X_train = np.random.randn(60000, 784)
    y_train = np.random.randint(0, 10, 60000)
    X_test = np.random.randn(1000, 784)
    y_test = np.random.randint(0, 10, 1000)

X20 = X_test[:1000]
yt20 = y_test[:1000]

class MLPClassifier:
    def __init__(self, input_size, hidden_size, output_size, learning_rate=None):
        if learning_rate is None:
            self.learning_rate = np.random.rand(4) * 0.01 + 0.001
        else:
            self.learning_rate = learning_rate
            
        self.W1 = np.random.randn(input_size, hidden_size) * np.sqrt(2.0 / input_size)
        self.b1 = np.zeros((1, hidden_size))
        self.W2 = np.random.randn(hidden_size, output_size) * np.sqrt(2.0 / hidden_size)
        self.b2 = np.zeros((1, output_size))
        
        # CCT State Tracker for ESD History
        self.w2_esd_history = []
        self.target_esd = 5.0
        
    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) + 1e-15)
    
    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 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)
        
        # 1. Measure Current ESD States
        current_w1_esd = self.calculate_esd(self.W1)
        current_w2_esd = self.calculate_esd(self.W2)
        self.w2_esd_history.append(current_w2_esd)
        if len(self.w2_esd_history) > 10:
            self.w2_esd_history.pop(0)
            
        # 2. Extract the Error Signal from the Oscillation Trajectory
        # We calculate the deviation from the target precision scale (5.0)
        error_w2 = self.target_esd - current_w2_esd
        error_w1 = self.target_esd - current_w1_esd
        
        # Compute volatility/oscillation dampening component (derivative-like control)
        if len(self.w2_esd_history) > 2:
            oscillation_amplitude = np.abs(self.w2_esd_history[-1] - self.w2_esd_history[-2])
        else:
            oscillation_amplitude = 0.0
            
        # 3. Formulate the Stabilization Scaling Factor
        # If ESD drops below 5.0 (positive error), it scales down the updates to compress the range
        # If it rises above 5.0, it allows larger updates to scale out minimum elements safely
        dampening_factor_w2 = np.exp(-0.5 * error_w2 - 0.2 * oscillation_amplitude)
        #dampening_factor_w1 = np.exp(-0.3 * error_w1)
        dampening_factor_w1 = np.exp(-0.5 * error_w1)
        
        norm_dW1 = np.linalg.norm(dW1) + 1e-15
        norm_dW2 = np.linalg.norm(dW2) + 1e-15
        
        # Base updates combined with adaptive CCT field constraints
        step1 = (self.learning_rate[0] * dampening_factor_w1) / (1.0 + np.log1p(norm_dW1))
        step2 = (self.learning_rate[1] * dampening_factor_w1) / (1.0 + np.log1p(norm_dW1))
        step3 = (self.learning_rate[2] * dampening_factor_w2) / (1.0 + np.log1p(norm_dW2))
        step4 = (self.learning_rate[3] * dampening_factor_w2) / (1.0 + np.log1p(norm_dW2))
        
        self.W1 -= step1 * dW1
        self.b1 -= step2 * db1
        self.W2 -= step3 * dW2
        self.b2 -= step4 * 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 calculate_esd(self, matrix):
        abs_mat = np.abs(matrix)
        max_val = np.max(abs_mat)
        min_val = np.min(abs_mat[abs_mat > 0]) if np.any(abs_mat > 0) else 1e-15
        return np.log10(max_val / min_val)

# Initialize and train
#learning_rate = np.array([0.05, 0.05, 0.05, 0.05])
learning_rate = np.random.rand(4)
f = MLPClassifier(input_size=784, hidden_size=128, output_size=10, learning_rate=learning_rate)

i = 0
print("Executing CCT Meta-Entropy Stabilized Field Engine...")
while True:
    idx = np.random.randint(0, X_train.shape[0], 128)
    X = X_train[idx]
    yt = y_train[idx]
    
    if i % 100 == 0:
        train_acc = f.score(X, yt)
        test_acc = f.score(X20, yt20)
        esd_w1 = f.calculate_esd(f.W1)
        esd_w2 = f.calculate_esd(f.W2)
        print(f"Iteration: {i:03d} | Train Acc: {train_acc:.4f} | Test Acc: {test_acc:.4f} | W1 ESD: {esd_w1:.2f} | W2 ESD: {esd_w2:.2f}")
        
    f.update(X, np.eye(10)[yt])
    i += 1