"""
Illustrative figures for:
1. Shape-valued treatment effects
2. Residual topology as a diagnostic
3. Distribution-level treatment effects

Outputs:
    fig_shape_tate.pdf / .png
    fig_residual_topology.pdf / .png
    fig_distribution_effect.pdf / .png

Dependencies:
    numpy
    matplotlib
    scipy

Optional:
    ripser  (for data-driven persistence diagrams)
Install with:
    pip install numpy matplotlib scipy ripser
"""

from pathlib import Path
import numpy as np
import matplotlib.pyplot as plt
from matplotlib import rcParams
from scipy.spatial import cKDTree

# Optional persistent-homology package
try:
    from ripser import ripser
    HAVE_RIPSER = True
except ImportError:
    HAVE_RIPSER = False


# ---------------------------------------------------------------------
# Global style
# ---------------------------------------------------------------------
OUT = Path(__file__).resolve().parent
OUT.mkdir(exist_ok=True)

rcParams.update({
    "font.family": "serif",
    "font.size": 10,
    "axes.labelsize": 10,
    "axes.titlesize": 11,
    "legend.fontsize": 8.5,
    "xtick.labelsize": 8.5,
    "ytick.labelsize": 8.5,
    "figure.dpi": 160,
    "savefig.dpi": 300,
    "savefig.bbox": "tight",
})

rng = np.random.default_rng(2026)


def save_both(fig, stem):
    fig.savefig(OUT / f"{stem}.pdf")
    fig.savefig(OUT / f"{stem}.png")
    plt.close(fig)


def clean_3d_axis(ax):
    ax.set_xticks([])
    ax.set_yticks([])
    ax.set_zticks([])
    ax.set_box_aspect((1, 1, 1))
    ax.grid(False)


def persistence_diagram(ax, points, maxdim=1, title=None):
    """
    Plot an empirical persistence diagram when ripser is available.
    Otherwise draw a clearly labelled schematic diagram.
    """
    if HAVE_RIPSER:
        diagrams = ripser(points, maxdim=maxdim)["dgms"]
        colors = ["C0", "C1", "C2"]
        all_finite = []

        for dim, dgm in enumerate(diagrams):
            finite = dgm[np.isfinite(dgm[:, 1])]
            if len(finite):
                all_finite.append(finite)
                ax.scatter(
                    finite[:, 0], finite[:, 1],
                    s=24, alpha=0.85, label=rf"$H_{dim}$",
                    color=colors[dim]
                )

        if all_finite:
            values = np.vstack(all_finite)
            upper = 1.08 * np.max(values[:, 1])
        else:
            upper = 1.0

        ax.plot([0, upper], [0, upper], "--", lw=1, color="0.45")
        ax.set_xlim(0, upper)
        ax.set_ylim(0, upper)
    else:
        # Schematic fallback: one long-lived H1 class plus short-lived noise.
        births = np.array([0.08, 0.12, 0.16, 0.20, 0.25])
        deaths = np.array([0.18, 0.24, 0.31, 0.95, 0.38])
        ax.scatter(births, deaths, s=28, alpha=0.85, label=r"$H_1$")
        upper = 1.0
        ax.plot([0, upper], [0, upper], "--", lw=1, color="0.45")
        ax.annotate(
            "persistent loop",
            xy=(0.20, 0.95), xytext=(0.43, 0.77),
            arrowprops=dict(arrowstyle="->", lw=0.9)
        )
        ax.set_xlim(0, upper)
        ax.set_ylim(0, upper)

    # Keep legend symbols outside the persistence-diagram data region.
    ax.legend(
        frameon=False,
        loc="upper left",
        bbox_to_anchor=(1.02, 1.0),
        borderaxespad=0.0,
    )

    ax.set_xlabel("Birth")
    ax.set_ylabel("Death")
    if title:
        ax.set_title(title)


# =====================================================================
# Figure 1: Shape-valued treatment effects
# =====================================================================
def make_tumour_voxels(n=34):
    x, y, z = np.mgrid[-1:1:complex(n), -1:1:complex(n), -1:1:complex(n)]

    # Potential outcome under a=0: irregular but connected solid tumour
    base0 = (
        (x / 0.78) ** 2
        + (y / 0.66) ** 2
        + (z / 0.72) ** 2
        < 1
    )
    lobe0 = (
        ((x - 0.46) / 0.34) ** 2
        + ((y + 0.08) / 0.29) ** 2
        + ((z - 0.05) / 0.32) ** 2
        < 1
    )
    tumour0 = base0 | lobe0

    # Potential outcome under a=1:
    # smaller principal mass, an enclosed cavity, a tunnel, and a satellite.
    base1 = (
        ((x + 0.05) / 0.70) ** 2
        + ((y - 0.02) / 0.58) ** 2
        + (z / 0.64) ** 2
        < 1
    )

    cavity = (
        ((x + 0.12) / 0.23) ** 2
        + ((y - 0.02) / 0.20) ** 2
        + ((z + 0.03) / 0.22) ** 2
        < 1
    )

    # Cylindrical tunnel approximately parallel to the x-axis
    tunnel = ((y + 0.18) ** 2 + (z - 0.05) ** 2 < 0.12 ** 2) & (np.abs(x) < 0.72)

    satellite = (
        ((x - 0.78) / 0.16) ** 2
        + ((y + 0.50) / 0.14) ** 2
        + ((z - 0.18) / 0.16) ** 2
        < 1
    )

    tumour1 = (base1 & ~cavity & ~tunnel) | satellite
    return tumour0, tumour1


def figure_shape_tate():
    tumour0, tumour1 = make_tumour_voxels()

    fig = plt.figure(figsize=(12.0, 3.7))
    gs = fig.add_gridspec(1, 3, width_ratios=[1, 1, 1.25], wspace=0.18)

    ax0 = fig.add_subplot(gs[0, 0], projection="3d")
    ax1 = fig.add_subplot(gs[0, 1], projection="3d")
    ax2 = fig.add_subplot(gs[0, 2])

    ax0.voxels(tumour0, facecolors="C0", edgecolor="none", alpha=0.82)
    ax1.voxels(tumour1, facecolors="C1", edgecolor="none", alpha=0.82)

    ax0.set_title(r"Potential tumour shape $Y^0$")
    ax1.set_title(r"Potential tumour shape $Y^1$")
    clean_3d_axis(ax0)
    clean_3d_axis(ax1)

    # Stylized Banach-valued effect curve for one fixed homological degree.
    t = np.linspace(0, 1, 350)
    delta = (
        0.70 * np.exp(-((t - 0.22) / 0.10) ** 2)
        - 0.42 * np.exp(-((t - 0.50) / 0.13) ** 2)
        + 0.52 * np.exp(-((t - 0.77) / 0.11) ** 2)
    )
    ax2.axhline(0, lw=0.9, color="0.5")
    ax2.plot(t, delta, lw=2.2)
    ax2.fill_between(t, 0, delta, alpha=0.18)
    ax2.set_xlabel(r"Filtration scale $t$")
    ax2.set_ylabel(r"$\mathbb{E}[Z_k^1](t)-\mathbb{E}[Z_k^0](t)$")
    ax2.set_title(r"Functional effect at fixed degree $k$")

    ax2.annotate(
        "positive effect",
        xy=(0.22, delta[np.argmin(np.abs(t - 0.22))]),
        xytext=(0.05, 0.88),
        textcoords="axes fraction",
        arrowprops=dict(arrowstyle="->", lw=0.9),
        fontsize=8.5
    )
    ax2.annotate(
        "negative effect",
        xy=(0.50, delta[np.argmin(np.abs(t - 0.50))]),
        xytext=(0.37, 0.10),
        textcoords="axes fraction",
        arrowprops=dict(arrowstyle="->", lw=0.9),
        fontsize=8.5
    )
    ax2.annotate(
        "positive effect",
        xy=(0.77, delta[np.argmin(np.abs(t - 0.77))]),
        xytext=(0.67, 0.83),
        textcoords="axes fraction",
        arrowprops=dict(arrowstyle="->", lw=0.9),
        fontsize=8.5
    )
    ax2.spines[["top", "right"]].set_visible(False)

    fig.suptitle(
        "Shape-valued treatment effects: topology is computed for each potential outcome",
        y=1.02, fontsize=12
    )
    save_both(fig, "fig_shape_tate")


# =====================================================================
# Figure 2: Residual topology as a diagnostic
# =====================================================================
def figure_residual_topology():
    n = 330
    u = rng.uniform(0, 2 * np.pi, n)
    noise = 0.075
    v = np.cos(u) + rng.normal(0, noise, n)
    w = np.sin(u) + rng.normal(0, noise, n)
    cloud = np.column_stack([v, w])
    cov = np.cov(v, w, ddof=1)[0, 1]
    corr = np.corrcoef(v, w)[0, 1]

    fig, axes = plt.subplots(1, 3, figsize=(12.0, 3.45))

    axes[0].scatter(v, w, s=12, alpha=0.60)
    axes[0].axhline(0, lw=0.7, color="0.75")
    axes[0].axvline(0, lw=0.7, color="0.75")
    axes[0].set_aspect("equal")
    axes[0].set_xlabel(r"$V=\cos U+\varepsilon_V$")
    axes[0].set_ylabel(r"$W=\sin U+\varepsilon_W$")
    axes[0].set_title("Near-zero linear association")
    axes[0].text(
        0.04, 0.96,
        rf"sample covariance $={cov:.3f}$" + "\n" +
        rf"sample correlation $={corr:.3f}$",
        transform=axes[0].transAxes,
        va="top",
        bbox=dict(boxstyle="round,pad=0.25", fc="white", ec="0.8")
    )

    # A working linear model W ~ 1 + V and its residuals.
    X = np.column_stack([np.ones(n), v])
    beta = np.linalg.lstsq(X, w, rcond=None)[0]
    residual = w - X @ beta
    residual_cloud = np.column_stack([v, residual])

    axes[1].scatter(v, residual, s=12, alpha=0.60)
    axes[1].axhline(0, lw=0.8, color="0.5")
    axes[1].set_xlabel("Covariate $V$")
    axes[1].set_ylabel("Working-model residual")
    axes[1].set_title("Residual cloud retains structure")
    axes[1].text(
        0.04, 0.06,
        "A linear fit removes slope,\nnot the circular dependence.",
        transform=axes[1].transAxes,
        va="bottom",
        bbox=dict(boxstyle="round,pad=0.25", fc="white", ec="0.8")
    )

    persistence_diagram(
        axes[2], residual_cloud,
        maxdim=1,
        title=r"Persistent $H_1$ as a diagnostic"
    )

    for ax in axes:
        ax.spines[["top", "right"]].set_visible(False)

    fig.suptitle(
        "Residual topology can reveal nonlinear structure despite zero covariance",
        y=1.02, fontsize=12
    )
    save_both(fig, "fig_residual_topology")


# =====================================================================
# Figure 3: Distribution-level treatment effects
# =====================================================================
def approximate_dtm(points, gx, gy, mass_fraction=0.10):
    """
    Empirical distance-to-measure approximation:
    root mean square distance to the k nearest sample points.
    """
    query = np.column_stack([gx.ravel(), gy.ravel()])
    k = max(2, int(np.ceil(mass_fraction * len(points))))
    tree = cKDTree(points)
    dist, _ = tree.query(query, k=k)
    if k == 1:
        dtm = dist
    else:
        dtm = np.sqrt(np.mean(dist ** 2, axis=1))
    return dtm.reshape(gx.shape)


def figure_distribution_effect():
    n = 420

    # P^0: one persistent cluster centered at zero.
    p0 = rng.normal(loc=(0.0, 0.0), scale=(0.55, 0.55), size=(n, 2))

    # P^1: two separated clusters with equal weights and hence mean near zero.
    labels = rng.integers(0, 2, size=n)
    centers = np.where(labels[:, None] == 0, np.array([-1.35, 0.0]), np.array([1.35, 0.0]))
    p1 = centers + rng.normal(scale=(0.34, 0.43), size=(n, 2))

    # Recenter empirically so the plotted sample means are exactly equal.
    p0 = p0 - p0.mean(axis=0)
    p1 = p1 - p1.mean(axis=0)

    x = np.linspace(-2.7, 2.7, 165)
    y = np.linspace(-2.1, 2.1, 145)
    gx, gy = np.meshgrid(x, y)
    dtm0 = approximate_dtm(p0, gx, gy, mass_fraction=0.10)
    dtm1 = approximate_dtm(p1, gx, gy, mass_fraction=0.10)

    # Common contour levels make geometric differences comparable.
    pooled = np.concatenate([dtm0.ravel(), dtm1.ravel()])
    levels = np.quantile(pooled, [0.08, 0.16, 0.28, 0.43, 0.62])

    fig, axes = plt.subplots(1, 3, figsize=(12.2, 3.55))

    for ax, pts, dtm, title in [
        (axes[0], p0, dtm0, r"$P_Y^0$: one persistent cluster"),
        (axes[1], p1, dtm1, r"$P_Y^1$: two persistent clusters"),
    ]:
        ax.scatter(pts[:, 0], pts[:, 1], s=8, alpha=0.30)
        ax.contour(gx, gy, dtm, levels=levels, linewidths=1.2)
        ax.scatter(
            [0], [0],
            marker="x", s=80, linewidths=2, color="C1", zorder=5
        )
        ax.annotate(
            "common mean",
            xy=(0, 0),
            xytext=(0.68, 0.88),
            textcoords="axes fraction",
            color="C1",
            arrowprops=dict(arrowstyle="->", lw=1.0, color="C1"),
            bbox=dict(boxstyle="round,pad=0.2", fc="white", ec="none", alpha=0.85),
            fontsize=8.5,
        )
        ax.set_aspect("equal")
        ax.set_xlim(x.min(), x.max())
        ax.set_ylim(y.min(), y.max())
        ax.set_xlabel(r"$Y_1$")
        ax.set_ylabel(r"$Y_2$")
        ax.set_title(title)

    # Stylized topology-valued function of filtration scale.
    t = np.linspace(0, 1, 350)
    phi0 = 0.15 + 0.78 / (1 + np.exp(18 * (t - 0.43)))
    phi1 = (
        0.12
        + 0.42 / (1 + np.exp(20 * (t - 0.28)))
        + 0.55 / (1 + np.exp(20 * (t - 0.63)))
    )
    effect = phi1 - phi0

    axes[2].plot(t, phi0, lw=1.8, label=r"$T_{\mathrm{dist}}(P_Y^0)$")
    axes[2].plot(t, phi1, lw=1.8, label=r"$T_{\mathrm{dist}}(P_Y^1)$")
    axes[2].plot(t, effect, "--", lw=2.1, label=r"$\Delta_{\mathrm{dist}}$")
    axes[2].axhline(0, lw=0.8, color="0.55")
    axes[2].set_xlabel("Filtration scale")
    axes[2].set_ylabel("Topological summary")
    axes[2].set_title("Topology after forming each law")
    axes[2].legend(frameon=False)
    axes[2].annotate(
        r"$\mathbb{E}[Y^1]=\mathbb{E}[Y^0]$" + "\n" +
        "but the laws have different geometry",
        xy=(0.53, effect[np.argmin(np.abs(t - 0.53))]),
        xytext=(0.34, 0.12),
        textcoords="axes fraction",
        arrowprops=dict(arrowstyle="->", lw=0.9),
        fontsize=8.5
    )

    for ax in axes:
        ax.spines[["top", "right"]].set_visible(False)

    fig.suptitle(
        "Distribution-level treatment effects: equal means need not imply equal topology",
        y=1.02, fontsize=12
    )
    save_both(fig, "fig_distribution_effect")


if __name__ == "__main__":
    figure_shape_tate()
    figure_residual_topology()
    figure_distribution_effect()

    print(f"Figures written to: {OUT.resolve()}")
    if not HAVE_RIPSER:
        print(
            "Note: ripser was not installed, so the persistence diagram "
            "uses a labelled schematic fallback."
        )