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]

# CCT-CC-SI MLP with Conditional Manifestation Oscillator
class CC_MLPClassifier:
    def __init__(self, input_size, hidden_size, output_size, 
                 collapse_threshold=0.5,  # Entropy threshold for M update
                 energy_budget=1.0,        # Energy per iteration
                 cycle_check=10):          # Periodicity detection window
        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))
        
        # CCT: Manifestation Matrix (Stationary Truth for each class)
        self.M = np.random.randn(output_size, hidden_size) * 0.01
        self.b = np.zeros((1, hidden_size))
        
        # CCT: Thresholds and Energy Budget
        self.collapse_threshold = collapse_threshold
        self.energy_budget = energy_budget
        self.cycle_check = cycle_check
        
        # CCT: Oscillation History (for periodicity detection)
        self.M_history = []
        self.error_history = []
        
    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
        self.output = self.softmax(self.z2)
        return self.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 -= 0.01 * dW1
        self.b1 -= 0.01 * db1
        self.W2 -= 0.01 * dW2
        self.b2 -= 0.01 * 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)
    
    # === CCT Extension: Conditional Manifestation Collapse ===
    
    def collapse_manifestation_M(self, X, yt, energy_weight=0.01):
        """
        Conditional Collapse of M based on entropy/error.
        
        CCT Principles:
        - Only update when error > collapse_threshold
        - Detect periodicity to avoid oscillation
        - Pay with energy only when necessary
        """
        activation = np.dot(X, self.W1) + self.b  # Current state
        err = self.M[yt] - activation  # Distance from truth
        
        # CCT: Compute entropy (error magnitude)
        entropy = np.mean(np.linalg.norm(err, axis=1))
        
        # CCT: Periodicity Detection
        self.error_history.append(entropy)
        if len(self.error_history) > self.cycle_check:
            self.error_history.pop(0)
            
        is_periodic = self._detect_oscillation()
        
        # CCT: Conditional Collapse Decision
        if entropy > self.collapse_threshold and not is_periodic:
            # Pay energy to collapse
            self.M[yt] -= energy_weight * err
            self.b += energy_weight * err.mean(axis=0)
            collapse_occurred = True
        else:
            # System is stable or oscillating - conserve energy
            collapse_occurred = False
            
        return entropy, collapse_occurred, is_periodic
    
    def _detect_oscillation(self):
        """Detect if M is oscillating (periodic behavior)."""
        if len(self.error_history) < self.cycle_check:
            return False
            
        # Check for periodic pattern in error history
        history = np.array(self.error_history)
        
        # Simple periodicity check: variance over window vs variance of differences
        variance = np.var(history)
        diff_variance = np.var(np.diff(history))
        
        # If variance is low but diff variance is high, we're oscillating
        if variance < 0.1 and diff_variance > 0.01:
            return True
        return False
    
    def update_M_adaptive(self, X, yt):
        """
        Adaptive M update using CCT energy efficiency.
        
        Strategy: 
        - If system is stable (low entropy): minimal update
        - If system is unstable (high entropy): full update
        - If system is oscillating: freeze M, let MLP adapt
        """
        activation = np.dot(X, self.W1) + self.b
        err = self.M[yt] - activation
        entropy = np.mean(np.linalg.norm(err, axis=1))

        self.error_history.append(entropy)
        if len(self.error_history) > self.cycle_check:
            self.error_history.pop(0)

        is_periodic = self._detect_oscillation()
        
        # CCT: Determine update magnitude based on entropy
        if is_periodic:
            # Oscillating regime - freeze M updates and let the MLP adapt.
            lr = 0.0
        elif entropy < 0.1:
            # Stable regime - minimal energy expenditure
            lr = 0.001
        elif entropy < self.collapse_threshold:
            # Medium - standard update
            lr = 0.01
        else:
            # High entropy - maximum collapse
            lr = 0.05

        # Apply update with adaptive learning rate.
        if lr > 0.0:
            self.M[yt] -= lr * err
            self.b += lr * err.mean(axis=0)

        return entropy, lr, is_periodic

# === CCT-CC-SI Training Loop ===

f = CC_MLPClassifier(
    input_size=784, 
    hidden_size=100, 
    output_size=10,
    collapse_threshold=0.8,   # CCT: Only collapse when error exceeds this
    energy_budget=1.0,
    cycle_check=20
)

# Track CCT metrics
total_collapse_ops = 0
total_energy_saved = 0
oscillation_freezes = 0

i = 0
while True:
    idx = np.random.randint(0, 60000, 100)
    X = X_train[idx]
    yt = y_train[idx]
    
    # === Standard MLP Backprop (always runs) ===
    f.forward(X)
    f.update(X, np.eye(10)[yt])
    
    # === CCT: Conditional M Collapse ===
    entropy, lr_used, is_periodic = f.update_M_adaptive(X, yt)
    
    # Track oscillation state
    if is_periodic:
        oscillation_freezes += 1
    
    # Score check
    score = f.score(X, yt)
    total_error = np.mean((f.M[yt] - (np.dot(X, f.W1) + f.b))**2)
    
    if i % 100 == 0:
        energy_status = "ACTIVE" if lr_used > 0.01 else "CONSERVED"
        osc_status = "OSCILLATING" if is_periodic else "STABLE"
        print(f"Step {i:4d} | Score: {score:.4f} | "
              f"Entropy: {entropy:.4f} | "
              f"LR: {lr_used:.3f} | "
              f"Energy: {energy_status} | "
              f"State: {osc_status}")
    
    # CCT: Early termination if stable and accurate
    if score > 0.98 and entropy < 0.1 and not is_periodic:
        print(f"\n✓ COLLAPSE ACHIEVED: Score {score:.4f}, Entropy {entropy:.4f}")
        break
        
    i += 1