"""Draw the PIVOT overview with editable vector paths and embedded fonts."""

from pathlib import Path
import numpy as np
import matplotlib

matplotlib.use("Agg")
import matplotlib.pyplot as plt
from matplotlib import font_manager as fm
from matplotlib.patches import Ellipse, FancyArrowPatch, PathPatch, FancyBboxPatch
from matplotlib.path import Path as MPath

ROOT = Path(__file__).resolve().parents[1]
for f in (ROOT / "assets/fonts").glob("*.ttf"):
    fm.fontManager.addfont(str(f))
plt.rcParams.update(
    {
        "font.family": "Ubuntu",
        "mathtext.fontset": "cm",
        "pdf.fonttype": 42,
        "ps.fonttype": 42,
        "font.size": 13,
    }
)
blue = "#62AEDD"
navy = "#245E84"
red = "#9D2944"
ink = "#243746"
pale = "#E9F4FB"
gray = "#9EB0BD"
fig, ax = plt.subplots(figsize=(13.6, 4.8))
ax.set(xlim=(0, 13.6), ylim=(0, 4.8))
ax.axis("off")


def text(x, y, s, size=13, color=ink, **kw):
    ax.text(
        x,
        y,
        s,
        fontsize=size,
        color=color,
        ha=kw.pop("ha", "center"),
        va="center",
        **kw,
    )


def arrow(start, end, color=blue, rad=0, lw=2.2):
    ax.add_patch(
        FancyArrowPatch(
            start,
            end,
            arrowstyle="-|>",
            mutation_scale=16,
            connectionstyle=f"arc3,rad={rad}",
            color=color,
            lw=lw,
        )
    )


def cloud(x, y, w, h, color, seed):
    rng = np.random.default_rng(seed)
    ax.add_patch(Ellipse((x, y), w, h, facecolor=color, alpha=0.10, edgecolor="none"))
    z = rng.normal(size=(36, 2))
    z = z[np.linalg.norm(z, axis=1) < 2]
    ax.scatter(
        x + z[:, 0] * w / 5,
        y + z[:, 1] * h / 5,
        s=rng.uniform(9, 24, len(z)),
        color=color,
        alpha=0.7,
        linewidth=0,
    )


cloud(1.45, 3.08, 2.1, 1.72, blue, 2)
text(1.45, 4.19, "Control population", 18, navy, weight="bold")
text(1.45, 2.02, r"$c_0\sim\rho_0$", 18)
# One broad map across the figure, with faint parallel sample trajectories.
for off in [-0.32, -0.12, 0.12, 0.32]:
    verts = [(2.55, 3.1 + off), (4.3, 3.8 + off), (6.7, 3.9 + off), (8.75, 3.15 + off)]
    ax.add_patch(
        PathPatch(
            MPath(verts, [MPath.MOVETO, MPath.CURVE4, MPath.CURVE4, MPath.CURVE4]),
            fill=False,
            edgecolor=blue,
            alpha=0.20,
            lw=10,
        )
    )
arrow((2.6, 3.1), (8.8, 3.14), blue, 0.20, 3)
text(5.6, 3.72, r"$X_\theta(0,1,c_0,e_u)$", 23)
text(5.6, 3.17, "Predict the response", 18, navy, weight="bold")
text(5.6, 4.38, r"$u=\{(g_j,o_j)\}_{j=1}^{M}\quad\longmapsto\quad e_u$", 17)
arrow((5.6, 4.12), (5.6, 3.92), gray, 0, 1.2)
cloud(9.4, 3.18, 1.75, 1.5, blue, 9)
text(9.4, 4.19, "Predicted cells", 18, navy, weight="bold")
cloud(11.97, 3.7, 1.65, 1.18, red, 3)
text(11.97, 4.56, "Target cells", 18, red, weight="bold")
text(12, 2.88, r"$c^\star\sim\rho^\star$", 18)
arrow((10.2, 3.23), (11.22, 3.63), red, -0.08, 1.6)
text(11, 2.57, "Endpoint reward", 17, red)
# Reward gradients return to the intervention coordinates.
arrow((11.3, 2.31), (5.22, 1.89), red, -0.17, 2.6)
text(7.95, 1.36, "Optimize intervention embeddings", 17, red)
text(3.05, 1.72, "Rank admissible interventions", 17, navy, weight="bold")
arrow((5.01, 2.03), (3.13, 2.25), navy, -0.10, 1.5)
text(3.45, 2.48, r"$e^{(L)}\;\longrightarrow\;u_{1:K}$", 18)
ax.add_patch(
    FancyBboxPatch(
        (0.18, 0.12),
        13.22,
        1.03,
        boxstyle="round,pad=0.02,rounding_size=.08",
        facecolor=pale,
        edgecolor="none",
    )
)
text(0.51, 0.78, "PIVOT", 18, navy, ha="left", weight="bold")
text(
    7.8,
    0.74,
    r"$\hat c_1=X_\theta(0,1,c_0,e)\qquad g_e=J_eX_\theta^{\mathsf{T}}\nabla_{\hat c_1}r\qquad e^+=e+\gamma\,\dfrac{g_e}{\|g_e\|_2+\epsilon}$",
    20,
)
text(
    7.6,
    0.29,
    "Endpoint prediction, reward gradients, and admissible interventions.",
    16,
)
fig.subplots_adjust(left=0, right=1, bottom=0, top=1)
for ext in ["pdf", "png", "svg"]:
    fig.savefig(ROOT / "assets" / f"figure1.{ext}", dpi=250, facecolor="white")