"""Flow matching on MNIST, from noise to a handwritten digit in one script.

We draw a Gaussian sample at time zero and a real image at time one, join them by a straight line,
and train a U-Net to predict the velocity of that line at a random time in between. Then we start
from fresh noise and take Euler steps along the learned field. Chapter 2, Section 2.20 of
the notes walks through this file section by section.

    python lecture_2/flow_matching_mnist.py                  # 20 epochs on MNIST, downloads on first use
    python lecture_2/flow_matching_mnist.py --steps 20       # 20 optimizer steps, for a first look
    python lecture_2/flow_matching_mnist.py --epochs 1       # a single pass, about two minutes on a CPU
    python lecture_2/flow_matching_mnist.py --dataset toy    # no download, for a first look or a test

The network class and the checkpoint keys match the published `flow_unet_mnist.pt`, so the saved
five-epoch run still loads.
"""

# %% 01. Imports
import argparse
from pathlib import Path

import torch  # Tensors, gradients, and optimization.
import torch.nn as nn  # Layers and model classes.
import torch.nn.functional as F  # Resize feature maps in the decoder.


# %% 02. Configuration
BATCH_SIZE = 128  # Images per training batch.
EPOCHS = 20  # Passes through the training set.
LEARNING_RATE = 2e-4  # Optimizer step size.
BASE_CHANNELS = 32  # U-Net widths: 32, 64, 128.
SAMPLE_STEPS = 100  # Euler steps from t=0 to t=1.
NUM_SAMPLES = 16  # Generate a 4-by-4 image grid.
SEED = 7  # Seed for initialization and training.
DATA_DIR = Path("data")  # MNIST download location.
OUTPUT_DIR = Path("outputs/flow_matching_mnist")  # Checkpoint and generated image folder.

DEVICE = torch.device(
    "cuda" if torch.cuda.is_available()  # NVIDIA GPU.
    else "mps" if torch.backends.mps.is_available()  # Apple GPU.
    else "cpu"  # CPU fallback.
)


# %% 03. Load real images
def make_loader(dataset="mnist", batch_size=BATCH_SIZE):
    """Batches of images normalized to [-1, 1], the range the velocity field is trained on."""
    from torch.utils.data import DataLoader, TensorDataset

    if dataset == "toy":  # Shapes like MNIST, drawn offline, so a test runs without a download.
        from cis6270.data import toy_images

        images = toy_images(num_images=512, size=28, seed=SEED)  # [512, 1, 28, 28].
        labels = torch.zeros(len(images), dtype=torch.long)  # Unused; generation is unconditional.
        return DataLoader(TensorDataset(images, labels), batch_size=batch_size, shuffle=True)

    from torchvision import datasets
    from torchvision.transforms import v2

    transform = v2.Compose([
        v2.ToImage(),  # Channel-first image: [1, 28, 28].
        v2.ToDtype(torch.float32, scale=True),  # Integer pixels 0-255 -> floats 0-1.
        v2.Normalize(mean=(0.5,), std=(0.5,)),  # Real data pixels: [0, 1] -> [-1, 1].
    ])
    mnist = datasets.MNIST(root=DATA_DIR, train=True, download=True, transform=transform)
    return DataLoader(
        mnist, batch_size=batch_size, shuffle=True,  # Batch shape: [B, 1, 28, 28].
        num_workers=0, pin_memory=(DEVICE.type == "cuda"),  # Load in the main process.
    )


# %% 04. Sample the conditional path
def sample_conditional_path(x1):  # x1 contains normalized real images.
    x0 = torch.randn_like(x1)  # Independent Gaussian prior; same shape as x1.
    t = torch.rand(x1.shape[0], device=x1.device)  # One random time per image; [B].
    t_image = t[:, None, None, None]  # [B, 1, 1, 1]; share time across pixels.
    xt = (1.0 - t_image) * x0 + t_image * x1  # Noise at t=0; real image at t=1.
    ut = x1 - x0  # Derivative of the straight interpolation.
    return xt, t, ut  # Network input, time, and target velocity.


# %% 05. Define a simple convolution block
def conv_block(in_channels, out_channels):
    return nn.Sequential(
        nn.Conv2d(in_channels, out_channels, 3, padding=1), nn.SiLU(),  # Keep H and W.
        nn.Conv2d(out_channels, out_channels, 3, padding=1), nn.SiLU(),  # Refine features.
    )


# %% 06. Define the small U-Net
class FlowUNet(nn.Module):
    def __init__(self, base_channels=BASE_CHANNELS):
        super().__init__()
        c = base_channels  # Base width; c=32 by default.
        self.encoder1 = conv_block(2, c)  # Image + time -> c channels.
        self.encoder2 = conv_block(c, 2 * c)  # Encoder: c -> 2c.
        self.middle = conv_block(2 * c, 4 * c)  # Bottleneck: 2c -> 4c.
        self.decoder2 = conv_block(6 * c, 2 * c)  # Upsampled 4c + skip 2c -> 2c.
        self.decoder1 = conv_block(3 * c, c)  # Upsampled 2c + skip c -> c.
        self.output = nn.Conv2d(c, 1, 1)  # One signed velocity per pixel.
        self.pool = nn.MaxPool2d(2)  # Halve the image height and width.

    # %% 07. Compute the velocity field
    def forward(self, xt, t):
        t_image = t[:, None, None, None].expand_as(xt)  # Constant time channel.
        x = torch.cat([xt, t_image], dim=1)  # [B, 2, 28, 28].
        skip1 = self.encoder1(x)  # Keep full-resolution features.
        skip2 = self.encoder2(self.pool(skip1))  # 28x28 -> 14x14.
        x = self.middle(self.pool(skip2))  # 14x14 -> 7x7.
        x = F.interpolate(x, size=skip2.shape[-2:], mode="nearest")  # 7x7 -> 14x14.
        x = self.decoder2(torch.cat([x, skip2], dim=1))  # Restore encoder detail.
        x = F.interpolate(x, size=skip1.shape[-2:], mode="nearest")  # 14x14 -> 28x28.
        x = self.decoder1(torch.cat([x, skip1], dim=1))  # Restore fine detail.
        v_pred = self.output(x)  # Predict v_theta(xt, t).
        return v_pred  # [B, 1, 28, 28], matching xt.


# %% 08. Train for one epoch
def train_one_epoch(model, loader, optimizer, max_steps=0):
    model.train()  # Select training mode.
    total_loss, num_images = 0.0, 0  # Accumulate an image-weighted epoch loss.
    for step, (x1, _) in enumerate(loader):  # Ignore digit labels, since generation is unconditional.
        if max_steps and step >= max_steps:
            break  # The step budget the caller left for this epoch is spent.
        x1 = x1.to(DEVICE)  # Move real images to the model device.
        xt, t, ut = sample_conditional_path(x1)  # Fresh noise and times for this batch.
        v_pred = model(xt, t)  # The model receives only xt and t.
        squared_error = (v_pred - ut).square()  # Velocity error at every pixel.
        loss = squared_error.flatten(1).sum(1).mean()  # Sum pixels, average images.
        optimizer.zero_grad(set_to_none=True)  # Clear old gradients.
        loss.backward()  # Compute gradients with respect to weights.
        optimizer.step()  # Update the U-Net weights.
        total_loss += loss.item() * x1.shape[0]  # Weight by the actual batch size.
        num_images += x1.shape[0]  # Include the smaller final batch.
    return total_loss / num_images  # Average training loss per image.


# %% 09. Generate images with Euler integration
# Snapshots land every quarter of the default 100-step run. The ODE states stay unclipped.
@torch.no_grad()  # Sampling runs without gradient tracking.
def sample_images(model, num_samples=NUM_SAMPLES, steps=SAMPLE_STEPS):
    if num_samples < 1 or steps < 1:
        raise ValueError("num_samples and steps must be positive")
    model.eval()  # Select evaluation mode.
    device = next(model.parameters()).device  # Use the model device.
    xt = torch.randn(num_samples, 1, 28, 28, device=device)  # Initial x0.
    dt = 1.0 / steps  # Positive time increment.
    snapshots = [xt.cpu().clone()]  # Keep the initial noise for plotting.
    for step in range(steps):  # Integrate forward in time.
        t = torch.full((num_samples,), step * dt, device=device)  # Shared time grid.
        v_pred = model(xt, t)  # Reevaluate velocity at the current state.
        xt = xt + dt * v_pred  # Euler: state += time increment * velocity.
        if (step + 1) % max(steps // 4, 1) == 0 or step + 1 == steps:
            snapshots.append(xt.cpu().clone())  # Save intermediate and final states.
    return xt, torch.stack(snapshots, dim=1)  # Raw images; [B, snapshots, 1, 28, 28].


# %% 10. Set up the complete run
def main():
    parser = argparse.ArgumentParser(description=__doc__.split("\n")[0])
    parser.add_argument("--epochs", type=int, default=EPOCHS)
    parser.add_argument(
        "--steps", type=int, default=0,
        help="optimizer steps in the whole run; 0 runs every requested epoch in full",
    )
    parser.add_argument("--batch-size", type=int, default=BATCH_SIZE)
    parser.add_argument("--dataset", choices=["mnist", "toy"], default="mnist")
    parser.add_argument("--sample-steps", type=int, default=SAMPLE_STEPS)
    parser.add_argument("--samples", type=int, default=NUM_SAMPLES)
    parser.add_argument("--seed", type=int, default=SEED)
    parser.add_argument("--out", type=str, default=str(OUTPUT_DIR))
    parser.add_argument("--quiet", action="store_true", help="print only the final report")
    args = parser.parse_args()

    torch.manual_seed(args.seed)  # Seed training randomness.
    out = Path(args.out)
    out.mkdir(parents=True, exist_ok=True)  # Create the output folder.
    loader = make_loader(args.dataset, args.batch_size)  # Download, normalize, and batch.
    model = FlowUNet().to(DEVICE)  # Create the velocity network.
    optimizer = torch.optim.AdamW(
        model.parameters(), lr=LEARNING_RATE, weight_decay=0.0,  # Optimize the flow loss.
    )
    num_parameters = sum(p.numel() for p in model.parameters())  # Count weights and biases.
    if not args.quiet:
        print(f"device={DEVICE} dataset={args.dataset} parameters={num_parameters:,}")

    # %% 11. Train and save the learned parameters
    losses = []  # Record the training loss after each epoch.
    budget = args.steps  # Optimizer steps left in the run, counted across all epochs.
    for epoch in range(1, args.epochs + 1):  # Complete the requested passes through the data.
        loss = train_one_epoch(model, loader, optimizer, budget)  # Run the training loop.
        losses.append(loss)  # Keep the loss history.
        if not args.quiet:
            print(f"epoch={epoch:02d} squared_norm_loss={loss:.3f}")
        if args.steps:  # An epoch runs the smaller of the budget and the number of batches.
            budget -= min(budget, len(loader))  # Optimizer steps this epoch used.
            if budget <= 0:
                break  # The run's budget is spent, so the remaining epochs do not run.
    torch.save({
        "model": model.state_dict(), "losses": losses,  # Weights and training curve.
        "base_channels": BASE_CHANNELS,  # Width needed to reconstruct the model.
        "seed": args.seed,  # Record the training seed.
        "epochs_completed": len(losses),  # How far this run got.
    }, out / "flow_unet_mnist.pt")

    # %% 12. Sample and save images
    torch.manual_seed(args.seed + 1)  # Fix the starting noise for generation.
    samples, trajectory = sample_images(model, args.samples, args.sample_steps)
    try:
        from torchvision.utils import save_image  # Only needed to write the PNGs.

        display_samples = ((samples.cpu() + 1.0) / 2.0).clamp(0.0, 1.0)  # Clip for display only.
        save_image(display_samples, out / "samples.png", nrow=4)  # Final images.
        display_path = ((trajectory[:4] + 1.0) / 2.0).clamp(0.0, 1.0)  # First four sample paths.
        save_image(
            display_path.flatten(0, 1), out / "trajectory.png",  # One row per sample.
            nrow=trajectory.shape[1],  # Time increases across columns.
        )
        if not args.quiet:
            print(f"wrote {out}/samples.png and {out}/trajectory.png")
    except ImportError:  # Without torchvision the run still trains and samples.
        if not args.quiet:
            print("torchvision is not installed, so no PNG was written")

    print(f"final loss {losses[-1]:.3f}  sample range "
          f"[{float(samples.min()):.2f}, {float(samples.max()):.2f}]")


# %% 13. Run the script
if __name__ == "__main__":  # Execute when launched directly.
    main()  # Load data, train, sample, and save.