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

class OverfitProofGeometricOptimizer(optim.Optimizer):
    """
    Implements a system of Multi-Point Derivative Objects engineered to prevent overfitting:
    - Itô Strip: Continuous exploratory diffusion to break memorization.
    - Wavelet Packet Tree (Simplified): Filters out fine-scale localized noise.
    - Caputo Fractional Arc: Power-law memory kernel to resist rapid local convergence.
    - Hessian Spectrum Diamond: Adaptive landscape flattening per dimension.
    """
    def __init__(self, params, lr=1e-3, alpha=0.6, drift_coef=0.7, diff_coef=0.1, memory_len=10):
        defaults = dict(lr=lr, alpha=alpha, drift_coef=drift_coef, diff_coef=diff_coef, memory_len=memory_len)
        super(OverfitProofGeometricOptimizer, 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]
                
                # 1. 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

                # 2. WAVELET PACKET TREE (Multi-scale Coarse Filter)
                # We separate the update into coarse (mean structural shifts) and fine variations.
                # To prevent overfitting to this one batch, we dampen the fine-scale details.
                coarse_trend = torch.mean(scaled_grad) * torch.ones_like(scaled_grad)
                fine_details = scaled_grad - coarse_trend
                filtered_grad = coarse_trend + (0.1 * fine_details)  # Attenuate noise leaves

                # 3. 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)

                # 4. 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)  # Power-law memory scaling
                    fractional_step.add_(past_update, alpha=weight)
                fractional_step.div_(h_len)

                # Apply the generalized multi-point calibrated step
                p.add_(fractional_step, alpha=-lr)
        return loss

# --- Architecture & Experiment 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)

# Load Data
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 = torch.utils.data.DataLoader(
    datasets.MNIST(root='../data', train=False, download=True, transform=transform), 
    batch_size=1000, shuffle=False
)

# Extract EXACTLY ONE isolated sample batch
single_batch_loader = torch.utils.data.DataLoader(train_dataset, batch_size=64, shuffle=True)
static_images, static_labels = next(iter(single_batch_loader))

# Initialize model and our custom system optimizer
model = SimpleMLP()
criterion = nn.CrossEntropyLoss()
optimizer = OverfitProofGeometricOptimizer(model.parameters(), lr=0.005, diff_coef=0.25, alpha=0.5)

print("Starting 100 loops of overtraining on the EXACT SAME sample batch...\n")

for i in range(1, 101):
    model.train()
    optimizer.zero_grad()
    
    output = model(static_images)
    loss = criterion(output, static_labels)
    loss.backward()
    optimizer.step()
    
    if i % 20 == 0 or i == 1:
        # Check generalized performance on the unseen validation set
        model.eval()
        correct, total = 0, 0
        with torch.no_grad():
            for test_images, test_labels in test_loader:
                preds = model(test_images).argmax(dim=1)
                correct += (preds == test_labels).sum().item()
                total += test_labels.size(0)
        
        print(f"Iteration {i:3d} | Current Train Batch Loss: {loss.item():.4f} | Unseen Test Accuracy: {100.*correct/total:.2f}%")