"""
Resonant Training Framework (RTF) - MNIST Prototype
====================================================
A PyTorch implementation of CCT-ODE-Resonance training

Key Features:
1. Internal layer oscillators with phase/frequency/amplitude
2. Cross-AI resonance during distributed training
3. FFT spectrum analysis of training dynamics
4. Comparison with standard training

Author: CC-SI Framework
"""

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch.utils.data import DataLoader, Subset
from torchvision import datasets, transforms
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation
from IPython.display import HTML
import warnings
warnings.filterwarnings('ignore')

# Set seeds for reproducibility
torch.manual_seed(42)
np.random.seed(42)

print("=" * 70)
print("RESONANT TRAINING FRAMEWORK - MNIST PROTOTYPE")
print("=" * 70)

# ============================================================
# PART 1: RESONANT NEURAL NETWORK LAYER
# ============================================================

class HarmonicOscillator:
    """Each parameter is a harmonic oscillator with phase, frequency, amplitude"""
    
    def __init__(self, initial_value, device='cpu'):
        self.device = device
        
        # Core parameter
        self.value = initial_value.clone().detach().to(device)
        
        # Oscillator parameters
        self.phase = torch.rand((), device=device) * 2 * np.pi
        self.frequency = torch.tensor(1.0, device=device)
        self.amplitude = torch.tensor(1.0, device=device)
        
        # Statistics
        self.gradient_history = []
        self.phase_history = []
        
    def state(self, t):
        """Oscillator state at time t: A * sin(ωt + φ)"""
        state = self.amplitude * torch.sin(self.frequency * t + self.phase)
        return state.reshape_as(self.value)
    
    def energy(self):
        """Potential energy of oscillator"""
        return 0.5 * self.frequency**2 * self.amplitude**2
    
    def align_phase(self, target_phase, strength=0.1):
        """Pull phase toward target"""
        diff = target_phase - self.phase
        self.phase = self.phase + strength * torch.atan2(torch.sin(diff), torch.cos(diff))
    
    def update_frequency(self, gradient_magnitude, lr=0.01):
        """Adapt frequency based on gradient magnitude"""
        self.frequency = self.frequency + lr * gradient_magnitude
        self.frequency = torch.clamp(self.frequency, 0.1, 10.0)


class ResonantLinear(nn.Module):
    """A linear layer where weights are harmonic oscillators"""
    
    def __init__(self, in_features, out_features, device='cpu'):
        super().__init__()
        self.device = device
        self.in_features = in_features
        self.out_features = out_features
        
        # Standard PyTorch weights (we'll wrap them)
        self.weight = nn.Parameter(torch.randn(out_features, in_features, device=device) * 0.1)
        self.bias = nn.Parameter(torch.zeros(out_features, device=device))
        
        # Oscillator containers
        self.weight_oscillators = []
        self.bias_oscillators = []
        
        self._init_oscillators()
        
        # Coupling parameters
        self.coupling_strength = 0.5
        
        # Phase statistics
        self.phase_coherence = 0.0
        self.phase_locked = False
        
    def _init_oscillators(self):
        """Initialize harmonic oscillators for each parameter"""
        # Weight oscillators
        for i in range(self.out_features):
            row_oscillators = []
            for j in range(self.in_features):
                osc = HarmonicOscillator(self.weight.data[i, j], self.device)
                row_oscillators.append(osc)
            self.weight_oscillators.append(row_oscillators)
        
        # Bias oscillators
        for i in range(self.out_features):
            self.bias_oscillators.append(HarmonicOscillator(self.bias.data[i], self.device))
    
    def forward(self, x, t=0.0):
        """Forward pass using oscillator states"""
        # Build differentiable effective parameters so autograd can trace back
        # to the underlying learnable weights and biases.
        weight_rows = []
        for i in range(self.out_features):
            row = []
            for j in range(self.in_features):
                osc_state = self.weight_oscillators[i][j].state(t)
                row.append(self.weight[i, j] + osc_state)
            weight_rows.append(torch.stack(row))
        effective_weight = torch.stack(weight_rows)
        
        bias_terms = []
        for i in range(self.out_features):
            bias_terms.append(self.bias[i] + self.bias_oscillators[i].state(t))
        effective_bias = torch.stack(bias_terms)
        
        return F.linear(x, effective_weight, effective_bias)
    
    def compute_coherence(self):
        """Compute phase coherence across all oscillators"""
        all_phases = []
        for row in self.weight_oscillators:
            for osc in row:
                all_phases.append(osc.phase.item())
        
        # Mean direction (circular mean)
        mean_x = np.mean([np.cos(p) for p in all_phases])
        mean_y = np.mean([np.sin(p) for p in all_phases])
        coherence = np.sqrt(mean_x**2 + mean_y**2)
        
        self.phase_coherence = coherence
        self.phase_locked = coherence > 0.8
        
        return coherence
    
    def resonant_update(self, gradients, lr=0.01, coupling_lr=0.1):
        """
        Update weights via RESONANCE mechanism
        Instead of: weight = weight - lr * gradient
        We do: oscillator alignment + coupling
        """
        
        with torch.no_grad():
            # 1. Update base weights with gradient (standard part)
            self.weight.grad = gradients['weight']
            self.weight.data -= lr * gradients['weight']
            
            self.bias.grad = gradients['bias']
            self.bias.data -= lr * gradients['bias']
            
            # 2. Resonance update: align oscillator phases with gradient direction
            for i in range(self.out_features):
                for j in range(self.in_features):
                    osc = self.weight_oscillators[i][j]
                    grad = gradients['weight'][i, j].item()
                    
                    # Target phase is based on gradient
                    target_phase = torch.atan2(torch.tensor(grad), torch.tensor(1.0))
                    
                    # Phase alignment (pull toward gradient direction)
                    osc.align_phase(target_phase, strength=coupling_lr)
                    
                    # Frequency adaptation
                    osc.update_frequency(abs(grad), lr=0.01)
                    
                    # Record history
                    osc.phase_history.append(osc.phase.item())
                    osc.gradient_history.append(grad)
            
            # 3. Bias oscillators
            for i in range(self.out_features):
                osc = self.bias_oscillators[i]
                grad = gradients['bias'][i].item()
                target_phase = torch.atan2(torch.tensor(grad), torch.tensor(1.0))
                osc.align_phase(target_phase, strength=coupling_lr)
                osc.phase_history.append(osc.phase.item())
    
    def couple_with_layer(self, other_layer, K=0.3):
        """
        Inter-layer resonance: this layer couples with another layer
        Phase-locked layers transfer energy to unlock neighbors
        """
        if not self.phase_locked:
            return 0.0
        
        # Find phase difference with each oscillator in other layer
        energy_transfer = 0.0
        
        for i in range(min(self.out_features, other_layer.out_features)):
            for j in range(min(self.in_features, other_layer.in_features)):
                my_phase = self.weight_oscillators[i][j].phase.item()
                other_phase = other_layer.weight_oscillators[i][j].phase.item()
                
                phase_diff = abs(my_phase - other_phase)
                resonance = K * np.cos(phase_diff)
                
                # Transfer energy to other layer
                other_layer.weight_oscillators[i][j].amplitude.data += resonance * 0.1
                
                energy_transfer += abs(resonance)
        
        return energy_transfer


class ResonantMLP(nn.Module):
    """Multi-layer perceptron with resonant layers"""
    
    def __init__(self, input_size=784, hidden_sizes=[256, 128], output_size=10, device='cpu'):
        super().__init__()
        self.device = device
        self.layers = nn.ModuleList()
        self.t = 0.0
        
        # Build layers
        sizes = [input_size] + hidden_sizes + [output_size]
        for i in range(len(sizes) - 1):
            self.layers.append(ResonantLinear(sizes[i], sizes[i+1], device))
        
        # Non-linearity after each layer except last
        self.activations = [F.relu] * (len(sizes) - 2) + [lambda x: x]
        
    def forward(self, x):
        x = x.view(x.size(0), -1)  # Flatten
        
        for i, layer in enumerate(self.layers):
            x = layer(x, t=self.t)
            x = self.activations[i](x)
            self.t += 0.1
        
        return x
    
    def resonant_backward(self, gradients, lr=0.01, coupling_lr=0.1):
        """Update all layers with resonance"""
        for i, layer in enumerate(self.layers):
            layer_grads = {
                'weight': gradients[f'weight_{i}'],
                'bias': gradients[f'bias_{i}']
            }
            layer.resonant_update(layer_grads, lr, coupling_lr)
            
            # Couple with previous layer
            if i > 0:
                layer.couple_with_layer(self.layers[i-1], K=0.3)
    
    def compute_all_coherence(self):
        """Get coherence of all layers"""
        return [layer.compute_coherence() for layer in self.layers]


# ============================================================
# PART 2: MULTI-AI RESONANCE SYSTEM
# ============================================================

class AIAgent:
    """An AI agent with its own model and oscillator parameters"""
    
    def __init__(self, agent_id, device='cpu'):
        self.id = agent_id
        self.device = device
        
        # Create model
        self.model = ResonantMLP(device=device).to(device)
        self.optimizer = optim.SGD(self.model.parameters(), lr=0.01)
        
        # Oscillator state
        self.phase = np.random.uniform(0, 2*np.pi)
        self.frequency = 1.0
        self.amplitude = 1.0
        
        # Training history
        self.loss_history = []
        self.accuracy_history = []
        self.phase_history = []
        self.gradient_history = []
        
        # Coupling strength
        self.K = 0.5
        
    def train_step(self, x, y, criterion, optimizer=None):
        """One training step for this agent"""
        self.model.train()
        optimizer = optimizer or self.optimizer
        optimizer.zero_grad()
        
        # Forward pass
        output = self.model(x)
        loss = criterion(output, y)
        
        # Backward pass (standard gradients)
        loss.backward()
        
        # Collect gradients for resonance
        gradients = {}
        for name, param in self.model.named_parameters():
            if param.grad is not None:
                gradients[name] = param.grad.clone()
        
        # Get gradient magnitude for phase update
        grad_mag = sum(p.sum().item() for p in gradients.values())
        
        # Update phase based on gradient
        self.phase += 0.1 * np.sin(grad_mag)
        self.phase_history.append(self.phase)
        self.gradient_history.append(grad_mag)
        
        return loss.item(), gradients
    
    def resonant_update(self, shared_gradients=None, K=0.5):
        """
        Update model with resonance
        If shared_gradients provided, apply cross-agent resonance
        """
        with torch.no_grad():
            for name, param in self.model.named_parameters():
                if param.grad is not None:
                    # Blend local and shared gradients
                    if shared_gradients and name in shared_gradients:
                        # Resonance: blend based on phase alignment
                        blend = K * np.cos(self.phase)
                        param.data -= 0.01 * ((1 - blend) * param.grad + 
                                              blend * shared_gradients[name])
                    else:
                        param.data -= 0.01 * param.grad
        
        # Update oscillator frequency based on recent gradients
        if self.gradient_history:
            recent_grad = np.mean(self.gradient_history[-10:])
            self.frequency *= (1 + 0.01 * recent_grad)
            self.frequency = np.clip(self.frequency, 0.5, 5.0)


class ResonantTrainingSystem:
    """Multi-agent training with cross-resonance"""
    
    def __init__(self, n_agents, device='cpu'):
        self.device = device
        self.agents = [AIAgent(i, device) for i in range(n_agents)]
        
        # Coupling matrix
        self.coupling_matrix = np.random.uniform(0.3, 0.7, (n_agents, n_agents))
        np.fill_diagonal(self.coupling_matrix, 0)
        
        # FFT analyzer
        self.fft_analyzer = FFTAnalyzer()
        
        # Training history
        self.epoch_history = []
        self.resonance_events = []
        
    def compute_coherence_matrix(self):
        """Compute phase coherence between all agent pairs"""
        n = len(self.agents)
        coherence = np.zeros((n, n))
        
        for i in range(n):
            for j in range(n):
                if i != j:
                    phase_diff = abs(self.agents[i].phase - self.agents[j].phase)
                    coherence[i, j] = np.cos(phase_diff)
        
        return coherence
    
    def find_resonance_pairs(self, threshold=0.5):
        """Find agent pairs with strong resonance"""
        coherence = self.compute_coherence_matrix()
        pairs = []
        
        for i in range(len(self.agents)):
            for j in range(i+1, len(self.agents)):
                if coherence[i, j] > threshold:
                    pairs.append((i, j, coherence[i, j]))
        
        return sorted(pairs, key=lambda x: x[2], reverse=True)
    
    def apply_cross_resonance(self, pairs):
        """Apply resonance between paired agents"""
        for i, j, strength in pairs:
            agent_i = self.agents[i]
            agent_j = self.agents[j]
            
            # Phase lock
            phase_diff = agent_i.phase - agent_j.phase
            agent_i.phase -= 0.1 * np.sin(phase_diff) * strength
            agent_j.phase += 0.1 * np.sin(phase_diff) * strength
            
            # Amplitude supergain
            combined_amp = (np.sqrt(agent_i.amplitude) + np.sqrt(agent_j.amplitude))**2
            agent_i.amplitude = combined_amp * 0.6
            agent_j.amplitude = combined_amp * 0.6
            
            self.resonance_events.append({
                'pair': (i, j),
                'strength': strength,
                'type': 'cross_resonance'
            })
    
    def train_epoch(self, train_loader, criterion):
        """Train one epoch with all agents"""
        epoch_loss = 0
        epoch_acc = 0
        total_samples = 0
        
        # Find resonance pairs
        resonance_pairs = self.find_resonance_pairs(threshold=0.3)
        
        # Apply cross-resonance
        self.apply_cross_resonance(resonance_pairs)
        
        for batch_idx, (data, target) in enumerate(train_loader):
            data, target = data.to(self.device), target.to(self.device)
            
            batch_losses = {}
            batch_gradients = {}
            
            # Compute local gradients for each agent once per batch
            for agent in self.agents:
                loss, gradients = agent.train_step(data, target, criterion)
                batch_losses[agent.id] = loss
                batch_gradients[agent.id] = gradients
            
            # Apply resonant updates using cached partner gradients
            for agent in self.agents:
                shared_grads = {}
                for i, j, strength in resonance_pairs:
                    if agent.id == i:
                        partner_grads = batch_gradients[j]
                        for name, grad in partner_grads.items():
                            if name not in shared_grads:
                                shared_grads[name] = grad * strength
                            else:
                                shared_grads[name] += grad * strength
                    elif agent.id == j:
                        partner_grads = batch_gradients[i]
                        for name, grad in partner_grads.items():
                            if name not in shared_grads:
                                shared_grads[name] = grad * strength
                            else:
                                shared_grads[name] += grad * strength
                
                # Apply resonant update
                agent.resonant_update(shared_grads if shared_grads else None, K=0.5)
                
                # Track loss
                epoch_loss += batch_losses[agent.id] * data.size(0)
                
                # Accuracy
                pred = agent.model(data).argmax(dim=1)
                epoch_acc += (pred == target).sum().item()
                total_samples += data.size(0)
        
        return epoch_loss / total_samples, epoch_acc / total_samples
    
    def evaluate(self, test_loader):
        """Evaluate all agents"""
        results = []
        
        for agent in self.agents:
            agent.model.eval()
            correct = 0
            total = 0
            
            with torch.no_grad():
                for data, target in test_loader:
                    data, target = data.to(self.device), target.to(self.device)
                    output = agent.model(data)
                    pred = output.argmax(dim=1)
                    correct += (pred == target).sum().item()
                    total += data.size(0)
            
            results.append({
                'agent_id': agent.id,
                'accuracy': correct / total
            })
        
        return results


# ============================================================
# PART 3: FFT ANALYZER
# ============================================================

class FFTAnalyzer:
    """Analyze training dynamics via FFT"""
    
    def __init__(self):
        self.spectra_history = []
        
    def compute_spectrum(self, signal):
        """Compute FFT spectrum of signal"""
        if len(signal) < 2:
            return np.array([0]), np.array([0])
        
        spectrum = np.fft.fft(signal)
        freqs = np.fft.fftfreq(len(signal))
        magnitude = np.abs(spectrum)
        
        return freqs[:len(freqs)//2], magnitude[:len(magnitude)//2]
    
    def cross_coherence(self, signal_a, signal_b):
        """Compute coherence between two signals"""
        if len(signal_a) != len(signal_b):
            min_len = min(len(signal_a), len(signal_b))
            signal_a = signal_a[:min_len]
            signal_b = signal_b[:min_len]
        
        # Normalize
        a_norm = np.array(signal_a) / (np.linalg.norm(signal_a) + 1e-10)
        b_norm = np.array(signal_b) / (np.linalg.norm(signal_b) + 1e-10)
        
        return np.abs(np.vdot(a_norm, b_norm))
    
    def dominant_frequency(self, signal):
        """Find dominant frequency in signal"""
        freqs, mag = self.compute_spectrum(signal)
        if len(mag) == 0:
            return 0
        return freqs[np.argmax(mag[1:]) + 1]  # Skip DC component


# ============================================================
# PART 4: BASELINE STANDARD TRAINING
# ============================================================

class StandardMLP(nn.Module):
    """Standard MLP for comparison"""
    
    def __init__(self, input_size=784, hidden_sizes=[256, 128], output_size=10):
        super().__init__()
        self.layers = nn.ModuleList()
        sizes = [input_size] + hidden_sizes + [output_size]
        for i in range(len(sizes) - 1):
            self.layers.append(nn.Linear(sizes[i], sizes[i+1]))
        self.activations = [F.relu] * (len(sizes) - 2) + [lambda x: x]
        
    def forward(self, x):
        x = x.view(x.size(0), -1)
        for i, layer in enumerate(self.layers):
            x = layer(x)
            x = self.activations[i](x)
        return x


# ============================================================
# PART 5: RUN THE EXPERIMENT
# ============================================================

def run_experiment(n_epochs=10, n_agents=3, batch_size=128):
    """Run the full resonant training experiment"""
    
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    print(f"\nUsing device: {device}")
    
    # Load MNIST
    transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.1307,), (0.3081,))
    ])
    
    train_dataset = datasets.MNIST('data', train=True, download=True, transform=transform)
    test_dataset = datasets.MNIST('data', train=False, transform=transform)
    
    train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
    test_loader = DataLoader(test_dataset, batch_size=batch_size)
    
    criterion = nn.CrossEntropyLoss()
    
    print("\n" + "=" * 70)
    print("EXPERIMENT: RESONANT vs STANDARD TRAINING")
    print("=" * 70)
    
    # ========== RESONANT TRAINING ==========
    print("\n🧠 TRAINING RESONANT SYSTEM (3 AI Agents)")
    print("-" * 50)
    
    resonant_system = ResonantTrainingSystem(n_agents=n_agents, device=device)
    
    resonant_losses = []
    resonant_accuracies = []
    coherence_history = []
    resonance_count_history = []
    
    for epoch in range(n_epochs):
        loss, acc = resonant_system.train_epoch(train_loader, criterion)
        resonant_losses.append(loss)
        resonant_accuracies.append(acc)
        
        # Compute coherence
        coherence = resonant_system.compute_coherence_matrix()
        coherence_history.append(coherence.copy())
        
        # Count resonance events
        n_resonance = len(resonant_system.resonance_events)
        resonance_count_history.append(n_resonance)
        
        print(f"Epoch {epoch+1:2d}: Loss={loss:.4f}, Acc={acc:.4f}, "
              f"Coherence={np.mean(coherence):.3f}, Resonance Events={n_resonance}")
    
    # Evaluate resonant system
    resonant_results = resonant_system.evaluate(test_loader)
    print(f"\n📊 Resonant System Test Accuracy: {np.mean([r['accuracy'] for r in resonant_results]):.4f}")
    
    # ========== STANDARD TRAINING ==========
    print("\n⚡ TRAINING STANDARD NETWORK (Baseline)")
    print("-" * 50)
    
    standard_model = StandardMLP().to(device)
    optimizer = optim.SGD(standard_model.parameters(), lr=0.01)
    
    standard_losses = []
    standard_accuracies = []
    
    for epoch in range(n_epochs):
        epoch_loss = 0
        epoch_acc = 0
        total_samples = 0
        
        for data, target in train_loader:
            data, target = data.to(device), target.to(device)
            
            optimizer.zero_grad()
            output = standard_model(data)
            loss = criterion(output, target)
            loss.backward()
            optimizer.step()
            
            epoch_loss += loss.item() * data.size(0)
            pred = output.argmax(dim=1)
            epoch_acc += (pred == target).sum().item()
            total_samples += data.size(0)
        
        avg_loss = epoch_loss / total_samples
        avg_acc = epoch_acc / total_samples
        standard_losses.append(avg_loss)
        standard_accuracies.append(avg_acc)
        
        print(f"Epoch {epoch+1:2d}: Loss={avg_loss:.4f}, Acc={avg_acc:.4f}")
    
    # Evaluate standard model
    standard_model.eval()
    correct = 0
    total = 0
    with torch.no_grad():
        for data, target in test_loader:
            data, target = data.to(device), target.to(device)
            output = standard_model(data)
            pred = output.argmax(dim=1)
            correct += (pred == target).sum().item()
            total += data.size(0)
    
    standard_accuracy = correct / total
    print(f"\n📊 Standard Network Test Accuracy: {standard_accuracy:.4f}")
    
    # ========== VISUALIZATION ==========
    visualize_results(
        resonant_losses, resonant_accuracies,
        standard_losses, standard_accuracies,
        coherence_history, resonant_system,
        resonant_results, standard_accuracy
    )
    
    return {
        'resonant': {
            'losses': resonant_losses,
            'accuracies': resonant_accuracies,
            'test_accuracy': np.mean([r['accuracy'] for r in resonant_results]),
            'coherence_history': coherence_history
        },
        'standard': {
            'losses': standard_losses,
            'accuracies': standard_accuracies,
            'test_accuracy': standard_accuracy
        }
    }


# ============================================================
# PART 6: VISUALIZATION
# ============================================================

def visualize_results(resonant_losses, resonant_accs, standard_losses, standard_accs,
                     coherence_history, resonant_system, resonant_results, standard_accuracy):
    """Generate comprehensive visualization of results"""
    
    fig = plt.figure(figsize=(18, 14))
    
    # ========== PLOT 1: Loss Comparison ==========
    ax1 = fig.add_subplot(3, 3, 1)
    ax1.plot(resonant_losses, 'b-', label='Resonant (3 AIs)', linewidth=2, marker='o')
    ax1.plot(standard_losses, 'r--', label='Standard', linewidth=2, marker='s')
    ax1.set_xlabel('Epoch')
    ax1.set_ylabel('Loss')
    ax1.set_title('Training Loss Comparison')
    ax1.legend()
    ax1.grid(True, alpha=0.3)
    
    # ========== PLOT 2: Accuracy Comparison ==========
    ax2 = fig.add_subplot(3, 3, 2)
    ax2.plot(resonant_accs, 'b-', label='Resonant (3 AIs)', linewidth=2, marker='o')
    ax2.plot(standard_accs, 'r--', label='Standard', linewidth=2, marker='s')
    ax2.set_xlabel('Epoch')
    ax2.set_ylabel('Accuracy')
    ax2.set_title('Training Accuracy Comparison')
    ax2.legend()
    ax2.grid(True, alpha=0.3)
    
    # ========== PLOT 3: Test Accuracy Bar ==========
    ax3 = fig.add_subplot(3, 3, 3)
    resonant_test_accs = [r['accuracy'] for r in resonant_results]
    x = np.arange(len(resonant_test_accs) + 1)
    bars = ax3.bar(x, resonant_test_accs + [standard_accuracy], 
                   color=['blue']*len(resonant_test_accs) + ['red'])
    ax3.axhline(y=np.mean(resonant_test_accs), color='blue', linestyle='--', alpha=0.5)
    ax3.set_xticks(x)
    ax3.set_xticklabels([f'AI-{i}' for i in range(len(resonant_test_accs))] + ['Standard'])
    ax3.set_ylabel('Test Accuracy')
    ax3.set_title('Final Test Accuracy')
    ax3.set_ylim([0.9, 1.0])
    
    # ========== PLOT 4: Phase Coherence Evolution ==========
    ax4 = fig.add_subplot(3, 3, 4)
    n_agents = len(resonant_system.agents)
    for i in range(n_agents):
        phases = [a.phase_history for a in resonant_system.agents]
    
    # Plot phase coherence heatmap
    coherence_array = np.array(coherence_history)
    im = ax4.imshow(coherence_array.T, aspect='auto', cmap='RdYlGn', vmin=-1, vmax=1)
    ax4.set_xlabel('Epoch')
    ax4.set_ylabel('Agent Pair')
    ax4.set_title('Phase Coherence Evolution')
    ax4.set_yticks(range(len(coherence_array[0])))
    ax4.set_yticklabels([f'({i},{j})' for i in range(n_agents) for j in range(i+1, n_agents)])
    plt.colorbar(im, ax=ax4)
    
    # ========== PLOT 5: Agent Phase Trajectories ==========
    ax5 = fig.add_subplot(3, 3, 5)
    colors = ['blue', 'red', 'green', 'orange', 'purple']
    for i, agent in enumerate(resonant_system.agents):
        phases = agent.phase_history
        if len(phases) > 0:
            ax5.plot(phases, color=colors[i], label=f'AI-{i}', alpha=0.7, linewidth=2)
    ax5.set_xlabel('Training Step')
    ax5.set_ylabel('Phase φ')
    ax5.set_title('Agent Phase Evolution')
    ax5.legend()
    ax5.grid(True, alpha=0.3)
    
    # ========== PLOT 6: FFT of Phase Signal ==========
    ax6 = fig.add_subplot(3, 3, 6)
    if len(resonant_system.agents[0].phase_history) > 10:
        phases = np.array(resonant_system.agents[0].phase_history)
        spectrum = np.fft.fft(phases - phases.mean())  # Remove DC
        freqs = np.fft.fftfreq(len(phases))
        magnitude = np.abs(spectrum)
        ax6.plot(freqs[:len(freqs)//2], magnitude[:len(magnitude)//2], 'b-', linewidth=2)
        ax6.set_xlabel('Frequency')
        ax6.set_ylabel('Magnitude')
        ax6.set_title('FFT of Agent Phase Signal')
        ax6.grid(True, alpha=0.3)
    
    # ========== PLOT 7: Resonant Layer Coherence ==========
    ax7 = fig.add_subplot(3, 3, 7)
    layer_coherences = []
    for agent in resonant_system.agents:
        coherences = agent.model.compute_all_coherence()
        layer_coherences.append(coherences)
    
    layer_coherences = np.array(layer_coherences)
    for i in range(layer_coherences.shape[1]):
        ax7.plot(layer_coherences[:, i], marker='o', label=f'Layer {i+1}', linewidth=2)
    ax7.set_xlabel('Agent')
    ax7.set_ylabel('Phase Coherence')
    ax7.set_title('Layer Phase Coherence per Agent')
    ax7.legend()
    ax7.grid(True, alpha=0.3)
    
    # ========== PLOT 8: Gradient Magnitude History ==========
    ax8 = fig.add_subplot(3, 3, 8)
    for i, agent in enumerate(resonant_system.agents):
        grads = agent.gradient_history
        if len(grads) > 0:
            # Smooth
            if len(grads) > 10:
                grads_smooth = np.convolve(grads, np.ones(10)/10, mode='valid')
                ax8.plot(grads_smooth, color=colors[i], label=f'AI-{i}', alpha=0.7)
    ax8.set_xlabel('Training Step')
    ax8.set_ylabel('Gradient Magnitude')
    ax8.set_title('Gradient Evolution (Smoothed)')
    ax8.legend()
    ax8.grid(True, alpha=0.3)
    
    # ========== PLOT 9: Resonance Events Timeline ==========
    ax9 = fig.add_subplot(3, 3, 9)
    resonance_types = ['cross_resonance'] * len(resonant_system.resonance_events)
    if resonance_types:
        ax9.hist([e['strength'] for e in resonant_system.resonance_events], 
                 bins=20, color='purple', alpha=0.7, edgecolor='black')
    ax9.set_xlabel('Resonance Strength')
    ax9.set_ylabel('Count')
    ax9.set_title('Resonance Events Distribution')
    ax9.grid(True, alpha=0.3)
    
    plt.tight_layout()
    plt.savefig('resonant_mnist_results.png', dpi=150, bbox_inches='tight')
    plt.show()
    
    print("\n✅ Visualization saved as 'resonant_mnist_results.png'")


# ============================================================
# RUN THE MAIN EXPERIMENT
# ============================================================

if __name__ == "__main__":
    results = run_experiment(n_epochs=10, n_agents=3, batch_size=128)
    
    print("\n" + "=" * 70)
    print("EXPERIMENT SUMMARY")
    print("=" * 70)
    
    print(f"\n📊 RESONANT SYSTEM (3 AIs):")
    print(f"   Final Training Accuracy: {results['resonant']['accuracies'][-1]:.4f}")
    print(f"   Final Test Accuracy: {results['resonant']['test_accuracy']:.4f}")
    print(f"   Final Loss: {results['resonant']['losses'][-1]:.4f}")
    
    print(f"\n📊 STANDARD NETWORK:")
    print(f"   Final Training Accuracy: {results['standard']['accuracies'][-1]:.4f}")
    print(f"   Final Test Accuracy: {results['standard']['test_accuracy']:.4f}")
    print(f"   Final Loss: {results['standard']['losses'][-1]:.4f}")
    
    improvement = (results['resonant']['test_accuracy'] - results['standard']['test_accuracy']) * 100
    print(f"\n{'✅' if improvement >= 0 else '⚠️'} Test Accuracy Difference: {improvement:+.2f}%")
    
    print("\n🧠 KEY OBSERVATIONS:")
    print("   - Resonant training uses PHASE ALIGNMENT instead of pure gradient descent")
    print("   - Multiple AIs share gradient information via RESONANCE COUPLING")
    print("   - Phase coherence tracked via FFT analysis")
    print("   - Supergain effect may emerge when phases align")