"""
MNIST with a Learned Update Rule (Discrete Generalized ODE)
------------------------------------------------------------
A single shared convolutional block plays the role of the universal
update operator U_theta(S).  The network "understands" the image by
evolving the state field S via repeated discrete updates:

    S_{t+1} = S_t + U_theta(S_t)

Depth is not layer count; it is integration steps (trajectory length).
"""

import torch
import torch.nn as nn
import torch.nn.functional as F
from torchvision import datasets, transforms
from torch.utils.data import DataLoader
import matplotlib.pyplot as plt

# ----------------------------------------------------------------------
# Config
# ----------------------------------------------------------------------
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

channels   = 64      # dimension of the latent state field S(x,t)
num_steps  = 4       # integration depth T (how many update steps)
batch_size = 128
lr         = 1e-3
epochs     = 5

# ----------------------------------------------------------------------
# Data
# ----------------------------------------------------------------------
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.1307,), (0.3081,))
])

train_loader = DataLoader(
    datasets.MNIST('../data', train=True, download=True, transform=transform),
    batch_size=batch_size, shuffle=True, num_workers=2, pin_memory=True
)
test_loader = DataLoader(
    datasets.MNIST('../data', train=False, download=True, transform=transform),
    batch_size=batch_size, shuffle=False, num_workers=2, pin_memory=True
)

# ----------------------------------------------------------------------
# Model
# ----------------------------------------------------------------------
class UpdateRule(nn.Module):
    """
    The learned generalized derivative / state update operator.
    This is the 'df' of the framework: dS = U_theta(S).
    """
    def __init__(self, c):
        super().__init__()
        # Two-layer local convolution = emergent interaction kernel
        self.net = nn.Sequential(
            nn.Conv2d(c, c, kernel_size=3, padding=1, bias=False),
            nn.ReLU(inplace=True),
            nn.Conv2d(c, c, kernel_size=3, padding=1, bias=False),
        )

    def forward(self, S):
        return self.net(S)


class GeneralizedODELearner(nn.Module):
    """
    Integrates the learned update rule over T steps.
    Equivalent to solving dS/dt = U(S) with Euler discretization.
    """
    def __init__(self, channels=64, num_steps=4):
        super().__init__()
        self.num_steps = num_steps

        # Encoder: raw image -> initial field S(t=0)
        self.encoder = nn.Sequential(
            nn.Conv2d(1, channels, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
        )

        # Shared update rule (single law, applied recurrently)
        self.update = UpdateRule(channels)

        # Decoder / collapse-to-prediction readout
        self.pool = nn.AdaptiveAvgPool2d(1)
        self.fc = nn.Linear(channels, 10)

    def forward(self, x):
        # Initial state S_0
        S = self.encoder(x)

        # Generalized integral: accumulate updates
        # S_final = S_0 + sum_{t=1}^{T} dS_t
        for _ in range(self.num_steps):
            dS = self.update(S)   # df = U(S)  [the derivative is the update]
            S = S + dS            # state evolution (discrete Euler step)

        # Collapse field to logits
        out = self.pool(S).view(S.size(0), -1)
        out = self.fc(out)
        return out


# ----------------------------------------------------------------------
# Train & Test
# ----------------------------------------------------------------------
model = GeneralizedODELearner(channels=channels, num_steps=num_steps).to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=lr)
criterion = nn.CrossEntropyLoss()

train_losses = []
test_accs = []

print("Training...")
for epoch in range(epochs):
    model.train()
    for batch_idx, (data, target) in enumerate(train_loader):
        data, target = data.to(device), target.to(device)

        optimizer.zero_grad()
        output = model(data)
        loss = criterion(output, target)
        loss.backward()
        optimizer.step()

        train_losses.append(loss.item())

    # ---- evaluate ----
    model.eval()
    correct = 0
    total = 0
    with torch.no_grad():
        for data, target in test_loader:
            data, target = data.to(device), target.to(device)
            logits = model(data)
            pred = logits.argmax(dim=1)
            total += target.size(0)
            correct += (pred == target).sum().item()

    acc = 100.0 * correct / total
    test_accs.append(acc)
    print(f"Epoch {epoch+1}/{epochs}  |  Test Accuracy: {acc:.2f}%")

# ----------------------------------------------------------------------
# Visualize
# ----------------------------------------------------------------------
fig, ax = plt.subplots(1, 2, figsize=(11, 4))

ax[0].plot(train_losses, color='crimson', lw=1.2)
ax[0].set_title('Training Loss (per batch)')
ax[0].set_xlabel('Step')
ax[0].set_ylabel('Cross-Entropy')
ax[0].grid(True, alpha=0.3)

ax[1].plot(range(1, epochs + 1), test_accs, marker='o', color='teal', lw=2)
ax[1].set_title('Test Accuracy')
ax[1].set_xlabel('Epoch')
ax[1].set_ylabel('%')
ax[1].set_ylim(90, 100)
ax[1].grid(True, alpha=0.3)

fig.suptitle(f'MNIST via Learned Update Rule  |  {num_steps} integration steps  |  {channels} channels',
             fontsize=13)
plt.tight_layout()
plt.show()

# ----------------------------------------------------------------------
# Inspect: how does the state field evolve?
# ----------------------------------------------------------------------
@torch.no_grad()
def trace_feature_trajectory(model, image):
    """
    Returns the mean activation (proxy for 'energy' in the field)
    at each integration step.
    """
    model.eval()
    image = image.unsqueeze(0).to(device)
    S = model.encoder(image)
    means = [S.mean().item()]
    for _ in range(model.num_steps):
        S = S + model.update(S)
        means.append(S.mean().item())
    return means


# Grab one image and watch the generalized integral accumulate
x_batch, _ = next(iter(test_loader))
traj = trace_feature_trajectory(model, x_batch[0])
print("\nState field mean activation across integration steps:")
print("  t=0 (initial):  ", f"{traj[0]:+.4f}")
for i in range(1, len(traj)):
    delta = traj[i] - traj[i-1]
    print(f"  t={i} (update):  {traj[i]:+.4f}  (dS = {delta:+.4f})")