"""Lesson 11: conditional flow matching of populations, not biological time."""
import argparse
import json
from pathlib import Path
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np
import torch
from torch import nn
from common import make_dataset, prepare_embedding


def run(output_dir: Path) -> dict:
    output_dir.mkdir(parents=True, exist_ok=True)
    torch.manual_seed(7)
    torch.set_num_threads(1)
    rng = np.random.default_rng(7)
    data = make_dataset()
    _, Z, _, _ = prepare_embedding(data)
    mean, std = Z[data["train"]].mean(0), Z[data["train"]].std(0)
    Z = (Z - mean) / std
    early = data["time"] <= .15
    late = data["time"] >= .85
    A = torch.tensor(Z[early & data["train"]], dtype=torch.float32)
    B = torch.tensor(Z[late & data["train"]], dtype=torch.float32)
    field = nn.Sequential(nn.Linear(3, 64), nn.Tanh(), nn.Linear(64, 64),
                          nn.Tanh(), nn.Linear(64, 2))
    optimizer = torch.optim.Adam(field.parameters(), lr=.002)
    losses = []
    for step in range(700):
        # Independent endpoint pairs; this is not an OT coupling.
        x0 = A[torch.randint(len(A), (128,))]
        x1 = B[torch.randint(len(B), (128,))]
        t = torch.rand(128, 1)
        xt = (1 - t) * x0 + t * x1
        loss = ((field(torch.cat([xt, t], dim=1)) - (x1 - x0)) ** 2).mean()
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        losses.append(float(loss.detach()))
    test_A, test_B = Z[early & ~data["train"]], Z[late & ~data["train"]]
    initial = test_A[rng.integers(len(test_A), size=400)]
    target = test_B[rng.integers(len(test_B), size=400)]
    @torch.no_grad()
    def rollout(steps):
        x = torch.tensor(initial, dtype=torch.float32)
        trace = [x.numpy().copy()]
        for k in range(steps):
            t = torch.full((len(x), 1), k / steps)
            x = x + field(torch.cat([x, t], dim=1)) / steps
            trace.append(x.numpy().copy())
        return np.asarray(trace)
    trace = rollout(80)
    refined = rollout(160)
    directions = rng.normal(size=(2, 64))
    directions /= np.linalg.norm(directions, axis=0)
    def sliced_w1(x, y):
        return float(np.mean(abs(np.sort(x @ directions, axis=0)
                                 - np.sort(y @ directions, axis=0))))
    fig, axes = plt.subplots(1, 2, figsize=(10, 4.2))
    axes[0].plot(losses, color="#167b8c", alpha=.7)
    axes[0].set(xlabel="Training update", ylabel="Conditional velocity MSE",
                title="Regression loss (not a trajectory error)")
    axes[1].scatter(*target.T, s=8, c="#dc5a32", alpha=.35, label="Held-out target")
    axes[1].scatter(*trace[-1].T, s=7, c="#167b8c", alpha=.4, label="Generated endpoint")
    for i in range(0, 400, 25):
        axes[1].plot(trace[:, i, 0], trace[:, i, 1], color="gray", alpha=.3)
    axes[1].set(xlabel="Standardized PC 1", ylabel="Standardized PC 2",
                title="Artificial interpolation time 0 to 1")
    axes[1].legend(fontsize=8)
    fig.tight_layout()
    fig.savefig(output_dir / "lesson11.png", dpi=160)
    plt.close(fig)
    return {"final_50_loss_mean": float(np.mean(losses[-50:])),
            "endpoint_sliced_w1": sliced_w1(trace[-1], target),
            "initial_sliced_w1": sliced_w1(initial, target),
            "euler_refinement_rmse": float(np.sqrt(np.mean((trace[-1]-refined[-1])**2))),
            "training_updates": 700}


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--output", type=Path, default=Path("outputs"))
    args = parser.parse_args()
    print(json.dumps(run(args.output), indent=2, allow_nan=False))