"""
FLUX MNIST Classifier
=====================

A Neural ODE implementation of the FLUX architecture where:
- Encoder  = LOAD instruction (image -> initial register values)
- Dynamics = Instruction Field (program counter t flows from 0 to 1)
  Different instruction windows are active at different program counter
  positions, exactly like the FLUX eta_i(p) blending.
- Integrator = EXECUTE (RK4 ODE solver)
- Decoder  = READ (extract classification register at HALT)

This is a depth-variant (non-autonomous) Neural ODE where the
vector field explicitly changes along the program counter axis.
"""

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torchvision import datasets, transforms
from torch.utils.data import DataLoader

# Device setup
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')


# =============================================================================
# FLUX DYNAMICS: The Instruction Field (the Vector Field F)
# =============================================================================
class FLUXDynamics(nn.Module):
    """
    Defines the vector field dh/dt = f(t, h) where t is the program counter.
    In FLUX terms, this is the blended instruction field:
        F(h) = sum_i eta_i(t) * F_i(h)
    where eta_i(t) are smooth instruction windows centered at different
    depths of the program.
    """
    def __init__(self, hidden_dim: int, num_instructions: int = 4):
        super().__init__()
        self.hidden_dim = hidden_dim
        self.num_instructions = num_instructions

        # Multiple instruction modules (different operations at different depths)
        self.instructions = nn.ModuleList([
            nn.Sequential(
                nn.Linear(hidden_dim, hidden_dim),
                nn.Tanh(),
                nn.Linear(hidden_dim, hidden_dim)
            ) for _ in range(num_instructions)
        ])

        # Instruction window centers (program counter positions)
        centers = torch.linspace(0.0, 1.0, num_instructions)
        self.register_buffer('centers', centers)
        self.register_buffer('width', torch.tensor(0.25))

    def bump_window(self, t):
        """Smooth instruction window function (Gaussian bump)."""
        t = torch.as_tensor(t, device=self.centers.device, dtype=self.centers.dtype)
        weights = torch.exp(-((t - self.centers) / self.width) ** 2)
        weights = weights / (weights.sum() + 1e-8)
        return weights  # [num_instructions]

    def forward(self, t, h):
        """
        t : scalar program counter position (time)
        h : [batch_size, hidden_dim] register state
        """
        # Get instruction window weights (blend coefficients)
        weights = self.bump_window(t)  # [num_instructions]

        # Evaluate each instruction field on the current registers
        contributions = torch.stack([instr(h) for instr in self.instructions], dim=0)
        # -> [num_instructions, batch_size, hidden_dim]

        # Blend contributions: FLUX F(z) = sum_i eta_i(p) * F_i(z)
        w = weights.view(-1, 1, 1)  # [num_instructions, 1, 1]
        dhdt = (w * contributions).sum(dim=0)  # [batch_size, hidden_dim]

        return dhdt


# =============================================================================
# FLUX ENCODER: LOAD Instruction
# =============================================================================
class FLUXEncoder(nn.Module):
    """
    Loads the MNIST image into the register file.
    Maps 784 pixels -> hidden_dim initial register values.
    """
    def __init__(self, input_dim: int = 784, hidden_dim: int = 64):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(input_dim, hidden_dim * 2),
            nn.ReLU(),
            nn.Linear(hidden_dim * 2, hidden_dim)
        )

    def forward(self, x):
        # x: [batch, 1, 28, 28]
        x = x.view(x.size(0), -1)
        return self.net(x)


# =============================================================================
# FLUX DECODER: READ Instruction
# =============================================================================
class FLUXDecoder(nn.Module):
    """
    Reads the classification register at HALT (program counter = 1.0).
    Maps hidden registers -> 10 digit classes.
    """
    def __init__(self, hidden_dim: int = 64, num_classes: int = 10):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(hidden_dim, hidden_dim // 2),
            nn.ReLU(),
            nn.Linear(hidden_dim // 2, num_classes)
        )

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


# =============================================================================
# FLUX INTEGRATOR: Program Execution (ODE Solver)
# =============================================================================
class FLUXIntegrator:
    """
    Executes the FLUX program by integrating the ODE from t=0 to t=1.
    Program counter flows forward (sequential execution).
    Uses RK4 for stable trajectory evolution.
    """
    def __init__(self, dynamics, t_start: float = 0.0, t_end: float = 1.0, num_steps: int = 20):
        self.dynamics = dynamics
        self.t_start = t_start
        self.t_end = t_end
        self.num_steps = num_steps
        self.dt = (t_end - t_start) / num_steps

    def integrate(self, h0):
        """
        h0: [batch_size, hidden_dim] initial register state
        Returns hT: final register state at HALT (t=1.0)
        """
        h = h0
        t = self.t_start

        for _ in range(self.num_steps):
            # RK4 integration step
            k1 = self.dynamics(t, h)
            k2 = self.dynamics(t + self.dt / 2, h + self.dt * k1 / 2)
            k3 = self.dynamics(t + self.dt / 2, h + self.dt * k2 / 2)
            k4 = self.dynamics(t + self.dt, h + self.dt * k3)

            h = h + (self.dt / 6.0) * (k1 + 2 * k2 + 2 * k3 + k4)
            t = t + self.dt

        return h


# =============================================================================
# FLUX MNIST CLASSIFIER: Complete Program
# =============================================================================
class FLUXMNISTClassifier(nn.Module):
    """
    Full FLUX program execution:
        1. LOAD  : Image -> Registers
        2. EXEC : Program counter 0 -> 1 (flow through instruction field)
        3. HALT : Read classification register
    """
    def __init__(self, input_dim: int = 784, hidden_dim: int = 64,
                 num_classes: int = 10, num_instructions: int = 4):
        super().__init__()
        self.encoder = FLUXEncoder(input_dim, hidden_dim)
        self.dynamics = FLUXDynamics(hidden_dim, num_instructions)
        self.decoder = FLUXDecoder(hidden_dim, num_classes)
        self.integrator = FLUXIntegrator(self.dynamics, t_start=0.0, t_end=1.0, num_steps=20)

    def forward(self, x):
        # Step 1: LOAD (initialize registers from input image)
        h0 = self.encoder(x)

        # Step 2: EXECUTE (program counter flows from 0 to 1 through instruction field)
        hT = self.integrator.integrate(h0)

        # Step 3: HALT (read output register at program counter = 1.0)
        logits = self.decoder(hT)
        return logits


# =============================================================================
# Training & Testing Routines
# =============================================================================
def train_epoch(model, loader, optimizer, criterion):
    model.train()
    total_loss = 0.0
    correct = 0
    total = 0

    for batch_idx, (data, target) in enumerate(loader):
        data, target = data.to(device), target.to(device)

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

        total_loss += loss.item()
        pred = output.argmax(dim=1)
        correct += pred.eq(target).sum().item()
        total += target.size(0)

        if batch_idx % 200 == 0:
            print(f'  Batch {batch_idx:3d}/{len(loader)}  Loss: {loss.item():.4f}')

    return total_loss / len(loader), 100.0 * correct / total


def test_epoch(model, loader, criterion):
    model.eval()
    test_loss = 0.0
    correct = 0
    total = 0

    with torch.no_grad():
        for data, target in loader:
            data, target = data.to(device), target.to(device)
            output = model(data)
            test_loss += criterion(output, target).item()
            pred = output.argmax(dim=1)
            correct += pred.eq(target).sum().item()
            total += target.size(0)

    avg_loss = test_loss / len(loader)
    accuracy = 100.0 * correct / total
    return avg_loss, accuracy


# =============================================================================
# Main Entry Point
# =============================================================================
if __name__ == '__main__':
    # Hyperparameters
    BATCH_SIZE = 128
    HIDDEN_DIM = 64
    NUM_INSTRUCTIONS = 5
    EPOCHS = 5
    LR = 1e-3

    # MNIST dataset
    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=BATCH_SIZE, shuffle=True)
    test_loader = DataLoader(test_dataset, batch_size=BATCH_SIZE, shuffle=False)

    # Build FLUX model
    model = FLUXMNISTClassifier(
        input_dim=784,
        hidden_dim=HIDDEN_DIM,
        num_classes=10,
        num_instructions=NUM_INSTRUCTIONS
    ).to(device)

    criterion = nn.CrossEntropyLoss()
    optimizer = optim.Adam(model.parameters(), lr=LR)

    # Print architecture summary
    total_params = sum(p.numel() for p in model.parameters())
    print("=" * 50)
    print("FLUX MNIST Classifier")
    print("=" * 50)
    print(f"Hidden dimension     : {HIDDEN_DIM}")
    print(f"Instruction windows  : {NUM_INSTRUCTIONS}")
    print(f"Program counter      : [0.0, 1.0] (LOAD to HALT)")
    print(f"Integrator steps     : 20 (RK4)")
    print(f"Total parameters     : {total_params:,}")
    print(f"Device               : {device}")
    print("=" * 50)

    for epoch in range(1, EPOCHS + 1):
        print(f"\nEpoch {epoch}/{EPOCHS}")
        train_loss, train_acc = train_epoch(model, train_loader, optimizer, criterion)
        test_loss, test_acc = test_epoch(model, test_loader, criterion)

        print(f"  Train -> Loss: {train_loss:.4f} | Acc: {train_acc:.2f}%")
        print(f"  Test  -> Loss: {test_loss:.4f} | Acc: {test_acc:.2f}%")

    print("\nTraining complete. The FLUX program has optimized its vector field.")