"""Reproduce the analytic Gaussian counterexample in the diffusion/flow post.

Requires NumPy and Matplotlib. Run this file to write fig_rectified_gaussian.svg
next to it. Neither panel shows a trained neural network or an empirical claim.
"""
from pathlib import Path

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np


def variance(t):
    return (1 - t) ** 2 + t ** 2


def velocity(x, t):
    return (2 * t - 1) * x / variance(t)


def make_figure():
    plt.rcParams.update({
        "font.family": "sans-serif",
        "font.size": 12,
        "axes.spines.top": False,
        "axes.spines.right": False,
        "axes.edgecolor": "#b5b6b8",
        "text.color": "#222222",
        "axes.labelcolor": "#444444",
        "xtick.color": "#555555",
        "ytick.color": "#555555",
        "svg.fonttype": "none",
        "svg.hashsalt": "rectified-gaussian-example",
    })
    t = np.linspace(1, 0, 401)
    std = np.sqrt(variance(t))
    rng = np.random.default_rng(23)
    data, noise = rng.normal(size=(2, 16))
    fig, axes = plt.subplots(2, 1, figsize=(7.2, 7.2), layout="constrained")
    fig.set_facecolor("#fbfbf8")
    for ax in axes:
        ax.set_facecolor("#fbfbf8")
        ax.axhline(0, color="#d5d5d2", linewidth=0.8)
        ax.set_xlim(1, 0)
        ax.set_ylim(-2.65, 2.65)
        ax.set_xticks([1, 0.5, 0], ["1 · noise", "0.5", "0 · data"])
        ax.set_yticks([-2, 0, 2])
        ax.set_ylabel("State")
        ax.grid(axis="x", color="#e7e7e2", linewidth=0.7)
        ax.fill_between(t, -std, std, color="#dddcd3", alpha=0.35)
    for x0, eps in zip(data, noise):
        axes[0].plot(t, (1-t)*x0 + t*eps, color="#2f6db1", alpha=0.55, lw=1.3)
    axes[0].set_title("Conditional interpolants: each pair gives a straight line", loc="left", fontsize=12, pad=12)
    for z1 in [-2.2, -1.6, -1.0, -0.4, 0.4, 1.0, 1.6, 2.2]:
        axes[1].plot(t, std*z1, color="#2e7d4f", alpha=0.8, lw=1.6)
    axes[1].set_title("Marginal ODE: exact trajectories have changing speed", loc="left", fontsize=12, pad=12)
    axes[1].plot([1, 0], [2.2, 0], "--", color="#c96a2a", lw=2, label="One backward Euler step")
    axes[1].scatter([0], [0], color="#c96a2a", s=28, zorder=5, clip_on=False)
    axes[1].legend(loc="lower center", frameon=False, fontsize=11)
    axes[1].set_xlabel("Time t (generation runs from 1 to 0)", labelpad=12)
    fig.suptitle("Straight training lines do not guarantee accurate one-step sampling", fontsize=13, fontweight="normal")
    output = Path(__file__).with_name("fig_rectified_gaussian.svg")
    fig.savefig(output, metadata={
        "Date": None,
        "Title": "Rectified-flow Gaussian counterexample",
        "Description": "Independent Gaussian endpoint interpolants and their exact marginal ODE trajectories. A single Euler step collapses to zero.",
    })
    output.write_text("\n".join(line.rstrip() for line in output.read_text().splitlines()) + "\n")
    plt.close(fig)
    return output


if __name__ == "__main__":
    print(make_figure())