"""Figures for Gradiance's AI analysis articles.

python reproduce.py --puromycin Puromycin.csv --output generated
Requires numpy, scipy, matplotlib. No application imports or network requests.
The donor example is constructed teaching data; Puromycin is published data;
the agent diagram describes a conceptual workflow and contains no measurements.
"""
import argparse
import csv
import hashlib
import json
from pathlib import Path

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
from matplotlib.patches import FancyArrowPatch, FancyBboxPatch
import numpy as np
import scipy
from scipy.optimize import curve_fit

INK = "#1e3338"
TEAL = "#246b68"
RUST = "#b4623d"
MUTED = "#637477"
PAPER = "#faf9f5"
COLORS = ["#246b68", "#bc6848", "#6076a1", "#8c6e95", "#76814d", "#737a80"]
plt.rcParams.update({"font.family": "DejaVu Sans", "font.size": 12,
                     "text.color": INK, "axes.labelcolor": INK,
                     "axes.edgecolor": "#bcc6c4", "xtick.color": MUTED,
                     "ytick.color": MUTED, "axes.spines.top": False,
                     "axes.spines.right": False, "axes.facecolor": PAPER,
                     "figure.facecolor": PAPER, "svg.fonttype": "none"})


def save(fig, out, name):
    fig.savefig(out / f"{name}.png", dpi=160)
    fig.savefig(out / f"{name}.svg")
    plt.close(fig)


def donor_example(out):
    # Deliberately specified values, not observations or a fitted generative model.
    control = np.array([80, 100, 120, 140, 160, 180], dtype=float)
    changes = np.array([8, 4, 10, -2, 6, 12], dtype=float)
    offsets = np.array([-2, 0, 2], dtype=float)
    rows = []
    fig, axes = plt.subplots(1, 2, figsize=(11, 6.6), layout="constrained")
    fig.suptitle("36 wells. Six paired donors.", fontsize=23, fontweight="medium")
    for i, (baseline, delta, color) in enumerate(zip(control, changes, COLORS)):
        for condition, mean in [("control", baseline), ("treated", baseline + delta)]:
            for well, offset in enumerate(offsets, start=1):
                rows.append([f"D{i+1}", condition, well, mean + offset])
        for condition, mean, x in [("control", baseline, 0), ("treated", baseline + delta, 1)]:
            axes[0].scatter(x + np.array([-0.055, 0, 0.055]), mean + offsets,
                            color=color, s=40, alpha=0.85, zorder=3)
        axes[1].plot([0, 1], [baseline, baseline + delta], color=color,
                     linewidth=2, marker="o", markersize=7)
        axes[1].text(1.08, baseline + delta, f"D{i+1}  {delta:+.0f}",
                     va="center", color=color, fontsize=11)
    axes[0].set_title("Each dot is one technical well", loc="left", fontsize=13, pad=16)
    axes[1].set_title("Each line connects one donor’s means", loc="left", fontsize=13, pad=16)
    for ax in axes:
        ax.set_xticks([0, 1], ["Control", "Treated"])
        ax.set_ylim(65, 210)
        ax.set_yticks([80, 120, 160, 200])
        ax.grid(axis="y", color="#e0e5e1", linewidth=0.7)
        ax.set_axisbelow(True)
        ax.tick_params(length=0, pad=8)
    axes[0].set(xlim=(-0.5, 1.5), ylabel="Signal (arbitrary units)")
    axes[1].set_xlim(-0.25, 1.75)
    fig.supxlabel("CONSTRUCTED EXAMPLE  ·  Three wells per donor and condition  ·  Mean paired change: +6.33 units",
                  fontsize=10, color=MUTED)
    save(fig, out, "replicates")
    with (out / "paired-donors-example.csv").open("w", newline="") as f:
        writer = csv.writer(f)
        writer.writerow(["donor_id", "condition", "technical_well", "signal_au"])
        writer.writerows(rows)
    return {"evidence": "Constructed teaching data; no experimental measurements",
            "donors": 6, "wells_per_donor_condition": 3, "total_wells": len(rows),
            "paired_changes": changes.tolist(), "mean_paired_change": float(changes.mean())}


def residuals(source, out):
    digest = hashlib.sha256(source.read_bytes()).hexdigest()
    expected = "389be46927d2217b13d48b5b9ec56374007a87ea607c8e6f004c17683c427c70"
    if digest != expected:
        raise ValueError("Puromycin source checksum differs from the published reference")
    with source.open() as f:
        rows = list(csv.DictReader(f))
    def mm(x, vmax, km):
        return vmax * x / (km + x)
    fig, axes = plt.subplots(1, 2, figsize=(11, 6.6), layout="constrained", sharey=True)
    fig.suptitle("What the fitted curve leaves behind", fontsize=23, fontweight="medium")
    results, residual_rows = {}, []
    for ax, state, color in zip(axes, ["treated", "untreated"], [TEAL, RUST]):
        obs = [row for row in rows if row["state"] == state]
        x = np.array([float(row["conc"]) for row in obs])
        y = np.array([float(row["rate"]) for row in obs])
        p, _ = curve_fit(mm, x, y, p0=[200, 0.05], bounds=(0, np.inf), maxfev=10000)
        predicted = mm(x, *p)
        residual = y - predicted
        rmse = float(np.sqrt(np.mean(residual ** 2)))
        results[state] = {"n_observations": len(x), "vmax": float(p[0]), "km": float(p[1]),
                          "rmse": rmse, "r_squared": float(1 - np.sum(residual**2) / np.sum((y-y.mean())**2))}
        for dose, observed, fitted, error in zip(x, y, predicted, residual):
            residual_rows.append([state, dose, observed, fitted, error])
        ax.axhline(0, color=MUTED, linewidth=1)
        ax.vlines(x, 0, residual, color=color, alpha=0.3, linewidth=1.5)
        ax.scatter(x, residual, color=color, s=58, zorder=3, edgecolors=PAPER, linewidths=0.7)
        ax.set(xscale="log", xlabel="Concentration (ppm)",
               title=f"{state.capitalize()} · {len(x)} observations\nRMSE {rmse:.2f} counts/min/min")
        ax.set_xticks([0.02, 0.1, 1], ["0.02", "0.1", "1"])
        ax.tick_params(which="both", length=0, pad=7)
        ax.grid(axis="y", color="#e0e5e1", linewidth=0.7)
        ax.set_axisbelow(True)
    limit = 5 * np.ceil(max(abs(row[-1]) for row in residual_rows) / 5)
    for ax in axes:
        ax.set_ylim(-limit, limit)
        ax.set_yticks(np.arange(-limit, limit + 1, 10))
    axes[0].set_ylabel("Observed − fitted rate (counts/min/min)")
    fig.supxlabel("PUBLISHED PUROMYCIN DATA  ·  Separate unweighted Michaelis–Menten fits  ·  Independent SciPy calculation",
                  fontsize=9.7, color=MUTED)
    save(fig, out, "residuals")
    with (out / "puromycin-residuals.csv").open("w", newline="") as f:
        writer = csv.writer(f)
        writer.writerow(["state", "concentration_ppm", "observed_rate", "fitted_rate", "residual"])
        writer.writerows(residual_rows)
    return {"evidence": "Published Puromycin data, independently refitted; not a new app run",
            "source": "https://vincentarelbundock.github.io/Rdatasets/csv/datasets/Puromycin.csv",
            "sha256": digest, "method": "Unweighted curve_fit, nonnegative bounds, initial [200, 0.05]",
            "results": results}


def agent_diagram(out):
    fig, ax = plt.subplots(figsize=(11, 7.5))
    fig.subplots_adjust(left=0.03, right=0.97, top=0.87, bottom=0.09)
    ax.set(xlim=(0, 11), ylim=(0, 6.4)); ax.axis("off")
    fig.suptitle("Inside an analysis agent", x=0.075, ha="left", fontsize=23, fontweight="medium")
    def box(x, y, width, title, body, accent=False):
        ax.add_patch(FancyBboxPatch((x, y), width, 1.15,
            boxstyle="round,pad=0.02,rounding_size=0.09", linewidth=1,
            facecolor="#e5eeea" if accent else "#ffffff", edgecolor="#b5c8c0"))
        ax.text(x+0.18, y+0.84, title, fontsize=13, weight="bold", va="center")
        ax.text(x+0.18, y+0.43, body, fontsize=11, va="center", linespacing=1.45, color=MUTED)
    def arrow(start, end, style="arc3,rad=0"):
        ax.add_patch(FancyArrowPatch(start, end, connectionstyle=style,
            arrowstyle="-|>", mutation_scale=15, color=TEAL, linewidth=1.7))
    box(0.4, 4.65, 3, "1  Inspect the input", "Read the table, units,\ncontrols and requested endpoint.")
    box(4, 4.65, 3, "2  Choose the next action", "The model selects a tool\nand supplies its arguments.", True)
    box(7.6, 4.65, 3, "3  Run the calculation", "Software reads data, fits\na model or reports an error.")
    box(4, 2.15, 3, "4  Review the result", "Use the returned values\nto decide what comes next.", True)
    box(0.4, 2.15, 3, "Ask the researcher", "Resolve missing design facts\nbefore depending on them.")
    box(7.6, 2.15, 3, "Report the result", "Return parameters, plots\nand the unresolved limits.")
    arrow((3.44, 5.22), (3.97, 5.22))
    arrow((7.04, 5.22), (7.57, 5.22))
    arrow((9.1, 4.6), (6.8, 3.33))
    arrow((5.5, 3.34), (5.5, 4.6))
    ax.text(5.67, 3.98, "another action", fontsize=10, color=TEAL)
    arrow((3.96, 2.72), (3.44, 2.72))
    arrow((1.9, 3.34), (4.3, 4.59))
    arrow((7.04, 2.72), (7.57, 2.72))
    ax.text(0.4, 0.96, "The loop follows actual tool output.", fontsize=15, color=TEAL)
    ax.text(0.4, 0.48, "A failed calculation or an unanswered question leaves work unfinished.", fontsize=12, color=MUTED)
    fig.text(0.075, 0.045, "CONCEPTUAL DIAGRAM  ·  Simplified assay workflow  ·  Arrows do not represent timing or performance",
             fontsize=10, color=MUTED)
    save(fig, out, "agent-loop")


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--puromycin", type=Path, required=True)
    parser.add_argument("--output", type=Path, default=Path("generated"))
    args = parser.parse_args()
    args.output.mkdir(parents=True, exist_ok=True)
    result = {"versions": {"numpy": np.__version__, "scipy": scipy.__version__, "matplotlib": matplotlib.__version__},
              "donors": donor_example(args.output), "puromycin": residuals(args.puromycin, args.output)}
    agent_diagram(args.output)
    (args.output / "results.json").write_text(json.dumps(result, indent=2) + "\n")
    print(json.dumps(result, indent=2))
