"""
============================================================
  10-ATTACTOR N-BODY MNIST CLASSIFIER
  Based on Conditional Collapse Theory (CCT) + 
  Gravitational N-Body Optimization Framework (GNBOF)
============================================================

  Concept: Each digit (0-9) is a gravitational attractor.
  - Parameters orbit around these attractors
  - Black holes = overfitting singularities (memorization)
  - Lagrange points = decision boundaries
  - Energy conservation = convergence guarantee
============================================================
"""

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch.utils.data import DataLoader, TensorDataset
import torchvision
import torchvision.transforms as transforms
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.patches import Circle
import time
import math
from dataclasses import dataclass, field
from typing import List, Dict, Tuple, Optional

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

@dataclass
class AttractorConfig:
    """Configuration for each digit attractor"""
    digit: int
    mass: float = 1.0  # Strength of attractor
    position: torch.Tensor = None  # In latent space
    is_black_hole: bool = False  # Overfitting detected
    event_horizon_radius: float = 0.0
    orbital_energy: float = 0.0
    
@dataclass
class GNBOFConfig:
    """Global N-Body Optimization Config"""
    G: float = 0.1  # Gravitational constant
    beta: float = 0.9  # Friction/momentum coefficient
    singularity_threshold: float = 5.0  # Mass threshold for black hole
    hawking_rate: float = 0.01  # Regularization evaporation rate
    escape_velocity_scale: float = 1.0
    lagrange_threshold: float = 0.5  # Decision boundary proximity
    periodicity_tolerance: float = 1e-4
    compute_budget: float = 1e6
    entropy_threshold: float = 0.1
    
class NBodyOptimizer:
    """
    N-Body Gravitational Optimizer
    
    Instead of standard gradient descent, parameters experience
    gravitational pull from all 10 digit attractors.
    """
    
    def __init__(self, parameters, config: GNBOFConfig, num_attractors=10):
        self.params = list(parameters)
        self.config = config
        self.num_attractors = num_attractors
        
        # State tracking
        self.velocity = [torch.zeros_like(p) for p in self.params]
        self.position_history = []  # For periodicity detection
        self.attractors: List[AttractorConfig] = []
        self.current_entropy = float('inf')
        self.total_energy_spent = 0.0
        self.collapse_achieved = False
        
        # Black hole registry
        self.black_holes: List[int] = []  # Indices of overfitting attractors
        
    def initialize_attractors(self, model, dataloader):
        """
        Initialize 10 attractors in latent space.
        Each attractor = mean embedding of its digit class.
        """
        print("\n[GNBOF] Initializing 10 digit attractors...")
        
        # Get model embeddings
        model.eval()
        all_embeddings = []
        all_labels = []
        max_samples = 500
        collected = 0
        
        with torch.no_grad():
            for batch_idx, (inputs, targets) in enumerate(dataloader):
                if collected >= max_samples:
                    break
                # Get embeddings before final layer
                if hasattr(model, 'get_embedding'):
                    emb = model.get_embedding(inputs)
                else:
                    emb = inputs.view(inputs.size(0), -1)  # Fallback

                take = min(emb.size(0), max_samples - collected)
                all_embeddings.append(emb[:take])
                all_labels.append(targets[:take])
                collected += take
        
        all_embeddings = torch.cat(all_embeddings, dim=0)
        all_labels = torch.cat(all_labels, dim=0)
        
        # Create one attractor per digit
        self.attractors = []
        for digit in range(10):
            mask = all_labels == digit
            if mask.sum() > 0:
                center = all_embeddings[mask].mean(dim=0)
            else:
                center = torch.randn(all_embeddings.shape[1]) * 0.1
            
            # Mass based on class separation (more separated = stronger attractor)
            mass = 1.0  # Base mass
            
            self.attractors.append(AttractorConfig(
                digit=digit,
                mass=mass,
                position=center.clone().detach().requires_grad_(True)
            ))
        
        print(f"[GNBOF] Attractors initialized with masses: {[a.mass for a in self.attractors]}")
        return self.attractors
    
    def compute_gravitational_force(self, position: torch.Tensor) -> torch.Tensor:
        """
        Compute total gravitational force on parameter position from all attractors.
        
        F_total = Σ (G * m_i * (p_i - pos) / ||p_i - pos||³)
        """
        total_force = torch.zeros_like(position)
        
        for i, attractor in enumerate(self.attractors):
            if i in self.black_holes:
                # Black holes exert repulsive force (avoid overfitting)
                direction = position - attractor.position
                distance = torch.norm(direction) + 1e-6
                repulsive = -self.config.G * attractor.mass / (distance ** 2 + 1.0)
                total_force += repulsive * (direction / distance)
            else:
                # Normal gravitational attraction
                direction = attractor.position - position
                distance = torch.norm(direction) + 1e-6
                attraction = self.config.G * attractor.mass / (distance ** 2 + 1e-6)
                total_force += attraction * (direction / distance)
        
        return total_force
    
    def update_attractor_mass(self, model, dataloader, epoch):
        """
        Update attractor masses based on training dynamics.
        - High loss = weak attractor
        - Overfitting detected = convert to black hole
        """
        model.eval()
        digit_losses = {i: [] for i in range(10)}
        
        with torch.no_grad():
            for inputs, targets in dataloader:
                outputs = model(inputs)
                for d in range(10):
                    mask = targets == d
                    if mask.sum() > 0:
                        digit_losses[d].append(F.cross_entropy(outputs[mask], targets[mask], reduction='mean').item())
        
        # Update masses and detect black holes
        for i, attractor in enumerate(self.attractors):
            if len(digit_losses[i]) > 0:
                avg_loss = np.mean(digit_losses[i])
                # Attractor strength inversely proportional to loss
                new_mass = 1.0 / (1.0 + avg_loss)
                
                # Singularity detection: mass growing too fast = overfitting
                if new_mass > self.config.singularity_threshold and not attractor.is_black_hole:
                    attractor.is_black_hole = True
                    self.black_holes.append(i)
                    print(f"[GNBOF] ⚫ BLACK HOLE DETECTED: Digit {i} overfitting!")
                
                attractor.mass = new_mass
                
                # Hawking radiation: slowly evaporate black holes
                if attractor.is_black_hole:
                    attractor.mass -= self.config.hawking_rate
                    if attractor.mass < self.config.singularity_threshold:
                        attractor.is_black_hole = False
                        self.black_holes.remove(i)
                        print(f"[GNBOF] ⭐ Black hole {i} evaporated via Hawking radiation")
    
    def step(self, closure):
        """
        Single N-Body optimization step.
        """
        # Compute loss
        loss = closure()
        
        # Get parameter position
        position = torch.cat([p.flatten() for p in self.params])
        
        # Compute gravitational force from attractors
        gravity = self.compute_gravitational_force(position)
        
        # Gradient from loss (standard)
        grads = torch.autograd.grad(loss, self.params, retain_graph=True)
        grad_flat = torch.cat([g.flatten() for g in grads])
        
        # Combine: gradient direction + gravitational pull
        combined_force = grad_flat + gravity * 0.1  # Gravity strength factor
        
        # Update momentum (orbital velocity)
        for i, p in enumerate(self.params):
            v = self.velocity[i]
            v_new = self.config.beta * v + (1 - self.config.beta) * grads[i]
            self.velocity[i] = v_new
            p.data -= self.config.escape_velocity_scale * v_new
        
        # Track energy spent
        self.total_energy_spent += loss.item()
        
        return loss
    
    def detect_periodicity(self) -> bool:
        """
        Check if parameters are in stable orbit (converged).
        Periodicity = params approximately repeat after k steps.
        """
        current_pos = torch.cat([p.flatten().clone().detach() for p in self.params])
        
        if len(self.position_history) > 10:
            # Check for match with history
            for t in range(len(self.position_history) - 1, max(0, len(self.position_history) - 20), -1):
                diff = torch.norm(current_pos - self.position_history[t])
                if diff < self.config.periodicity_tolerance:
                    return True
        
        self.position_history.append(current_pos)
        if len(self.position_history) > 50:
            self.position_history.pop(0)
        
        return False
    
    def get_cct_entropy(self, model, dataloader) -> float:
        """
        Compute CCT entropy: uncertainty about digit classification.
        Lower = more confident = more collapsed.
        """
        model.eval()
        all_probs = []
        
        with torch.no_grad():
            for inputs, _ in dataloader:
                outputs = model(inputs)
                probs = F.softmax(outputs, dim=1)
                all_probs.append(probs)
        
        all_probs = torch.cat(all_probs, dim=0)
        
        # Entropy: H = -Σ p * log(p)
        entropy = -torch.sum(all_probs * torch.log(all_probs + 1e-10)) / all_probs.size(0)
        self.current_entropy = entropy.item()
        
        return entropy.item()
    
    def check_collapse(self, model, dataloader) -> bool:
        """
        CCT Collapse Condition:
        - Entropy below threshold
        - OR periodicity detected (stable orbit)
        """
        entropy = self.get_cct_entropy(model, dataloader)
        periodic = self.detect_periodicity()
        
        if entropy < self.config.entropy_threshold:
            print(f"[GNBOF] ✓ COLLAPSE: Entropy {entropy:.4f} < {self.config.entropy_threshold}")
            self.collapse_achieved = True
            return True
        
        if periodic:
            print(f"[GNBOF] ✓ COLLAPSE: Stable orbit detected (periodicity)")
            self.collapse_achieved = True
            return True
        
        return False
    
    def get_lagrange_points(self) -> List[torch.Tensor]:
        """
        Find Lagrange points between attractors = decision boundaries.
        These are equilibrium points where forces balance.
        """
        lagrange_points = []
        n = len(self.attractors)
        
        for i in range(n):
            for j in range(i + 1, n):
                # Midpoint between two attractors
                mid = (self.attractors[i].position + self.attractors[j].position) / 2
                
                # Check if it's approximately equidistant
                d1 = torch.norm(mid - self.attractors[i].position)
                d2 = torch.norm(mid - self.attractors[j].position)
                
                if abs(d1 - d2) < self.config.lagrange_threshold:
                    lagrange_points.append(mid)
        
        return lagrange_points


# ============================================================
# MODEL ARCHITECTURE: Attractor-Based Classifier
# ============================================================

class AttractorMNIST(nn.Module):
    """
    MNIST classifier with attractor-based dynamics.
    Hidden activations orbit around digit attractors.
    """
    
    def __init__(self, latent_dim=64, num_attractors=10):
        super().__init__()
        self.latent_dim = latent_dim
        self.num_attractors = num_attractors
        
        # Feature extraction
        self.conv1 = nn.Conv2d(1, 32, 3, padding=1)
        self.conv2 = nn.Conv2d(32, 64, 3, padding=1)
        self.pool = nn.MaxPool2d(2, 2)
        self.fc1 = nn.Linear(64 * 7 * 7, latent_dim)
        
        # Attractor layer: projects latent to attractor space
        self.attractor_projection = nn.Linear(latent_dim, num_attractors)
        
        # Per-attractor classifiers
        self.classifiers = nn.ModuleList([
            nn.Sequential(
                nn.Linear(latent_dim, 32),
                nn.ReLU(),
                nn.Linear(32, 1)
            ) for _ in range(num_attractors)
        ])
        
        # Register attractor positions as learnable
        self.attractor_positions = nn.Parameter(torch.randn(num_attractors, latent_dim) * 0.1)
        
    def get_embedding(self, x):
        x = self.pool(F.relu(self.conv1(x)))
        x = self.pool(F.relu(self.conv2(x)))
        x = x.view(x.size(0), -1)
        x = F.relu(self.fc1(x))
        return x
    
    def forward(self, x):
        # Get latent embedding
        z = self.get_embedding(x)
        
        # Compute attractor-guided logits
        # Each attractor "pulls" the embedding toward its class
        
        # Standard classification
        logits = self.attractor_projection(z)
        
        # Add attractor-based correction
        # Embedding feels gravitational pull from each attractor
        attractor_logits = []
        for i in range(self.num_attractors):
            # Distance to attractor i
            dist = torch.norm(z - self.attractor_positions[i], dim=1, keepdim=True)
            # Closer to attractor = higher logit for that class
            correction = -dist  # Negative distance = attraction
            attractor_logits.append(correction)
        
        attractor_logits = torch.cat(attractor_logits, dim=1)
        
        # Combine
        final_logits = logits + attractor_logits * 0.1
        
        return final_logits
    
    def get_attractor_states(self) -> Dict:
        """Return current state of all attractors"""
        return {
            'positions': self.attractor_positions.data.clone(),
            'gradients': self.attractor_positions.grad.clone()
                         if self.attractor_positions.grad is not None
                         else None
        }


# ============================================================
# TRAINING LOOP: CCT-GNBOF Style
# ============================================================

class CCTNBodyTrainer:
    """
    Trainer implementing Conditional Collapse Theory + N-Body Optimization.
    """
    
    def __init__(self, model, optimizer, config: GNBOFConfig):
        self.model = model
        self.optimizer = optimizer
        self.config = config
        
        self.history = {
            'epoch': [],
            'train_loss': [],
            'test_acc': [],
            'entropy': [],
            'energy_spent': [],
            'black_holes': [],
            'attractor_masses': []
        }
        
    def train_epoch(self, train_loader, epoch):
        self.model.train()
        total_loss = 0
        correct = 0
        total = 0
        
        for batch_idx, (inputs, targets) in enumerate(train_loader):
            # CCT: Check entropy before processing
            if batch_idx % 100 == 0:
                entropy = self.optimizer.get_cct_entropy(self.model, train_loader)
                print(f"  [Batch {batch_idx}] Entropy: {entropy:.4f}, Black Holes: {len(self.optimizer.black_holes)}")
            
            # Forward pass
            outputs = self.model(inputs)
            loss = F.cross_entropy(outputs, targets)
            
            # N-Body update
            self.model.zero_grad()
            loss.backward()
            
            # Custom N-Body gradient modification
            with torch.no_grad():
                for i, p in enumerate(self.model.parameters()):
                    if p.grad is not None:
                        # Add gravitational component
                        if p.shape == self.model.attractor_positions.shape:
                            # Attractor positions get extra gravitational influence
                            gravity = self.optimizer.compute_gravitational_force(p)
                            p.grad += gravity * 0.01
            
            # Manual update with momentum
            with torch.no_grad():
                for i, p in enumerate(self.model.parameters()):
                    if p.grad is not None and i < len(self.optimizer.velocity):
                        v = self.optimizer.velocity[i]
                        v.mul_(self.config.beta).add_(p.grad, alpha=1 - self.config.beta)
                        p.sub_(v, alpha=self.config.escape_velocity_scale)
            
            total_loss += loss.item()
            _, predicted = outputs.max(1)
            total += targets.size(0)
            correct += predicted.eq(targets).sum().item()
        
        # Update attractor masses based on training
        self.optimizer.update_attractor_mass(self.model, train_loader, epoch)
        
        # CCT: Check for collapse
        collapsed = self.optimizer.check_collapse(self.model, train_loader)
        
        return total_loss / len(train_loader), 100. * correct / total, collapsed
    
    def test(self, test_loader):
        self.model.eval()
        correct = 0
        total = 0
        
        with torch.no_grad():
            for inputs, targets in test_loader:
                outputs = self.model(inputs)
                _, predicted = outputs.max(1)
                total += targets.size(0)
                correct += predicted.eq(targets).sum().item()
        
        return 100. * correct / total
    
    def train(self, train_loader, test_loader, epochs=20):
        print("\n" + "="*60)
        print("  CCT-GNBOF TRAINING: 10-ATTRACTOR MNIST")
        print("="*60)
        
        # Initialize attractors
        self.optimizer.initialize_attractors(self.model, train_loader)
        
        for epoch in range(epochs):
            start_time = time.time()
            
            # Train
            train_loss, train_acc, collapsed = self.train_epoch(train_loader, epoch)
            
            # Test
            test_acc = self.test(test_loader)
            
            # Record
            entropy = self.optimizer.get_cct_entropy(self.model, test_loader)
            
            self.history['epoch'].append(epoch)
            self.history['train_loss'].append(train_loss)
            self.history['test_acc'].append(test_acc)
            self.history['entropy'].append(entropy)
            self.history['energy_spent'].append(self.optimizer.total_energy_spent)
            self.history['black_holes'].append(len(self.optimizer.black_holes))
            self.history['attractor_masses'].append([a.mass for a in self.optimizer.attractors])
            
            # Print
            print(f"\n[Epoch {epoch+1}/{epochs}] {time.time() - start_time:.1f}s")
            print(f"  Train Loss: {train_loss:.4f} | Train Acc: {train_acc:.2f}%")
            print(f"  Test Acc:   {test_acc:.2f}%")
            print(f"  Entropy:    {entropy:.4f}")
            print(f"  Energy:     {self.optimizer.total_energy_spent:.2f}")
            print(f"  Black Holes: {len(self.optimizer.black_holes)}")
            print(f"  Attractor Masses: {[f'{a.mass:.3f}' for a in self.optimizer.attractors]}")
            
            if collapsed:
                print("\n✓ COLLAPSE ACHIEVED - Theory understood!")
                if epoch < epochs - 1:
                    print("  Early stopping - entropy threshold reached")
                    break
        
        return self.history
    
    def visualize_attractors(self, save_path='attractor_evolution.png'):
        """Visualize the attractor dynamics during training"""
        fig, axes = plt.subplots(2, 2, figsize=(14, 10))
        
        # Plot 1: Test Accuracy
        axes[0, 0].plot(self.history['epoch'], self.history['test_acc'], 'b-', linewidth=2)
        axes[0, 0].set_xlabel('Epoch')
        axes[0, 0].set_ylabel('Test Accuracy (%)')
        axes[0, 0].set_title('Classification Performance')
        axes[0, 0].grid(True)
        
        # Plot 2: Entropy Collapse
        axes[0, 1].plot(self.history['epoch'], self.history['entropy'], 'r-', linewidth=2)
        axes[0, 1].axhline(y=self.config.entropy_threshold, color='g', linestyle='--', 
                          label=f'Threshold ({self.config.entropy_threshold})')
        axes[0, 1].set_xlabel('Epoch')
        axes[0, 1].set_ylabel('CCT Entropy')
        axes[0, 1].set_title('Semantic Collapse Progress')
        axes[0, 1].legend()
        axes[0, 1].grid(True)
        
        # Plot 3: Attractor Masses
        masses = np.array(self.history['attractor_masses'])
        for i in range(10):
            axes[1, 0].plot(self.history['epoch'], masses[:, i], 
                           label=f'Digit {i}', linewidth=1.5)
        axes[1, 0].set_xlabel('Epoch')
        axes[1, 0].set_ylabel('Attractor Mass')
        axes[1, 0].set_title('Gravitational Mass Evolution')
        axes[1, 0].legend(loc='upper right', fontsize=8, ncol=2)
        axes[1, 0].grid(True)
        
        # Plot 4: Black Hole Detection
        axes[1, 1].bar(self.history['epoch'], self.history['black_holes'], color='purple', alpha=0.7)
        axes[1, 1].set_xlabel('Epoch')
        axes[1, 1].set_ylabel('Number of Black Holes')
        axes[1, 1].set_title('Singularity Detection (Overfitting)')
        axes[1, 1].grid(True)
        
        plt.tight_layout()
        plt.savefig(save_path, dpi=150)
        print(f"\n[GNBOF] Visualization saved to {save_path}")
        plt.show()
    
    def visualize_loss_landscape(self, save_path='loss_landscape.png'):
        """
        Visualize the 2D loss landscape with attractors and black holes.
        """
        # Create a 2D grid in weight space
        fig, ax = plt.subplots(figsize=(10, 10))
        
        # Sample weights from model
        weights = torch.cat([p.flatten() for p in self.model.parameters()])
        if weights.shape[0] > 1000:
            weights = weights[:1000]  # Use subset for visualization
        
        # PCA-like reduction to 2D (just use first two dimensions for demo)
        w1 = weights[:len(weights)//2].cpu().numpy() if len(weights) > 2 else np.random.randn(500)
        w2 = weights[len(weights)//2:].cpu().numpy() if len(weights) > 2 else np.random.randn(500)
        
        if len(w1) < 10 or len(w2) < 10:
            w1 = np.random.randn(500)
            w2 = np.random.randn(500)
        
        # Grid for loss evaluation
        x_range = np.linspace(w1.min() - 1, w1.max() + 1, 50)
        y_range = np.linspace(w2.min() - 1, w2.max() + 1, 50)
        X, Y = np.meshgrid(x_range, y_range)
        
        # Simulated loss landscape (in real implementation, evaluate actual loss)
        Z = np.zeros_like(X)
        for i in range(X.shape[0]):
            for j in range(X.shape[1]):
                # Simplified: multiple basins (attractors)
                z = 0
                for k, attractor in enumerate(self.optimizer.attractors):
                    mass = attractor.mass if not attractor.is_black_hole else -attractor.mass
                    pos_x = np.random.uniform(w1.min(), w1.max())
                    pos_y = np.random.uniform(w2.min(), w2.max())
                    dist = np.sqrt((X[i,j] - pos_x)**2 + (Y[i,j] - pos_y)**2) + 0.1
                    if attractor.is_black_hole:
                        z -= mass / dist  # Repulsive (black hole)
                    else:
                        z += mass / dist  # Attractive (normal attractor)
                Z[i,j] = -z  # Invert for loss surface
        
        # Plot contours
        contour = ax.contourf(X, Y, Z, levels=20, cmap='viridis', alpha=0.6)
        plt.colorbar(contour, ax=ax, label='Loss')
        
        # Plot attractors
        for i, attractor in enumerate(self.optimizer.attractors):
            if i < len(w1) and i < len(w2):
                x_pos = w1[i % len(w1)]
                y_pos = w2[i % len(w2)]
                
                if attractor.is_black_hole:
                    # Black hole: black circle with event horizon
                    circle = Circle((x_pos, y_pos), radius=0.5 * attractor.mass,
                                   color='black', alpha=0.8)
                    ax.add_patch(circle)
                    ax.plot(x_pos, y_pos, 'kx', markersize=15, markeredgewidth=3)
                else:
                    # Normal attractor: colored circle
                    circle = Circle((x_pos, y_pos), radius=0.3 * attractor.mass,
                                   color=plt.cm.Set1(i/10), alpha=0.7)
                    ax.add_patch(circle)
                    ax.text(x_pos, y_pos, str(attractor.digit), ha='center', va='center',
                           fontsize=10, fontweight='bold', color='white')
        
        ax.set_xlabel('Weight Dimension 1')
        ax.set_ylabel('Weight Dimension 2')
        ax.set_title('10-Attractor Loss Landscape\n(Black Circles = Overfitting Singularities)')
        
        plt.tight_layout()
        plt.savefig(save_path, dpi=150)
        print(f"\n[GNBOF] Loss landscape saved to {save_path}")
        plt.show()


# ============================================================
# MAIN: RUN TRAINING
# ============================================================

def main():
    print("\n" + "="*60)
    print("  CCT-GNBOF: 10-Attractor MNIST Classification")
    print("="*60)
    
    # Load MNIST
    transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.5,), (0.5,))
    ])
    
    train_dataset = torchvision.datasets.MNIST(
        root='../data', train=True, download=True, transform=transform
    )
    test_dataset = torchvision.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=128, shuffle=False)
    
    print(f"\n[DATA] Train: {len(train_dataset)} samples | Test: {len(test_dataset)} samples")
    
    # Initialize model
    model = AttractorMNIST(latent_dim=64, num_attractors=10)
    print(f"\n[MODEL] {sum(p.numel() for p in model.parameters())} parameters")
    
    # Initialize N-Body Optimizer
    config = GNBOFConfig(
        G=0.05,
        beta=0.9,
        singularity_threshold=3.0,
        hawking_rate=0.05,
        escape_velocity_scale=0.01,
        entropy_threshold=0.3
    )
    
    optimizer = NBodyOptimizer(model.parameters(), config, num_attractors=10)
    
    # Initialize trainer
    trainer = CCTNBodyTrainer(model, optimizer, config)
    
    # Train
    history = trainer.train(train_loader, test_loader, epochs=20)
    
    # Visualize results
    trainer.visualize_attractors('attractor_evolution.png')
    trainer.visualize_loss_landscape('loss_landscape.png')
    
    # Final metrics
    print("\n" + "="*60)
    print("  FINAL RESULTS")
    print("="*60)
    print(f"  Final Test Accuracy: {history['test_acc'][-1]:.2f}%")
    print(f"  Final Entropy:       {history['entropy'][-1]:.4f}")
    print(f"  Total Energy Spent:  {history['energy_spent'][-1]:.2f}")
    print(f"  Black Holes at End:  {history['black_holes'][-1]}")
    
    return model, history


if __name__ == '__main__':
    model, history = main()