"""Reproduce three Gradiance teaching examples. All data are SIMULATED.

Run: python reproduce.py --output ./generated
Dependencies used: numpy==1.26.4 scipy==1.17.1 matplotlib==3.10.9
No Gradiance code, network, experimental data, or saved app results are used.
"""
import argparse
import json
from pathlib import Path

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
from matplotlib.patches import Rectangle
import numpy as np
import scipy
from scipy.optimize import least_squares

BLUE, ORANGE, GRAY = "#315ba8", "#ba501e", "#8795a5"
plt.rcParams.update({"font.size": 11, "axes.spines.top": False,
                     "axes.spines.right": False, "figure.facecolor": "white"})


def calibration(x, p):
    bottom, span, center, hill, asymmetry = p
    return bottom + span / (1 + (center / x) ** hill) ** asymmetry


def inverse_calibration(y, p):
    bottom, span, center, hill, asymmetry = p
    fraction = (np.asarray(y) - bottom) / span
    if np.any((fraction <= 0) | (fraction >= 1)):
        raise ValueError("Signal outside fitted asymptotes; do not extrapolate")
    return center / (fraction ** (-1 / asymmetry) - 1) ** (1 / hill)


def fit_calibration(x, y, free_asymmetry):
    # Positive parameters via log coordinates. Span = Top - Bottom.
    lower = np.log([0.0001, 0.01, 0.0001, 0.1, 0.1])
    upper = np.log([1, 5, 10000, 5, 10])
    n = 5 if free_asymmetry else 4
    fits = []
    for center in [1, 4, 10]:
        for asymmetry in ([0.4, 1, 2] if free_asymmetry else [1]):
            start = np.log([0.08, 2.32, center, 1, asymmetry])[:n]
            def residual(q):
                p = np.exp(q) if free_asymmetry else np.r_[np.exp(q), 1]
                return calibration(x, p) - y
            result = least_squares(residual, start, bounds=(lower[:n], upper[:n]),
                                   max_nfev=10000, ftol=1e-12, xtol=1e-12, gtol=1e-12)
            if result.success:
                fits.append(result)
    if not fits:
        raise RuntimeError("No converged calibration fit")
    best = min(fits, key=lambda r: np.sum(r.fun ** 2))
    return np.exp(best.x) if free_asymmetry else np.r_[np.exp(best.x), 1]


def save_figure(fig, out, name):
    fig.savefig(out / name, dpi=160, bbox_inches="tight")
    plt.close(fig)


def standard_curve(out):
    rng = np.random.default_rng(608)
    truth = np.array([0.08, 2.32, 4, 1.1, 0.45])
    x = np.repeat(np.logspace(-2, 3, 12), 3)
    y = calibration(x, truth) + rng.normal(0, 0.015, x.size)
    np.savetxt(out / "standard-curve-simulated.csv", np.c_[x, y], delimiter=",",
               header="concentration_arbitrary_units,signal_arbitrary_units", comments="")
    fig, axes = plt.subplots(2, 1, figsize=(10, 7), sharex=True,
                             gridspec_kw={"height_ratios": [2, 1]}, layout="constrained")
    grid = np.logspace(-2, 3, 400)
    axes[0].scatter(x, y, s=23, color=GRAY, label="36 simulated observations")
    checks = np.array([0.1, 1, 10])
    results = {}
    for name, free, color in [("4PL", False, BLUE), ("5PL", True, ORANGE)]:
        p = fit_calibration(x, y, free)
        residual = y - calibration(x, p)
        recovered = inverse_calibration(calibration(checks, truth), p)
        results[name] = {"parameters_B_span_C_H_S": p.tolist(),
                         "rmse": float(np.sqrt(np.mean(residual ** 2))),
                         "check_concentrations": checks.tolist(),
                         "back_calculated": recovered.tolist(),
                         "check_error_percent": ((recovered / checks - 1) * 100).tolist()}
        axes[0].plot(grid, calibration(grid, p), color=color, label=name)
        axes[1].scatter(x, residual, s=22, alpha=0.8, color=color, label=name)
        assert np.allclose(inverse_calibration(calibration(checks, p), p), checks)
    axes[0].set(title="4PL vs. 5PL on an asymmetric teaching dataset", ylabel="Signal (arbitrary units)")
    axes[0].legend(frameon=False)
    axes[1].axhline(0, color=GRAY, linewidth=1)
    axes[1].set(xscale="log", xlabel="Concentration (arbitrary units)", ylabel="Observed − fitted")
    fig.supxlabel("SIMULATED DATA · Independent SciPy fits · Not an experimental or app result", fontsize=10)
    save_figure(fig, out, "4pl-vs-5pl.png")
    return {"seed": 608, "truth_B_span_C_H_S": truth.tolist(), "fits": results}


def inhibition(x, bottom, top, center, hill):
    return bottom + (top - bottom) / (1 + (x / center) ** hill)


def dose_response(out):
    rng = np.random.default_rng(609)
    x = np.repeat(np.logspace(-2, 3, 11), 3)
    y = inhibition(x, 20, 100, 3, 1.2) + rng.normal(0, 1.2, x.size)
    np.savetxt(out / "ic50-simulated.csv", np.c_[x, y], delimiter=",",
               header="concentration_uM,activity_percent", comments="")
    def residual(q):
        return inhibition(x, q[0], q[1], np.exp(q[2]), np.exp(q[3])) - y
    result = least_squares(residual, [15, 100, np.log(3), np.log(1)],
                           bounds=([-20, 70, np.log(0.001), np.log(0.1)],
                                   [60, 130, np.log(10000), np.log(5)]),
                           ftol=1e-12, xtol=1e-12, gtol=1e-12)
    assert result.success
    bottom, top, center, hill = result.x[0], result.x[1], np.exp(result.x[2]), np.exp(result.x[3])
    absolute = center * ((top - bottom) / (50 - bottom) - 1) ** (1 / hill)
    assert abs(inhibition(absolute, bottom, top, center, hill) - 50) < 1e-8
    grid = np.logspace(-2, 3, 400)
    fig, ax = plt.subplots(figsize=(10, 6), layout="constrained")
    low = x <= 1
    ax.axvspan(0.01, 1, color="#f5eadf", label="Restricted range: ≤1 µM")
    ax.scatter(x[~low], y[~low], color=GRAY, s=32, label="Additional simulated measurements")
    ax.scatter(x[low], y[low], color=ORANGE, s=32, label="Restricted-range measurements")
    ax.plot(grid, inhibition(grid, bottom, top, center, hill), color=BLUE, label="4PL fit to all 33 observations")
    ax.axhline(50, color=GRAY, linestyle=":")
    ax.scatter([center, absolute], [(top + bottom) / 2, 50], color=BLUE, marker="D", zorder=5)
    ax.annotate(f"Relative midpoint\n{center:.2f} µM", (center, (top + bottom) / 2),
                xytext=(12, 82), arrowprops={"arrowstyle": "->", "color": BLUE})
    ax.annotate(f"50% of control\n{absolute:.2f} µM", (absolute, 50),
                xytext=(25, 56), arrowprops={"arrowstyle": "->", "color": BLUE})
    ax.set(xscale="log", xlabel="Concentration (µM)", ylabel="Activity (% of control)",
           title="A fitted midpoint and 50% of control answer different questions", ylim=(0, 110))
    ax.legend(loc="lower left", fontsize=9, frameon=False)
    fig.supxlabel("SIMULATED DATA · The restricted subset was not fitted separately", fontsize=10)
    save_figure(fig, out, "ic50-troubleshooting.png")
    return {"seed": 609, "observations": len(x), "restricted_observations": int(low.sum()),
            "restricted_min_activity": float(y[low].min()), "bottom": bottom, "top": top,
            "relative_midpoint_uM": center, "absolute_50_percent_uM": absolute, "hill": hill}


def flow_gating(out):
    rng = np.random.default_rng(610)
    n = 5000
    debris = rng.random(n) < 0.12
    fsc = rng.normal(5, 0.9, n); ssc = rng.normal(3, 0.6, n)
    fsc[debris] = rng.normal(1, 0.35, debris.sum())
    ssc[debris] = rng.normal(0.8, 0.3, debris.sum())
    ratio = rng.normal(1, 0.045, n)
    doublet = rng.random(n) < 0.1
    ratio[doublet] = rng.normal(0.65, 0.04, doublet.sum())
    height = fsc * ratio
    dead = rng.random(n) < 0.15
    viability = rng.normal(0.25, 0.12, n) + 1.5 * dead
    cd3 = rng.normal(0.25, 0.16, n) + 1.45 * (rng.random(n) < 0.7)
    cd4 = rng.normal(0.25, 0.2, n) + 1.4 * (rng.random(n) < 0.55)
    # A separate simulated FMO distribution, already in display coordinates.
    fmo = rng.normal(0.25, 0.2, 2000)
    threshold = float(np.quantile(fmo, 0.99))
    all_events = np.ones(n, dtype=bool)
    cells = (fsc >= 2.5) & (fsc <= 7.5) & (ssc >= 1.3) & (ssc <= 4.7)
    singlets = cells & (ratio >= 0.85) & (ratio <= 1.15)
    live = singlets & (viability < 0.8)
    t_cells = live & (cd3 > 0.8)
    cd4_cells = t_cells & (cd4 > threshold)
    masks = [all_events, cells, singlets, live, t_cells, cd4_cells]
    labels = ["All events", "Cells", "Singlets", "Live", "CD3+", "CD4+ within CD3+"]
    rows = []
    for i, (label, mask) in enumerate(zip(labels, masks)):
        parent = masks[max(0, i - 1)]
        assert np.all(~mask | parent)
        rows.append({"population": label, "count": int(mask.sum()),
                     "percent_parent": float(mask.sum() / parent.sum() * 100),
                     "percent_all": float(mask.mean() * 100)})
    np.savetxt(out / "flow-events-simulated.csv", np.c_[fsc, ssc, height, viability, cd3, cd4],
               delimiter=",", fmt="%.8f", comments="", header="FSC_A,SSC_A,FSC_H,Viability,CD3,CD4")
    np.savetxt(out / "flow-fmo-simulated.csv", fmo, delimiter=",", comments="", header="CD4_display_coordinate")
    fig, axes = plt.subplots(2, 2, figsize=(11, 9), layout="constrained")
    ax = axes[0, 0]
    ax.scatter(fsc, ssc, s=3, color=GRAY, alpha=0.3, rasterized=True)
    ax.add_patch(Rectangle((2.5, 1.3), 5, 3.4, fill=False, edgecolor=BLUE, linewidth=2))
    ax.set(title=f"1. Cells / all events: {cells.sum():,} / {n:,}", xlabel="FSC-A", ylabel="SSC-A")
    ax = axes[0, 1]
    ax.scatter(fsc[cells], height[cells], s=3, color=GRAY, alpha=0.3, rasterized=True)
    for slope in [0.85, 1.15]:
        ax.plot([2.5, 7.5], np.array([2.5, 7.5]) * slope, color=BLUE)
    ax.set(title=f"2. Singlets / cells: {singlets.sum():,} / {cells.sum():,}", xlabel="FSC-A", ylabel="FSC-H")
    ax = axes[1, 0]
    ax.hist(viability[singlets], bins=60, color=GRAY)
    ax.axvline(0.8, color=BLUE); ax.text(0, ax.get_ylim()[1] * 0.8, "Keep left", color=BLUE)
    ax.set(title=f"3. Live / singlets: {live.sum():,} / {singlets.sum():,}", xlabel="Viability dye", ylabel="Events")
    ax = axes[1, 1]
    ax.scatter(cd3[live], cd4[live], s=3, color=GRAY, alpha=0.3, rasterized=True)
    ax.axvline(0.8, color=BLUE); ax.axhline(threshold, color=ORANGE)
    ax.text(1.05, 2.3, f"CD3+ CD4+: {cd4_cells.sum():,}", color=BLUE)
    ax.set(title="4. Marker gates within live events", xlabel="CD3", ylabel="CD4")
    fig.suptitle("Flow gating: follow the parent population", fontsize=17)
    fig.supxlabel("SIMULATED EVENTS · Arbitrary display coordinates · Boundaries are teaching choices", fontsize=10)
    save_figure(fig, out, "flow-gating.png")
    return {"seed": 610, "cd4_fmo_99th_percentile": threshold, "populations": rows}


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--output", type=Path, default=Path("generated"))
    out = parser.parse_args().output
    out.mkdir(parents=True, exist_ok=True)
    result = {"evidence": "Simulated teaching data; not experimental or live app results",
              "versions": {"numpy": np.__version__, "scipy": scipy.__version__, "matplotlib": matplotlib.__version__},
              "standard_curve": standard_curve(out), "ic50": dose_response(out), "flow": flow_gating(out)}
    (out / "results.json").write_text(json.dumps(result, indent=2) + "\n")
    print(json.dumps(result, indent=2))
