"""The three residuals of a flow map, and the diagonal anchor added to each one.

The cocycle differentiated in the end time is the Lagrangian identity, differentiated in the start
time it is the Eulerian identity, and left finite it is the composition identity. Each becomes a
loss by squaring the amount the model violates it. None of the three mentions the data, so a model
that is self-consistent and wrong satisfies all of them. A flow matching loss on the diagonal
supplies the one piece of information about the target.

Chapter 6, Sections 6.15 to 6.17 and 6.23 of the notes. Trained on the four Gaussian blobs
of ``cis6270.data.gaussian_mixture_2d``. Prints the stay-still residuals of Section 6.17, 1, 1
and 0, and then four trained models side by side in one-step sampling.

    python lecture_6/residuals.py                 # four models, 600 steps each
    python lecture_6/residuals.py --steps 20      # a few seconds, for a first look
"""

# %% 01. Imports and configuration
import torch

from cis6270.data import GAUSSIAN_STD, gaussian_mixture_2d
from cis6270.metrics import mode_report
from cis6270.nets import TwoTimeMLP, count_parameters
from cis6270.runner import common_parser, pick_device, print_report, save_run, set_seed, train


torch.set_num_threads(1)  # These networks are small, and splitting one matrix product costs more
# than it saves. ``tests/conftest.py`` sets the same thing for the test session, and this line
# acts when a reader runs the script directly.

ON_MODE = 3.0 * GAUSSIAN_STD  # A sample counts as on a mode within three data standard deviations.
ODE_STEPS = 128  # Substeps used to integrate the model's own diagonal velocity accurately.


# %% 02. The two-time model
#
# The network produces the average velocity v_{s,t}(x). The map is then the state plus the elapsed
# time times that velocity. At t = s the second term vanishes whatever the weights are, so the loss
# spends no term on the identity axiom.
class Map(torch.nn.Module):
    """X_{s,t}(x) = x + (t - s) v_{s,t}(x), with v the two-time network of ``cis6270.nets``."""

    def __init__(self, dim: int = 2, width: int = 128):
        super().__init__()
        self.net = TwoTimeMLP(dim=dim, width=width)  # Reads x, s and t; returns a vector.

    def v(self, x, s, t):  # the network
        return self.net(x, s, t)  # [B, d] the average velocity over [s, t].

    def X(self, x, s, t):  # the map, where the identity axiom holds
        return self.net.flow_map(x, s, t)  # [B, d] = x + (t - s) v_{s,t}(x).


# %% 03. The diagonal anchor and the tangent axiom
#
# I_t = (1-t) x_0 + t x_1 and its derivative x_1 - x_0 are available in closed form, so the anchor
# costs one network call and no simulation. Its minimizer is E[I_dot | I_t], the true velocity.
def anchor(m, x0, x1, t):  # tangent axiom
    xt = (1 - t[:, None]) * x0 + t[:, None] * x1  # [B, d] the interpolant at time t.
    return ((m.v(xt, t, t) - (x1 - x0)) ** 2).mean()


# %% 04. The Lagrangian residual
#
# d/dt X = b at the output. The closure holds the state and the start time fixed, and a tangent of
# one in t returns the map and its end-time derivative from a single forward-mode pass. The target
# is the model's own diagonal velocity, detached, so the diagonal trains only through the anchor.
def lagrangian(m, xs, s, t):  # d/dt X = b at the output
    f = lambda tt: m.X(xs, s, tt)
    X, dXdt = torch.func.jvp(f, (t,), (torch.ones_like(t),))
    return ((dXdt - m.v(X, t, t).detach()) ** 2).mean()


# %% 05. The Eulerian residual
#
# d/ds X + b . grad X = 0. Both the state and the start time are open in the closure, and the
# tangent (v, 1) makes the one returned directional derivative the whole left-hand side. The
# identity is already written as zero, so the residual needs no target at all.
def eulerian(m, xs, s, t):  # d/ds X + b . grad X = 0
    vs = m.v(xs, s, s).detach()
    f = lambda xx, ss: m.X(xx, ss, t)
    _, dX = torch.func.jvp(f, (xs, s), (vs, torch.ones_like(s)))
    return (dX ** 2).mean()


# %% 06. The composition residual
#
# Two short maps against one long one, with no derivative of any kind. The detach covers the whole
# composed target, so the two legs are frozen and only the long map is corrected.
def composition(m, xs, s, u, t):  # two short maps, one long one
    tgt = m.X(m.X(xs, s, u), u, t).detach()
    return ((m.X(xs, s, t) - tgt) ** 2).mean()


# %% 07. The stay-still solution each residual has to rule out
#
# The average velocity equals the true b when the two times coincide and zero everywhere else, so
# the map is the identity for every pair of times and the sampler generates the source. The
# identity axiom, the cocycle and the anchor all hold on it, and the two differential residuals
# separate it from the truth by a full unit.
class StayStill(torch.nn.Module):
    """v_{s,t}(x) = b(x,t) on the diagonal and 0 off it, on the scalar flow dx/dt = x."""

    def v(self, x, s, t):
        on_diagonal = (s == t)[:, None]  # [B, 1] true where the two times meet.
        return torch.where(on_diagonal, x, torch.zeros_like(x))  # b(x,t) = x on this flow.

    def X(self, x, s, t):
        return x  # Off the diagonal (t-s) times zero, on it zero times b, so the map is the identity.


def stay_still_residuals():
    """The three residuals of Section 6.17 at x = 1, s = 0, t = 1, or 1, 1 and 0."""
    m = StayStill()
    x = torch.ones(1, 1)  # [1, 1] the state x = 1.
    s, t, u = torch.zeros(1), torch.ones(1), torch.full((1,), 0.5)
    diagonal_gap = float(((m.v(x, t, t) - x) ** 2).mean())  # The anchor, at its minimum.
    return {
        "stay-still Lagrangian": float(lagrangian(m, x, s, t)),
        "stay-still Eulerian": float(eulerian(m, x, s, t)),
        "stay-still composition": float(composition(m, x, s, u, t)),
        "stay-still diagonal gap": diagonal_gap,
    }


# %% 08. One training step, the same shape whichever residual we add
def make_loss(model, data, residual, batch_size, device):
    """Draw a coupling, draw times from the triangle s <= t, and add the residual to the anchor."""

    def loss_fn(step: int) -> torch.Tensor:
        index = torch.randint(0, len(data), (batch_size,), device=device)
        x1 = data[index]  # [B, 2] the data draw.
        x0 = torch.randn_like(x1)  # [B, 2] the independent source draw.
        s = torch.rand(batch_size, device=device)  # [B] start time.
        t = s + (1.0 - s) * torch.rand(batch_size, device=device)  # [B] end time, s <= t <= 1.
        u = s + (t - s) * torch.rand(batch_size, device=device)  # [B] intermediate time.
        total = anchor(model, x0, x1, torch.rand(batch_size, device=device))  # The data term.
        if residual is None:
            return total  # The diagonal alone, flow matching read in one step.
        xs = (1 - s[:, None]) * x0 + s[:, None] * x1  # [B, 2] the interpolant at the start time.
        if residual is composition:
            return total + residual(model, xs, s, u, t)
        return total + residual(model, xs, s, t)

    return loss_fn


# %% 09. One-step sampling, and the model's own ordinary differential equation
@torch.no_grad()
def evaluate(model, x0):
    """One-step output, the endpoint of the model's own diagonal velocity, and the mode statistics."""
    from cis6270.data import GAUSSIAN_CENTERS

    zeros = torch.zeros(len(x0), device=x0.device)
    one_step = model.X(x0, zeros, torch.ones_like(zeros))  # [N, 2] a single evaluation of the map.
    x, h = x0.clone(), 1.0 / ODE_STEPS  # Integrate dx/dt = v_{t,t}(x) with the same weights.
    for k in range(ODE_STEPS):
        t = torch.full((len(x0),), k * h, device=x0.device)
        x = x + h * model.v(x, t, t)  # Euler on the diagonal velocity the anchor trains.
    distance = float((one_step - x).norm(dim=1).mean())  # Column two of Table 6.2.
    nearest = torch.cdist(one_step, GAUSSIAN_CENTERS.to(one_step)).min(dim=1).values  # [N].
    on_mode = float((nearest < ON_MODE).to(torch.get_default_dtype()).mean())  # Column three of Table 6.2.
    return one_step, distance, on_mode, mode_report(one_step)


# %% 10. The report
def main() -> None:
    parser = common_parser(__doc__.split("\n")[0])
    parser.set_defaults(samples=512)
    args = parser.parse_args()
    set_seed(args.seed)
    device = pick_device(args.device)
    data = gaussian_mixture_2d(num_points=2048, seed=args.seed).to(device)  # [2048, 2] the target.

    report = stay_still_residuals()  # Section 6.17, before any training.
    if not args.quiet:
        print(f"stay-still at x=1, s=0, t=1: Lagrangian {report['stay-still Lagrangian']:.3f}  "
              f"Eulerian {report['stay-still Eulerian']:.3f}  "
              f"composition {report['stay-still composition']:.3f}")

    recipes = [("the diagonal alone", None), ("+ Lagrangian", lagrangian),
               ("+ Eulerian", eulerian), ("+ composition", composition)]
    samples: list[str] = []
    for name, residual in recipes:
        set_seed(args.seed)  # Same starting weights, same batches and same noise for all four.
        model = Map(dim=2, width=args.width).to(device)
        if not args.quiet:
            print(f"\ntraining {name}  parameters={count_parameters(model):,}")
        losses = train(model, make_loss(model, data, residual, args.batch_size, device),
                       steps=args.steps, lr=args.lr, quiet=args.quiet)
        torch.manual_seed(args.seed + 1)  # The same source points for every model.
        x0 = torch.randn(args.samples, 2, device=device)
        _, distance, on_mode, stats = evaluate(model, x0)
        report[f"{name}: distance to its own ODE"] = distance
        report[f"{name}: fraction on a mode"] = on_mode
        report[f"{name}: final loss"] = sum(losses[-20:]) / len(losses[-20:])
        samples.append(f"{name:<20s} distance {distance:.3f}  on a mode {on_mode:.2f}  "
                       f"modes found {stats['modes_found']}")

    print_report("Three residuals and the diagonal anchor, Chapter 6", report, samples)
    save_run(args.out, config=vars(args), report=report, samples=samples)


if __name__ == "__main__":
    main()