import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
import numpy as np
import matplotlib.pyplot as plt
from tqdm import tqdm

# ------------------------------
# 1. Define MLP with two hidden layers
# ------------------------------
class FractalMLP(nn.Module):
    def __init__(self, input_size=784, hidden1=256, hidden2=128, num_classes=10):
        super().__init__()
        self.fc1 = nn.Linear(input_size, hidden1)
        self.fc2 = nn.Linear(hidden1, hidden2)
        self.fc3 = nn.Linear(hidden2, num_classes)
        self.activation = nn.ReLU()

    def forward(self, x):
        h1 = self.activation(self.fc1(x))
        h2 = self.activation(self.fc2(h1))
        out = self.fc3(h2)
        return out, h1, h2   # return all layer outputs for analysis

# ------------------------------
# 2. Training setup
# ------------------------------
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
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=64, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=1000, shuffle=False)

model = FractalMLP().to(device)
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

def train(epochs=5):
    model.train()
    for epoch in range(epochs):
        for data, target in tqdm(train_loader, desc=f"Epoch {epoch+1}"):
            data, target = data.to(device), target.to(device)
            data = data.view(data.size(0), -1)
            optimizer.zero_grad()
            output, _, _ = model(data)
            loss = criterion(output, target)
            loss.backward()
            optimizer.step()

train(epochs=5)

# ------------------------------
# 3. Measure fractal exponents per layer
# ------------------------------
# For a given input, we add Gaussian noise of varying scales ε.
# We measure how much the hidden layer activations change as ε → 0.
# The scaling exponent α is defined by: ||Δh|| ~ ε^α
# (For a perfectly linear layer α=1; for a fractal mapping α<1.)

def compute_fractal_exponents(model, test_inputs, noise_scales):
    """
    Returns: (alpha1, alpha2, alpha_out) for each input sample
    """
    model.eval()
    exponents1, exponents2, exponents_out = [], [], []
    with torch.no_grad():
        for x in test_inputs:
            x = x.to(device)
            # clean forward pass
            out_clean, h1_clean, h2_clean = model(x)
            # store sensitivity per noise scale
            sens1, sens2, sens_out = [], [], []
            for eps in noise_scales:
                noise = torch.randn_like(x) * eps
                x_noisy = x + noise
                out_noisy, h1_noisy, h2_noisy = model(x_noisy)
                # relative change in each layer
                d1 = torch.norm(h1_noisy - h1_clean).item()
                d2 = torch.norm(h2_noisy - h2_clean).item()
                dout = torch.norm(out_noisy - out_clean).item()
                sens1.append(d1)
                sens2.append(d2)
                sens_out.append(dout)
            # Fit power law: log(d) = α * log(ε) + const
            log_eps = np.log(noise_scales)
            alpha1 = np.polyfit(log_eps, np.log(sens1), 1)[0]
            alpha2 = np.polyfit(log_eps, np.log(sens2), 1)[0]
            alpha_out = np.polyfit(log_eps, np.log(sens_out), 1)[0]
            exponents1.append(alpha1)
            exponents2.append(alpha2)
            exponents_out.append(alpha_out)
    return exponents1, exponents2, exponents_out

# Select a small subset of test images
test_data, test_labels = next(iter(test_loader))
test_data = test_data.view(test_data.size(0), -1)[:50]  # 50 samples
noise_scales = np.logspace(-3, 0, 15)   # ε from 0.001 to 1

alpha1, alpha2, alpha_out = compute_fractal_exponents(model, test_data, noise_scales)

# ------------------------------
# 4. Fractal Visualization: Complex plane mapping
# ------------------------------
# Real part = output exponent (α_out), Imag part = first hidden exponent (α1)
complex_points = [complex(alpha_out[i], alpha1[i]) for i in range(len(alpha_out))]

plt.figure(figsize=(10, 8))
plt.scatter([z.real for z in complex_points], [z.imag for z in complex_points], 
            c=alpha2, cmap='viridis', alpha=0.8, edgecolors='k')
plt.colorbar(label='α₂ (second hidden layer exponent)')
plt.xlabel('Real part: α_out (output layer)')
plt.ylabel('Imag part: α₁ (first hidden layer)')
plt.title('Fractal Signature of MLP on MNIST (complex plane mapping)\n'
          'Each point = one test image\n'
          'Color = exponent of second hidden layer')
plt.grid(True, alpha=0.3)
ax = plt.gca()
ax.set_aspect('equal')
plt.tight_layout()
plt.savefig('fractal_complex_mlp.png', dpi=150)
plt.show()

# ------------------------------
# 5. Interpretation and output of exponents
# ------------------------------
print("\n=== Fractal Exponents Summary ===")
print(f"Mean α₁ (first hidden): {np.mean(alpha1):.3f} ± {np.std(alpha1):.3f}")
print(f"Mean α₂ (second hidden): {np.mean(alpha2):.3f} ± {np.std(alpha2):.3f}")
print(f"Mean α_out (output):     {np.mean(alpha_out):.3f} ± {np.std(alpha_out):.3f}")
print("\nInterpretation:")
print("- α ≈ 1 : linear / smooth mapping (high precision reachable quickly)")
print("- α < 1 : fractal / sensitive mapping (residuals decay slowly, requires many iterations)")
print("- The spread across samples shows resistant information — why we can never reach 100%")