# examples/flowmatching.py
"""Seeded offline CPU mechanics for Lipman et al., Eq. 20–23.

Synthetic targets and a reduced MLP illustrate CFM, not image generation or
benchmark reproduction. Importing this module performs no training or sampling.
"""

import argparse
import hashlib
import json
import math
from pathlib import Path

import torch
from torch import nn


def _positive_int(value, name):
    """Reject empty dimensions and nonintegral counts."""
    if isinstance(value, bool) or not isinstance(value, int) or value < 1:
        raise ValueError(f"{name} must be a positive integer")


def _state(x):
    """Validate a finite floating CPU state [B, D]."""
    if (
        not isinstance(x, torch.Tensor)
        or x.ndim != 2
        or min(x.shape) < 1
        or not x.is_floating_point()
        or x.device.type != "cpu"
        or not torch.isfinite(x).all()
    ):
        raise ValueError("state must be finite floating CPU [B, D]")


def _sigma(sigma_min):
    """Keep Gaussian conditional densities positive, including at t=1."""
    if (
        isinstance(sigma_min, bool)
        or not isinstance(sigma_min, (float, int))
        or not 0 < sigma_min <= 1
    ):
        raise ValueError("sigma_min must be in (0, 1]")


def _pair(x0, x1):
    """Require matching endpoint shape and dtype."""
    _state(x0)
    _state(x1)
    if x0.shape != x1.shape or x0.dtype != x1.dtype:
        raise ValueError("endpoint shape and dtype must match")


def _time(x, t):
    """Validate one finite time per item, without accidental broadcasting."""
    if (
        not isinstance(t, torch.Tensor)
        or t.shape != (x.shape[0], 1)
        or t.dtype != x.dtype
        or t.device != x.device
        or not torch.isfinite(t).all()
        or not ((t >= 0) & (t <= 1)).all()
    ):
        raise ValueError(
            "time must match state dtype/device, shape [B, 1], in [0, 1]"
        )


# region: draw_endpoints
def draw_endpoints(batch=16, dimension=2, seed=7):
    """Draw independent Gaussian noise, synthetic two-mode data, and time."""
    _positive_int(batch, "batch")
    _positive_int(dimension, "dimension")
    generator = torch.Generator(device="cpu").manual_seed(seed)
    x0 = torch.randn(batch, dimension, generator=generator)
    modes = 2 * torch.randint(0, 2, (batch, 1), generator=generator) - 1
    x1 = 2 * modes + 0.2 * torch.randn(batch, dimension, generator=generator)
    t = torch.rand(batch, 1, generator=generator)
    return x0, x1, t


# endregion: draw_endpoints


# region: conditional_path
def conditional_path(x0, x1, t, sigma_min=0.05):
    """Evaluate Eq. 22; at t=1 the endpoint is x1 + sigma_min*x0."""
    _pair(x0, x1)
    _time(x0, t)
    _sigma(sigma_min)
    a = 1 - (1 - sigma_min) * t
    return a * x0 + t * x1


# endregion: conditional_path


# region: conditional_target
def conditional_target(x0, x1, sigma_min=0.05):
    """Differentiate Eq. 22 along a fixed conditional trajectory."""
    _pair(x0, x1)
    _sigma(sigma_min)
    return x1 - (1 - sigma_min) * x0


# endregion: conditional_target


# region: velocity_model
class VelocityModel(nn.Module):
    """Reduced field v(x,t); the unknown target x1 is never an input."""

    def __init__(self, dimension=2, hidden=32):
        super().__init__()
        _positive_int(dimension, "dimension")
        _positive_int(hidden, "hidden")
        self.dimension = dimension
        self.layers = nn.Sequential(
            nn.Linear(dimension + 1, hidden),
            nn.Tanh(),
            nn.Linear(hidden, dimension),
        )

    def forward(self, x, t):
        """Map state [B,D] and time [B,1] to velocity [B,D]."""
        _state(x)
        _time(x, t)
        if x.shape[1] != self.dimension:
            raise ValueError("state dimension must match model")
        return self.layers(torch.cat((x, t), dim=1))


# endregion: velocity_model


# region: cfm_loss
def cfm_loss(model, x0, x1, t, sigma_min=0.05):
    """Use Eq. 23's squared vector norm, then average over the batch."""
    xt = conditional_path(x0, x1, t, sigma_min)
    target = conditional_target(x0, x1, sigma_min)
    prediction = model(xt, t)
    if (
        not isinstance(prediction, torch.Tensor)
        or prediction.shape != target.shape
        or prediction.dtype != target.dtype
        or prediction.device != target.device
        or not torch.isfinite(prediction).all()
    ):
        raise ValueError("field must return finite matching [B,D] velocities")
    return (prediction - target).square().sum(dim=1).mean()


def update_field(model, optimizer, x0, x1, t, sigma_min=0.05):
    """Perform a real optimizer step and report the preceding gradient norm."""
    optimizer.zero_grad(set_to_none=True)
    loss = cfm_loss(model, x0, x1, t, sigma_min)
    loss.backward()
    gradients = [p.grad for p in model.parameters()]
    if any(g is None or not torch.isfinite(g).all() for g in gradients):
        raise ValueError("all parameter gradients must be finite")
    grad_norm = torch.sqrt(sum(g.square().sum() for g in gradients))
    optimizer.step()
    return float(loss.detach()), float(grad_norm)


# endregion: cfm_loss


# region: euler_integrator
@torch.no_grad()
def euler_integrator(field, x0, steps=32):
    """Integrate dx/dt=field(x,t) with pedagogical Euler step h=1/steps."""
    _state(x0)
    _positive_int(steps, "steps")
    x = x0.clone()
    trajectory = [x.clone()]
    h = 1 / steps
    for k in range(steps):
        t = x.new_full((x.shape[0], 1), k * h)
        velocity = field(x, t)
        if (
            not isinstance(velocity, torch.Tensor)
            or velocity.shape != x.shape
            or velocity.dtype != x.dtype
            or velocity.device != x.device
            or not torch.isfinite(velocity).all()
        ):
            raise ValueError(
                "field must return finite matching [B,D] velocities"
            )
        x = x + h * velocity
        _state(x)
        trajectory.append(x.clone())
    return torch.stack(trajectory)


# endregion: euler_integrator


# region: compare_paths
@torch.no_grad()
def compare_paths(model, x0, x1, sigma_min=0.05, steps=32):
    """Keep conditional training paths separate from field-only sampling."""
    _positive_int(steps, "steps")
    conditional = torch.stack(
        [
            conditional_path(
                x0, x1, x0.new_full((len(x0), 1), k / steps), sigma_min
            )
            for k in range(steps + 1)
        ]
    )
    generated = euler_integrator(model, x0, steps)
    return {"conditional_training": conditional, "generated": generated}


# endregion: compare_paths


def smoke(seed=7):
    """Execute training and independent analytic checks; return finite JSON."""
    torch.manual_seed(seed)
    torch.set_num_threads(1)
    x0, x1, t = draw_endpoints(seed=seed)
    sigma_min = 0.05
    model = VelocityModel()
    optimizer = torch.optim.SGD(model.parameters(), lr=0.02)
    before = torch.cat(
        [p.detach().flatten() for p in model.parameters()]
    ).clone()
    initial_loss = float(cfm_loss(model, x0, x1, t).detach())
    gradient_norms = []
    for _ in range(64):
        _, norm = update_field(model, optimizer, x0, x1, t, sigma_min)
        gradient_norms.append(norm)
    final_loss = float(cfm_loss(model, x0, x1, t).detach())
    delta = (
        torch.cat([p.detach().flatten() for p in model.parameters()]) - before
    )
    paths = compare_paths(model, x0, x1, sigma_min)
    xd, yd = x0.double(), x1.double()
    td = t.double()
    eps = 1e-5
    derivative = (
        conditional_path(xd, yd, td + eps, sigma_min)
        - conditional_path(xd, yd, td - eps, sigma_min)
    ) / (2 * eps)
    derivative_error = float(
        (derivative - conditional_target(xd, yd)).abs().max()
    )
    endpoint = conditional_path(xd, yd, td.new_ones(td.shape), sigma_min)
    endpoint_error = float((endpoint - (yd + sigma_min * xd)).abs().max())
    constant = euler_integrator(lambda x, time: torch.full_like(x, 2), xd, 32)
    growth32 = euler_integrator(lambda x, time: x, xd, 32)
    growth64 = euler_integrator(lambda x, time: x, xd, 64)
    constant_error = float((constant[-1] - (xd + 2)).abs().max())
    growth_errors = [
        float((p[-1] - math.e * xd).abs().max()) for p in (growth32, growth64)
    ]
    assert final_loss < initial_loss
    assert float(delta.norm()) > 0 and all(
        math.isfinite(g) and g > 0 for g in gradient_norms
    )
    assert derivative_error < 1e-8 and endpoint_error < 1e-12
    assert constant_error < 1e-12 and 0 < growth_errors[1] < growth_errors[0]
    assert not torch.allclose(
        paths["conditional_training"], paths["generated"]
    )
    observations = {
        "source_sha256": hashlib.sha256(
            Path(__file__).read_bytes()
        ).hexdigest(),
        "seed": seed,
        "boundary": "offline CPU synthetic mechanics; no pretrained model or benchmark",
        "sigma_min": sigma_min,
        "model_parameters": {
            name: value.detach().tolist()
            for name, value in model.named_parameters()
        },
        "first_sample": {
            "x0": x0[0].tolist(),
            "x1": x1[0].tolist(),
            "generated_trajectory": paths["generated"][:, 0].tolist(),
        },
        "reductions": {"cfm": "sum squared velocity components, mean batch"},
        "shapes": {
            "x0": list(x0.shape),
            "x1": list(x1.shape),
            "time": list(t.shape),
            "velocity": list(model(x0, t).shape),
            **{key: list(value.shape) for key, value in paths.items()},
        },
        "training": {
            "updates": 64,
            "learning_rate": 0.02,
            "initial_loss": initial_loss,
            "final_loss": final_loss,
            "gradient_norm_min": min(gradient_norms),
            "gradient_norm_max": max(gradient_norms),
            "finite_gradients": True,
            "parameter_delta_norm": float(delta.norm()),
        },
        "observations": {
            "path_derivative_max_error": derivative_error,
            "residual_endpoint_max_error": endpoint_error,
            "constant_euler_max_error": constant_error,
            "growth_euler_errors_32_64": growth_errors,
            "generated_endpoint": paths["generated"][-1, 0].tolist(),
            "conditional_endpoint": paths["conditional_training"][
                -1, 0
            ].tolist(),
            "generated_conditional_max_difference": float(
                (paths["generated"] - paths["conditional_training"])
                .abs()
                .max()
            ),
            "euler_steps": 32,
            "euler_step_size": 1 / 32,
        },
    }
    json.dumps(observations, allow_nan=False)
    return observations


def main():
    """Print smoke observations only when explicitly requested."""
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--smoke", action="store_true")
    parser.add_argument("--seed", type=int, default=7)
    args = parser.parse_args()
    if not args.smoke:
        parser.error("use --smoke for the offline mechanics run")
    print(json.dumps(smoke(args.seed), allow_nan=False, sort_keys=True))


if __name__ == "__main__":
    main()