"""Original scientific figures for the diffusion / flow-matching article.

Run with NumPy and Matplotlib. The cover is an analytic Gaussian transport,
not the output of a trained model. The finite-time example uses exact rotation.
"""
from pathlib import Path

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

OUT = Path(__file__).parent
BLUE = "#2f6db1"
GREEN = "#2e7d4f"
ORANGE = "#c96a2a"
BG = "#fbfbf8"
INK = "#272c32"
plt.rcParams.update({
    "font.family": "sans-serif",
    "font.size": 12,
    "text.color": INK,
    "axes.labelcolor": INK,
    "svg.fonttype": "none",
    "svg.hashsalt": "diffusion-flow-concepts",
})


def save_svg(fig, name, title):
    output = OUT / name
    fig.savefig(output, facecolor=BG,
                metadata={"Date": None, "Title": title})
    output.write_text("\n".join(line.rstrip() for line in output.read_text().splitlines()) + "\n")
    plt.close(fig)


def cover():
    # X_t=(1-t)*X_0+t*epsilon, X_0 ~ N(mu, diag(var)), epsilon ~ N(0,I).
    # Its exact marginal ODE carries z_1 to (1-t)*mu + sqrt(C_t)*z_1.
    rng = np.random.default_rng(10)
    z = rng.normal(size=(700, 2))
    mu = np.array([5.0, 0.0])
    var = np.array([0.20, 0.65])
    t = np.linspace(1, 0, 220)
    std = np.sqrt((1-t[:, None])**2 * var + t[:, None]**2)
    mean = (1-t[:, None]) * mu
    fig, ax = plt.subplots(figsize=(12, 6.2))
    fig.subplots_adjust(left=.025, right=.975, top=.85, bottom=.10)
    ax.set_facecolor(BG)
    ax.scatter(z[:, 0], z[:, 1], s=4, c=BLUE, alpha=.27, linewidths=0)
    end = mu + np.sqrt(var) * z
    ax.scatter(end[:, 0], end[:, 1], s=4, c=ORANGE, alpha=.33, linewidths=0)
    for initial in [(-.7, -1.6), (.5, -1.1), (-.2, -.5), (.0, .0), (.7, .6), (-.5, 1.2), (.1, 1.8)]:
        path = mean + std * initial
        ax.plot(path[:, 0], path[:, 1], color=GREEN, lw=1.2, alpha=.65)
        k = 118
        ax.annotate("", xy=path[k+5], xytext=path[k], arrowprops={"arrowstyle":"->", "color":GREEN, "lw":1.3})
    ax.text(0, 2.8, "Gaussian source", ha="center", fontsize=16, color=BLUE)
    ax.text(5, 2.8, "Gaussian target", ha="center", fontsize=16, color=ORANGE)
    ax.text(2.5, -3.1, "Exact ODE trajectories for an independent Gaussian interpolant", ha="center", fontsize=12, color="#58616a")
    ax.set(xlim=(-2.9, 7.0), ylim=(-3.4, 3.35), aspect="equal")
    ax.axis("off")
    fig.suptitle("A continuous flow transports probability", fontsize=23, y=.94)
    fig.savefig(OUT / "hero_gaussian_transport.png", dpi=160, facecolor=BG)
    plt.close(fig)


def framework():
    fig, ax = plt.subplots(figsize=(8, 7.2))
    ax.set(xlim=(0, 1), ylim=(0, 1))
    ax.axis("off")
    def box(x, y, w, h, title, body, color):
        ax.add_patch(FancyBboxPatch((x, y), w, h, boxstyle="round,pad=0.014", facecolor="white", edgecolor=color, linewidth=1.2))
        ax.text(x+w/2, y+h-.035, title, ha="center", va="top", fontsize=12, weight="bold", color=color)
        ax.text(x+w/2, y+h/2-.035, body, ha="center", va="center", fontsize=11, linespacing=1.6)
    box(.035, .72, .39, .17, "SCORE-BASED DIFFUSION", "Specify forward SDE\nLearn its marginal score", BLUE)
    box(.575, .72, .39, .17, "FLOW MATCHING", "Specify an interpolant\nRegress conditional velocities", ORANGE)
    box(.035, .40, .39, .19, "ASSOCIATED PF-ODE", r"$b_{\rm PF}=f-\frac{1}{2}g^2\nabla\log p_t$"+"\nShares the SDE marginals*", BLUE)
    box(.575, .40, .39, .19, "MARGINAL ODE", r"$b_{\rm FM}(x,t)=\mathbb{E}[\dot{X}_t\mid X_t=x]$"+"\nGenerates interpolant marginals*", ORANGE)
    for x in [.23, .77]:
        ax.annotate("", xy=(x,.61),xytext=(x,.70),arrowprops={"arrowstyle":"->", "color":"#7e858b", "lw":1.4})
    ax.annotate("",xy=(.54,.49),xytext=(.46,.49),arrowprops={"arrowstyle":"<->", "color":GREEN, "lw":1.4})
    ax.text(.5, .325, "Compatible construction can give equal fields.", ha="center", fontsize=12, color=GREEN)
    ax.text(.5, .26, "Equal marginals alone do not imply equal fields.", ha="center", fontsize=12, weight="bold")
    ax.text(.5, .17, r"Both satisfy $\partial_t p_t+\nabla\!\cdot(p_t b_t)=0$.", ha="center", fontsize=13)
    ax.text(.5, .075, "*Exact fields and evolution; suitable regularity and initialization.", ha="center", fontsize=10, color="#60676e")
    fig.suptitle("Where diffusion and flow matching meet", fontsize=16, y=.98)
    fig.subplots_adjust(left=.02,right=.98,bottom=.015,top=.95)
    save_svg(fig,"fig_framework_connections.svg","Diffusion and flow matching: conditional correspondence")


def finite_time():
    # z(t)=(cos(w*t),sin(w*t)), b(z)=w*J*z. Generate backward from t=1 to 0.
    omega = np.pi/2
    t = np.linspace(1, 0, 250)
    z = np.stack([np.cos(omega*t),np.sin(omega*t)],axis=1)
    start, end = z[0],z[-1]
    b_start=omega*np.array([-start[1],start[0]])
    euler=start-b_start
    fig, ax = plt.subplots(figsize=(8, 5.8))
    ax.set_facecolor(BG)
    ax.plot(z[:,0],z[:,1],color=BLUE,lw=2.5,label="Exact ODE trajectory")
    ax.annotate("",xy=z[115],xytext=z[105],arrowprops={"arrowstyle":"->","color":BLUE,"lw":2})
    ax.annotate("",xy=euler,xytext=start,arrowprops={"arrowstyle":"->","color":ORANGE,"linestyle":"--","lw":2})
    ax.annotate("",xy=end,xytext=start,arrowprops={"arrowstyle":"->","color":GREEN,"lw":2})
    ax.scatter([start[0],end[0]],[start[1],end[1]],c=[BLUE,GREEN],s=45,zorder=5)
    ax.scatter([euler[0]],[euler[1]],c=ORANGE,s=45,zorder=5)
    ax.text(-.04,1.07,r"$z_t$",color=BLUE,fontsize=16,ha="right")
    ax.text(1.03,-.015,r"$z_r$",color=GREEN,fontsize=16)
    ax.text(.76,1.09,"Local Euler update",ha="center",color=ORANGE,fontsize=12)
    ax.text(.31,.38,"Finite-interval\ndisplacement",ha="center",color=GREEN,fontsize=12)
    ax.text(1.21,.57,"Exact\ntrajectory",ha="center",color=BLUE,fontsize=12)
    ax.set(xlim=(-.25,1.9),ylim=(-.18,1.33),aspect="equal")
    ax.axis("off")
    fig.suptitle("A local velocity and an endpoint map are different objects",fontsize=15,y=.96)
    fig.text(.5,.08,r"$\Phi_{r\leftarrow t}(z_t)=z_t-(t-r)u(z_t,r,t)$",ha="center",fontsize=16,color=GREEN)
    fig.text(.5,.025,"Exact rotation, t = 1 to r = 0. A learned map only approximates the endpoint.",ha="center",fontsize=10,color="#60676e")
    fig.subplots_adjust(left=.04,right=.97,top=.89,bottom=.18)
    save_svg(fig,"fig_velocity_and_map.svg","Instantaneous velocity versus a finite-time flow map")


if __name__ == "__main__":
    cover()
    framework()
    finite_time()
    print("Wrote three scientific figures to",OUT)