import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms
from torch.utils.data import DataLoader
from collections import deque
import math

# ---------------------------------------------------------
# 1. Custom Anti-Overfitting Geometric Optimizer
# ---------------------------------------------------------
class AntiOverfitGeometricOptimizer(optim.Optimizer):
    def __init__(self, params, lr=1e-3, alpha=0.6, drift_coef=0.7, diff_coef=0.15, memory_len=8):
        defaults = dict(lr=lr, alpha=alpha, drift_coef=drift_coef, diff_coef=diff_coef, memory_len=memory_len)
        super(AntiOverfitGeometricOptimizer, self).__init__(params, defaults)
        
        for group in self.param_groups:
            for p in group['params']:
                state = self.state[p]
                state['history'] = deque(maxlen=memory_len)
                state['curvature_buffer'] = torch.zeros_like(p.data)

    @torch.no_grad()
    def step(self, closure=None):
        loss = None
        for group in self.param_groups:
            lr = group['lr']
            alpha = group['alpha']
            drift_coef = group['drift_coef']
            diff_coef = group['diff_coef']
            
            for p in group['params']:
                if p.grad is None:
                    continue
                
                grad = p.grad.clone()
                state = self.state[p]
                
                # Concept 7: Hessian Spectrum Diamond (Landscape Scaling)
                state['curvature_buffer'].mul_(0.95).add_(grad.pow(2), alpha=0.05)
                diamond_scale = 1.0 / (torch.sqrt(state['curvature_buffer']) + 1e-6)
                scaled_grad = grad * diamond_scale

                # Concept 17: Wavelet Packet Tree (Multi-scale Filtering)
                coarse_trend = torch.mean(scaled_grad) * torch.ones_like(scaled_grad)
                fine_details = scaled_grad - coarse_trend
                filtered_grad = coarse_trend + (0.15 * fine_details)  # Attenuate noise leaves

                # Concept 21: Itô Strip (Drift vs. Diffusion Exploration)
                drift = filtered_grad
                diffusion = torch.randn_like(grad) * diff_coef
                ito_step = (drift_coef * drift) + ((1.0 - drift_coef) * diffusion)

                # Concept 26: Caputo Fractional Arc (Memory Path Integration)
                state['history'].append(ito_step)
                
                fractional_step = torch.zeros_like(grad)
                h_len = len(state['history'])
                for i, past_update in enumerate(state['history']):
                    k = h_len - i
                    weight = 1.0 / math.pow(k, 1.0 - alpha)
                    fractional_step.add_(past_update, alpha=weight)
                fractional_step.div_(h_len)

                # Execute parameter shift
                p.add_(fractional_step, alpha=-lr)
        return loss

# ---------------------------------------------------------
# 2. Network & Training Pipeline Setup
# ---------------------------------------------------------
class SimpleMLP(nn.Module):
    def __init__(self):
        super().__init__()
        self.net = nn.Sequential(
            nn.Flatten(),
            nn.Linear(28*28, 128),
            nn.ReLU(),
            nn.Linear(128, 10)
        )
    def forward(self, x): 
        return self.net(x)

def main():
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    
    # Load Datasets
    transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))])
    train_dataset = datasets.MNIST(root='../data', train=True, download=True, transform=transform)
    test_loader = DataLoader(datasets.MNIST(root='../data', train=False, download=True, transform=transform), batch_size=1000, shuffle=False)
    
    # Isolate EXACTLY 100 sample batches into memory to reuse each epoch
    temp_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)
    fixed_batches = []
    for i, (data, target) in enumerate(temp_loader):
        if i >= 100: 
            break
        fixed_batches.append((data.to(device), target.to(device)))

    # Initialize Environment
    model = SimpleMLP().to(device)
    criterion = nn.CrossEntropyLoss()
    optimizer = AntiOverfitGeometricOptimizer(model.parameters(), lr=0.004, diff_coef=0.20, alpha=0.55)

    epochs = 3
    print(f"Beginning training across {epochs} Epochs utilizing the SAME 100 static batches...\n")

    for epoch in range(1, epochs + 1):
        model.train()
        epoch_loss = 0.0
        
        # Run through the identical set of 100 batches inside this epoch iteration
        for batch_idx, (data, target) in enumerate(fixed_batches):
            optimizer.zero_grad()
            output = model(data)
            loss = criterion(output, target)
            loss.backward()
            optimizer.step()
            
            epoch_loss += loss.item()

        # Evaluate Generalized Performance on the completely unseen global Validation Set
        model.eval()
        correct, total = 0, 0
        with torch.no_grad():
            for test_images, test_labels in test_loader:
                test_images, test_labels = test_images.to(device), test_labels.to(device)
                preds = model(test_images).argmax(dim=1)
                correct += (preds == test_labels).sum().item()
                total += test_labels.size(0)

        avg_loss = epoch_loss / len(fixed_batches)
        accuracy = 100. * correct / total
        print(f"Epoch {epoch}/{epochs} Summary Metrics:")
        print(f"    Avg Train Loss (Over the 100 Batches): {avg_loss:.4f}")
        print(f"    Unseen Global Test Set Accuracy      : {accuracy:.2f}%\n" + "-"*55)

if __name__ == '__main__':
    main()