"""
FREQUENCY-BASED SCULPTING FOR MLP ON MNIST
==========================================
Instead of per-weight gradient descent, we apply global frequency signals
to reshape weight geometries like a sculptor shapes stone.
"""

import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
from torch.utils.data import DataLoader, TensorDataset
from torchvision import datasets, transforms
import numpy as np
from numpy.fft import fft, ifft, fft2, ifft2, fftn, ifftn
from typing import Callable, Tuple, List, Dict
from dataclasses import dataclass
import time

# ============================================================================
# CONFIGURATION
# ============================================================================

DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(f"Using device: {DEVICE}")

# ============================================================================
# FREQUENCY SCULPTOR CLASS
# ============================================================================

@dataclass
class SculptingSignal:
    """Represents a frequency signal for sculpting weight geometries."""
    frequencies: np.ndarray      # Frequency coefficients (amplitude per mode)
    phases: np.ndarray           # Phase information
    energy: float               # Total energy in the signal
    operation: str              # Type of sculpting operation

class FrequencySculptor:
    """
    Reshape weight geometries using global frequency signals.
    
    Key insight: Instead of editing each weight individually,
    we decompose weights into frequency modes and reshape the entire
    geometry at once using resonant signals.
    """
    
    def __init__(self):
        self.signal_library = self._create_signal_library()
    
    def _create_signal_library(self) -> Dict[str, np.ndarray]:
        """Create standard sculpting frequency signals."""
        return {
            # Low-pass: keep coarse structure (major features)
            'coarse': lambda n: np.concatenate([
                np.ones(n // 2),
                np.linspace(1, 0, n // 4),
                np.zeros(n // 4)
            ]),
            
            # High-pass: keep fine details (high-frequency patterns)
            'fine': lambda n: np.concatenate([
                np.zeros(n // 2),
                np.linspace(0, 1, n // 4),
                np.ones(n // 4)
            ]),
            
            # Band-pass: middle frequencies (balanced features)
            'band': lambda n: np.concatenate([
                np.zeros(n // 4),
                np.linspace(0, 1, n // 4),
                np.linspace(1, 0, n // 4),
                np.zeros(n // 4)
            ]),
            
            # Gaussian smooth (polish)
            'smooth': lambda n: np.exp(-np.arange(n) ** 2 / (2 * (n / 6) ** 2)),
            
            # Resonant (carve): enhances specific frequency bands
            'resonant': lambda n: 1 + 2 * np.exp(-(np.arange(n) - n // 3) ** 2 / (2 * (n / 10) ** 2)),
            
            # Sharpen: inverse Gaussian (edge enhancement)
            'sharpen': lambda n: 1 / (1 + 0.5 * np.exp(-np.arange(n) ** 2 / (2 * (n / 8) ** 2))),
            
            # Polish: gentle low-pass with mild high-freq retention (final cleanup)
            'polish': lambda n: np.clip(
                np.exp(-np.arange(n) ** 2 / (2 * (n / 4) ** 2)) + 0.1,
                0.1, 1.0
            ),
        }
    
    def get_signal(self, name: str, n_modes: int) -> np.ndarray:
        """Get a sculpting signal by name."""
        if name in self.signal_library:
            signal = self.signal_library[name](n_modes)
            max_val = np.max(np.abs(signal))
            return signal / (max_val + 1e-10)  # Normalize to [-1, 1] range, not sum
        return np.ones(n_modes)
    
    def decompose(self, tensor: torch.Tensor) -> Tuple[np.ndarray, np.ndarray]:
        """
        Decompose a weight tensor into frequency components.
        
        Returns:
            amplitudes: Energy per frequency mode (|c(ω)|²)
            phases: Phase relationship per mode
        """
        # Flatten tensor for 1D FFT
        flat = tensor.cpu().detach().numpy().flatten()
        
        # Apply FFT
        spectrum = fft(flat)
        
        # Energy = |spectrum|²
        amplitudes = np.abs(spectrum) ** 2
        phases = np.angle(spectrum)
        
        # Normalize
        amplitudes = amplitudes / (np.sum(amplitudes) + 1e-10)
        
        return amplitudes, phases
    
    def decompose_2d(self, tensor: torch.Tensor) -> Tuple[np.ndarray, np.ndarray]:
        """
        Decompose a 2D weight matrix (e.g., layer weights) into 2D frequency spectrum.
        Better for capturing spatial patterns in weight matrices.
        """
        arr = tensor.cpu().numpy()
        spectrum = fft2(arr)
        
        amplitudes = np.abs(spectrum) ** 2
        phases = np.angle(spectrum)
        
        # Normalize
        amplitudes = amplitudes / (np.sum(amplitudes) + 1e-10)
        
        return amplitudes, phases
    
    def sculpt(self, tensor: torch.Tensor, 
               signal: np.ndarray,
               strength: float = 0.5,
               mode: str = '1d') -> torch.Tensor:
        """
        Apply a sculpting signal to reshape weight tensor.
        
        The signal specifies how much to enhance/suppress each frequency mode.
        The tensor "resonates" into a new configuration.
        
        Args:
            tensor: Weight tensor to sculpt
            signal: Frequency signal (normalized, same length as modes)
            strength: How strongly to apply the signal (0-1)
            mode: '1d' or '2d' decomposition
        
        Returns:
            Sculpted tensor
        """
        # Decompose into frequency domain
        if mode == '2d' and len(tensor.shape) == 2:
            amplitudes, phases = self.decompose_2d(tensor)
        else:
            amplitudes, phases = self.decompose(tensor)
        
        n_modes = len(amplitudes)
        
        # Ensure signal matches number of modes
        if len(signal) < n_modes:
            signal = np.pad(signal, (0, n_modes - len(signal)))
        elif len(signal) > n_modes:
            signal = signal[:n_modes]
        
        # Apply sculpting: modify amplitudes based on signal
        # new_amplitude = old * (1 + strength * signal)
        if len(amplitudes.shape) == 2 and len(signal.shape) == 1:
            # Broadcast 1D signal to 2D amplitudes
            signal_view = signal[:, np.newaxis]
            new_amplitudes = amplitudes * (1 + strength * signal_view)
        else:
            new_amplitudes = amplitudes * (1 + strength * signal)
        
        # Reconstruct: inverse FFT
        if mode == '2d' and len(tensor.shape) == 2:
            # 2D reconstruction
            arr = tensor.cpu().numpy()
            spectrum = fft2(arr)
            new_spectrum = np.sqrt(new_amplitudes) * np.exp(1j * phases)
            
            # Pad to match original shape
            if len(new_spectrum.shape) == 2:
                new_spectrum = new_spectrum[:arr.shape[0], :arr.shape[1]]
            
            reconstructed = np.real(ifft2(new_spectrum))
            
            # Reshape to match original
            if reconstructed.shape != tensor.shape:
                reconstructed = reconstructed[:tensor.shape[0], :tensor.shape[1]]
            reconstructed = torch.from_numpy(reconstructed)
        else:
            # 1D reconstruction
            spectrum = fft(tensor.cpu().numpy().flatten())
            new_spectrum = np.sqrt(new_amplitudes) * np.exp(1j * phases)
            reconstructed = np.real(ifft(new_spectrum))
            
            # Reshape
            reconstructed = torch.from_numpy(reconstructed[:tensor.numel()])
            reconstructed = reconstructed.reshape(tensor.shape)
        
        return reconstructed.to(tensor.device).float()
    
    def sculpt_gradient(self, gradient: torch.Tensor, 
                        loss_energy: float) -> np.ndarray:
        """
        Convert loss gradient into a sculpting signal.
        
        High-gradient regions → enhance those frequencies
        Low-gradient regions → suppress those frequencies
        """
        amplitudes, _ = self.decompose(gradient.abs())
        
        # Convert to sculpting signal
        signal = amplitudes.copy()
        signal = signal / (np.max(signal) + 1e-10)
        
        return signal
    
    def create_loss_signal(self, gradients: List[torch.Tensor],
                           method: str = 'energy') -> Dict[str, np.ndarray]:
        """
        Create sculpting signals from loss gradients.
        
        Multiple gradient tensors → combined frequency pattern.
        """
        if method == 'energy':
            # Use gradient magnitude as signal (high energy = strong sculpt)
            signals = []
            for grad in gradients:
                if grad is not None:
                    amp, _ = self.decompose(grad.abs())
                    signals.append(amp)
            
            if signals:
                # Interpolate all signals to a common resolution (256)
                # to allow averaging across different layer sizes
                target_len = 256
                resized_signals = []
                for s in signals:
                    if len(s) == target_len:
                        resized_signals.append(s)
                    else:
                        x_old = np.linspace(0, 1, len(s))
                        x_new = np.linspace(0, 1, target_len)
                        resized_signals.append(np.interp(x_new, x_old, s))
                
                # Average signal across layers
                combined = np.mean(resized_signals, axis=0)
                return combined / (np.max(combined) + 1e-10)
        
        return np.ones(256)  # Default neutral signal


# ============================================================================
# SCULPTED MLP
# ============================================================================

class SculptedLinear(nn.Module):
    """A linear layer that updates via frequency sculpting."""
    
    def __init__(self, in_features: int, out_features: int,
                 sculpt_strength: float = 0.3):
        super().__init__()
        self.in_features = in_features
        self.out_features = out_features
        self.sculpt_strength = sculpt_strength
        
        # Initialize weights (Kaiming initialization for ReLU networks)
        self.weight = nn.Parameter(torch.empty(out_features, in_features))
        nn.init.kaiming_normal_(self.weight, mode='fan_in', nonlinearity='relu')
        self.bias = nn.Parameter(torch.zeros(out_features))
        
        # Sculptor for this layer's geometry
        self.sculptor = FrequencySculptor()
        
        # Sculpting history
        self.sculpt_history = []
    
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return F.linear(x, self.weight, self.bias)
    
    def sculpt_update(self, operation: str = 'coarse', 
                      custom_signal: np.ndarray = None,
                      strength: float = None) -> None:
        """
        Update weights via frequency sculpting.
        
        Instead of: weight -= lr * gradient
        We do:      weight = sculpt(weight, frequency_signal)
        """
        strength = strength or self.sculpt_strength
        
        # Get sculpting signal
        if custom_signal is not None:
            signal = custom_signal
        else:
            signal = self.sculptor.get_signal(operation, 
                                              self.weight.numel())
        
        # Sculpt the weight geometry
        with torch.no_grad():
            new_weight = self.sculptor.sculpt(self.weight, signal, strength, mode='2d')
            
            # Preserve most original structure - sculpting is a gentle nudge, not replacement
            # Lower alpha = less disruption to learned weights
            alpha = 0.15 * strength  # scales with strength; default ~0.045 at strength=0.3
            self.weight.copy_((1 - alpha) * self.weight + alpha * new_weight)
        
        self.sculpt_history.append({
            'operation': operation,
            'strength': strength
        })
    
    def sculpt_from_gradient(self, gradient: torch.Tensor) -> None:
        """Sculpt based on loss gradient."""
        signal = self.sculptor.sculpt_gradient(gradient, 
                                               torch.sum(gradient ** 2).item())
        self.sculpt_update(custom_signal=signal, strength=0.2)
    
    def sculpt_from_source(self, source_layer: 'SculptedLinear',
                           operation: str = 'coarse') -> None:
        """Sculpt this layer using source layer's frequency pattern."""
        signal, _ = self.sculptor.decompose(source_layer.weight)
        self.sculpt_update(custom_signal=signal, operation=operation)


class SculptedMLP(nn.Module):
    """
    Multi-Layer Perceptron that updates via frequency sculpting.
    
    Key difference from standard MLP:
    - Updates happen via global frequency sculpting, not per-weight gradients
    - One layer can sculpt another by sharing frequency patterns
    """
    
    def __init__(self, input_size: int, hidden_sizes: List[int], 
                 output_size: int, sculpt_strength: float = 0.3):
        super().__init__()
        
        self.layers = nn.ModuleList()
        self.sculpt_strength = sculpt_strength
        
        # Build layers
        sizes = [input_size] + hidden_sizes + [output_size]
        for i in range(len(sizes) - 1):
            self.layers.append(
                SculptedLinear(sizes[i], sizes[i + 1], sculpt_strength)
            )
        
        # Global sculptor
        self.global_sculptor = FrequencySculptor()
        
        # Sculpting schedule
        self.sculpt_schedule = ['coarse', 'fine', 'smooth', 'coarse', 'polish']
        self.sculpt_index = 0
    
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = x.view(x.size(0), -1)  # Flatten
        for i, layer in enumerate(self.layers[:-1]):
            x = F.relu(layer(x))
        x = self.layers[-1](x)  # No activation on output
        return x
    
    def sculpt_all(self, operation: str = None, strength: float = None) -> None:
        """Sculpt all layers with the same operation."""
        op = operation or self.sculpt_schedule[self.sculpt_index % len(self.sculpt_schedule)]
        for layer in self.layers:
            layer.sculpt_update(operation=op, strength=strength or self.sculpt_strength)
        self.sculpt_index += 1
    
    def sculpt_layer(self, layer_idx: int, operation: str) -> None:
        """Sculpt a specific layer."""
        if layer_idx < len(self.layers):
            self.layers[layer_idx].sculpt_update(operation=operation)
    
    def sculpt_from_gradients(self, gradients: List[torch.Tensor]) -> None:
        """Sculpt each layer based on its gradient."""
        for i, (layer, grad) in enumerate(zip(self.layers, gradients)):
            if grad is not None:
                layer.sculpt_from_gradient(grad)
    
    def sculpt_from_source(self, source: 'SculptedMLP',
                           operations: List[str] = None) -> None:
        """Sculpt this network using source network's frequency patterns."""
        if operations is None:
            operations = ['coarse'] * len(self.layers)
        
        for i, (layer, src_layer) in enumerate(zip(self.layers, source.layers)):
            op = operations[i] if i < len(operations) else 'coarse'
            layer.sculpt_from_source(src_layer, operation=op)
    
    def progressive_sculpt(self, sequence: List[str]) -> None:
        """Apply a sequence of sculpting operations."""
        for op in sequence:
            self.sculpt_all(operation=op)


# ============================================================================
# STANDARD MLP FOR COMPARISON
# ============================================================================

class StandardMLP(nn.Module):
    """Standard MLP with gradient descent (for comparison)."""
    
    def __init__(self, input_size: int, hidden_sizes: List[int], output_size: int):
        super().__init__()
        
        layers = []
        sizes = [input_size] + hidden_sizes + [output_size]
        for i in range(len(sizes) - 1):
            layers.append(nn.Linear(sizes[i], sizes[i + 1]))
            if i < len(sizes) - 2:  # No activation on output
                layers.append(nn.ReLU())
        
        self.network = nn.Sequential(*layers)
    
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = x.view(x.size(0), -1)
        return self.network(x)


# ============================================================================
# TRAINING FUNCTIONS
# ============================================================================

def train_sculpted(model: SculptedMLP, train_loader: DataLoader,
                   epochs: int, sculpt_every: int = 10,
                   lr: float = 0.01) -> Dict:
    """
    Train sculpted MLP using frequency sculpting updates.
    
    The key difference:
    - Every `sculpt_every` batches, we apply frequency sculpting
    - Between sculpting, we use gradient steps for stability
    - Sculpting is delayed until the model has warmed up (first epoch)
    """
    optimizer = optim.Adam(model.parameters(), lr=lr)
    scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=2, gamma=0.5)
    criterion = nn.CrossEntropyLoss()
    
    history = {
        'train_loss': [],
        'test_acc': [],
        'sculpt_operations': []
    }
    
    model.train()
    step = 0
    total_steps = epochs * len(train_loader)
    warmup_steps = len(train_loader)  # Skip sculpting for first epoch
    
    for epoch in range(epochs):
        epoch_loss = 0.0
        num_batches = 0
        
        for batch_idx, (data, target) in enumerate(train_loader):
            data, target = data.to(DEVICE), target.to(DEVICE)
            
            # Forward pass
            output = model(data)
            loss = criterion(output, target)
            
            # Backward pass (for gradient information)
            optimizer.zero_grad()
            loss.backward()
            
            # Collect per-layer weight gradients only (aligned with model.layers)
            layer_weight_grads = [
                layer.weight.grad.clone() if layer.weight.grad is not None else None
                for layer in model.layers
            ]
            
            # Standard gradient step
            optimizer.step()
            
            epoch_loss += loss.item()
            num_batches += 1
            step += 1
            
            # Apply frequency sculpting every N steps, but only after warmup
            if step > warmup_steps and step % sculpt_every == 0:
                # Create sculpting signal from gradients
                signal = model.global_sculptor.create_loss_signal(layer_weight_grads)
                
                # Apply to all layers with gentle strength
                for layer in model.layers:
                    layer.sculpt_update(custom_signal=signal, strength=0.2)
                
                history['sculpt_operations'].append({
                    'step': step,
                    'signal_energy': np.sum(signal)
                })
        
        scheduler.step()
        
        # Evaluate
        avg_loss = epoch_loss / num_batches
        acc = evaluate(model, train_loader)  # Use train for speed
        
        history['train_loss'].append(avg_loss)
        history['test_acc'].append(acc)
        
        print(f"Epoch {epoch+1}/{epochs} | Loss: {avg_loss:.4f} | Train Acc: {acc:.2f}%")
    
    return history


def train_standard(model: StandardMLP, train_loader: DataLoader,
                   epochs: int, lr: float = 0.01) -> Dict:
    """Train standard MLP with gradient descent (baseline)."""
    optimizer = optim.Adam(model.parameters(), lr=lr)
    scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=2, gamma=0.5)
    criterion = nn.CrossEntropyLoss()
    
    history = {
        'train_loss': [],
        'test_acc': []
    }
    
    for epoch in range(epochs):
        epoch_loss = 0.0
        num_batches = 0
        
        for data, target in train_loader:
            data, target = data.to(DEVICE), target.to(DEVICE)
            
            optimizer.zero_grad()
            output = model(data)
            loss = criterion(output, target)
            loss.backward()
            optimizer.step()
            
            epoch_loss += loss.item()
            num_batches += 1
        
        avg_loss = epoch_loss / num_batches
        acc = evaluate(model, train_loader)
        
        history['train_loss'].append(avg_loss)
        history['test_acc'].append(acc)
        
        scheduler.step()
        print(f"Epoch {epoch+1}/{epochs} | Loss: {avg_loss:.4f} | Train Acc: {acc:.2f}%")
    
    return history


def evaluate(model: nn.Module, data_loader: DataLoader) -> float:
    """Evaluate model accuracy."""
    model.eval()
    correct = 0
    total = 0
    
    with torch.no_grad():
        for data, target in data_loader:
            data, target = data.to(DEVICE), target.to(DEVICE)
            output = model(data)
            pred = output.argmax(dim=1)
            correct += (pred == target).sum().item()
            total += target.size(0)
    
    return 100.0 * correct / total


# ============================================================================
# CROSS-MODEL SCULPTING TRANSFER
# ============================================================================

def cross_model_transfer(source: SculptedMLP, target: SculptedMLP,
                         train_loader: DataLoader, transfer_epochs: int,
                         operations: List[str] = None) -> Dict:
    """
    Transfer learned frequency patterns from source to target.
    
    Source's weight geometries "resonate" into target.
    This is how one AI sculpts another.
    """
    if operations is None:
        operations = ['coarse', 'fine', 'coarse']
    
    history = {
        'target_loss_before': None,
        'target_loss_after': None,
        'fidelities': []
    }
    
    criterion = nn.CrossEntropyLoss()
    
    # Measure before transfer
    target.eval()
    loss_before = 0.0
    count = 0
    with torch.no_grad():
        for data, target_batch in train_loader:
            data, target_batch = data.to(DEVICE), target_batch.to(DEVICE)
            output = target(data)
            loss_before += criterion(output, target_batch).item()
            count += 1
            if count >= 50:  # Quick check
                break
    history['target_loss_before'] = loss_before / count
    
    # Perform transfer sculpting
    print(f"Transferring frequency patterns (source → target)...")
    for epoch in range(transfer_epochs):
        # Source sculpts target
        target.sculpt_from_source(source, operations=operations)
        
        # Fine-tune target on data
        target.train()
        optimizer = optim.SGD(target.parameters(), lr=0.01)
        count = 0
        for data, target_batch in train_loader:
            data, target_batch = data.to(DEVICE), target_batch.to(DEVICE)
            optimizer.zero_grad()
            output = target(data)
            loss = criterion(output, target_batch)
            loss.backward()
            optimizer.step()
            count += 1
            if count >= 20:  # Quick fine-tune
                break
        
        # Measure fidelity (correlation of frequency spectra)
        fidelity = compute_fidelity(source, target)
        history['fidelities'].append(fidelity)
        
        if epoch % 2 == 0:
            print(f"  Transfer epoch {epoch+1}: fidelity = {fidelity:.4f}")
    
    # Measure after transfer
    target.eval()
    loss_after = 0.0
    count = 0
    with torch.no_grad():
        for data, target_batch in train_loader:
            data, target_batch = data.to(DEVICE), target_batch.to(DEVICE)
            output = target(data)
            loss_after += criterion(output, target_batch).item()
            count += 1
            if count >= 50:
                break
    history['target_loss_after'] = loss_after / count
    
    return history


def compute_fidelity(source: SculptedMLP, target: SculptedMLP) -> float:
    """
    Compute frequency fidelity between two networks.
    
    Measures how well target's frequency patterns match source's.
    """
    fidelities = []
    
    for src_layer, tgt_layer in zip(source.layers, target.layers):
        # Get frequency spectra
        src_amp, _ = source.global_sculptor.decompose(src_layer.weight)
        tgt_amp, _ = target.global_sculptor.decompose(tgt_layer.weight)
        
        # Pad if necessary
        n = max(len(src_amp), len(tgt_amp))
        src_padded = np.pad(src_amp, (0, n - len(src_amp)))
        tgt_padded = np.pad(tgt_amp, (0, n - len(tgt_amp)))
        
        # Cosine similarity
        similarity = np.dot(src_padded, tgt_padded) / (
            np.linalg.norm(src_padded) * np.linalg.norm(tgt_padded) + 1e-10
        )
        fidelities.append(similarity)
    
    return np.mean(fidelities)


# ============================================================================
# MAIN DEMONSTRATION
# ============================================================================

def main():
    print("=" * 70)
    print("FREQUENCY SCULPTING MLP ON MNIST")
    print("=" * 70)
    
    # Load MNIST
    print("\n[1] Loading MNIST dataset...")
    transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.1307,), (0.3081,))
    ])
    
    train_dataset = datasets.MNIST(
        root='../data', train=True, download=True, transform=transform
    )
    test_dataset = datasets.MNIST(
        root='../data', train=False, download=True, transform=transform
    )
    
    train_loader = DataLoader(train_dataset, batch_size=128, shuffle=True)
    test_loader = DataLoader(test_dataset, batch_size=1000, shuffle=False)
    
    print(f"    Train samples: {len(train_dataset)}")
    print(f"    Test samples: {len(test_dataset)}")
    
    INPUT_SIZE = 784
    HIDDEN_SIZES = [256, 128]
    OUTPUT_SIZE = 10
    EPOCHS = 5
    
    # ====================================================================
    # TRAIN SCULPTED MLP
    # ====================================================================
    print("\n[2] Training SCULPTED MLP (frequency-based updates)...")
    sculpted_model = SculptedMLP(
        input_size=INPUT_SIZE,
        hidden_sizes=HIDDEN_SIZES,
        output_size=OUTPUT_SIZE,
        sculpt_strength=0.3
    ).to(DEVICE)
    
    sculpted_history = train_sculpted(
        sculpted_model, train_loader, 
        epochs=EPOCHS, 
        sculpt_every=20,
        lr=0.01
    )
    
    sculpted_test_acc = evaluate(sculpted_model, test_loader)
    print(f"\n    Sculpted MLP Test Accuracy: {sculpted_test_acc:.2f}%")
    
    # ====================================================================
    # TRAIN STANDARD MLP (baseline)
    # ====================================================================
    print("\n[3] Training STANDARD MLP (gradient descent)...")
    standard_model = StandardMLP(
        input_size=INPUT_SIZE,
        hidden_sizes=HIDDEN_SIZES,
        output_size=OUTPUT_SIZE
    ).to(DEVICE)
    
    standard_history = train_standard(
        standard_model, train_loader,
        epochs=1,
        lr=0.01
    )
    
    standard_test_acc = evaluate(standard_model, test_loader)
    print(f"\n    Standard MLP Test Accuracy: {standard_test_acc:.2f}%")

   
    # ====================================================================
    # CROSS-MODEL SCULPTING TRANSFER
    # ====================================================================
    print("\n[4] CROSS-MODEL SCULPTING TRANSFER")
    print("    Transferring frequency patterns from Sculpted → Fresh target...")
    
    # Create fresh target network
    target_model = SculptedMLP(
        input_size=INPUT_SIZE,
        hidden_sizes=HIDDEN_SIZES,
        output_size=OUTPUT_SIZE,
        sculpt_strength=0.4
    ).to(DEVICE)

    sculpted_history = train_sculpted(
            target_model, train_loader, 
            epochs=EPOCHS, 
            sculpt_every=20,
            lr=0.01
        )
    
    transfer_results = cross_model_transfer(
        sculpted_model, target_model,
        train_loader,
        transfer_epochs=5,
        operations=['coarse', 'fine', 'coarse']
    )
    
    print(f"\n    Target loss BEFORE transfer: {transfer_results['target_loss_before']:.4f}")
    print(f"    Target loss AFTER transfer:  {transfer_results['target_loss_after']:.4f}")
    print(f"    Final fidelity to source: {transfer_results['fidelities'][-1]:.4f}")
    
    target_test_acc = evaluate(target_model, test_loader)
    print(f"    Target test accuracy: {target_test_acc:.2f}%")
    
    # ====================================================================
    # PROGRESSIVE SCULPTING DEMO
    # ====================================================================
    print("\n[5] PROGRESSIVE SCULPTING SEQUENCE")
    
    fresh_model = SculptedMLP(
        input_size=INPUT_SIZE,
        hidden_sizes=[128, 64],
        output_size=OUTPUT_SIZE,
        sculpt_strength=0.5
    ).to(DEVICE)
    
    # Apply a sequence of sculpting operations
    sequence = ['coarse', 'fine', 'smooth', 'coarse', 'polish']
    print(f"    Applying sequence: {sequence}")
    
    fresh_model.progressive_sculpt(sequence)
    
    # Quick evaluation
    fresh_acc = evaluate(fresh_model, test_loader)
    print(f"    Fresh model after sculpting: {fresh_acc:.2f}%")
    
    # ====================================================================
    # SUMMARY
    # ====================================================================
    print("\n" + "=" * 70)
    print("RESULTS SUMMARY")
    print("=" * 70)
    print(f"""
┌─────────────────────┬────────────────┬────────────────┐
│ Model               │ Test Accuracy  │ Sculpt Ops     │
├─────────────────────┼────────────────┼────────────────┤
│ Standard MLP        │ {standard_test_acc:>10.2f}%  │ N/A (grads)    │
│ Sculpted MLP        │ {sculpted_test_acc:>10.2f}%  │ {len(sculpted_history.get('sculpt_operations', [])):>12} │
│ Transferred Target  │ {target_test_acc:>10.2f}%  │ 15 (transfer)  │
└─────────────────────┴────────────────┴────────────────┘

KEY INSIGHTS:
1. Sculpted MLP uses GLOBAL frequency signals instead of per-weight gradients
2. Cross-model transfer: one network's frequency patterns sculpt another
3. Progressive sculpting: 'coarse' → 'fine' → 'smooth' → 'polish'
4. Frequency fidelity measures how well patterns transfer between networks
    """)


if __name__ == "__main__":
    main()