#!/usr/bin/env python3
"""
Generate Sepia Fig. 1 / DNA template figure from real processed CSV data.

Inputs expected in ./data/ or in the current working directory:
  - 260514_fig2_endpoint_well_table.csv
  - 260514_fig2_processed_long.csv

The figure uses real replicate-level endpoint data and kinetic replicate data.
Endpoint means and SD error bars are recomputed from the endpoint well table;
kinetic traces are recomputed from the long-format replicate table.
"""
from __future__ import annotations

from pathlib import Path
import math
import numpy as np
import pandas as pd
import matplotlib

matplotlib.use("Agg")
import matplotlib.pyplot as plt

from sepia_figure_style import (
    LINEAR_GREEN,
    PLASMID_GREEN,
    TEXT,
    style_axis,
)

# Match the current Shopify product-page Fig. 1 image slot.
# The OpenCFPS template declares width="1100" height="432" for this figure.
SCRIPT_VERSION = "2026-06-09.1"
FIG_W_PT = 1100.0
FIG_H_PT = 432.0
FIG_W_IN = FIG_W_PT / 72.0
FIG_H_IN = FIG_H_PT / 72.0

# Manual layout tuned for the wide, shallow Shopify figure slot.
LEFT_AX = [0.075, 0.205, 0.365, 0.620]
RIGHT_AX = [0.548, 0.205, 0.355, 0.620]
TARGETS = {
    ("Plasmid", 4.0): {"label": "Plasmid", "color": PLASMID_GREEN, "marker": "o"},
    ("Linear", 8.0): {"label": "Linear", "color": LINEAR_GREEN, "marker": "s"},
}

LABEL_BOX = None


def find_input_dir() -> Path:
    """Find input CSVs next to the script or in ./data."""
    here = Path(__file__).resolve().parent
    candidates = [here / "data", here, Path.cwd() / "data", Path.cwd()]
    for candidate in candidates:
        if (candidate / "260514_fig2_endpoint_well_table.csv").exists() and (
            candidate / "260514_fig2_processed_long.csv"
        ).exists():
            return candidate
    raise FileNotFoundError(
        "Could not find required CSVs. Put the two input CSV files in ./data or next to this script."
    )


def summarize_endpoint(endpoint_wells: pd.DataFrame) -> pd.DataFrame:
    required = {"DNA_Type", "DNA_Actual_nM", "RFE_uM"}
    missing = required - set(endpoint_wells.columns)
    if missing:
        raise ValueError(f"Endpoint table missing columns: {sorted(missing)}")
    summary = (
        endpoint_wells.groupby(["DNA_Type", "DNA_Actual_nM"], as_index=False)
        .agg(N=("RFE_uM", "count"), Mean_RFE_uM=("RFE_uM", "mean"), SD_RFE_uM=("RFE_uM", "std"))
        .sort_values(["DNA_Type", "DNA_Actual_nM"])
    )
    summary["SD_RFE_uM"] = summary["SD_RFE_uM"].fillna(0.0)
    summary["x_nM"] = summary["DNA_Actual_nM"].astype(float)
    return summary


def summarize_kinetics(long_df: pd.DataFrame) -> pd.DataFrame:
    required = {"Assay", "Time_h", "DNA_Type", "DNA_Actual_nM", "RFE_uM", "Raw_610_650_RFU"}
    missing = required - set(long_df.columns)
    if missing:
        raise ValueError(f"Long kinetic table missing columns: {sorted(missing)}")
    df = long_df.loc[long_df["Assay"].eq("TXTL")].copy()
    keep = pd.Series(False, index=df.index)
    for dna_type, dna_nM in TARGETS:
        keep |= df["DNA_Type"].eq(dna_type) & np.isclose(df["DNA_Actual_nM"], dna_nM)
    df = df.loc[keep]
    if df.empty:
        raise ValueError("No kinetic rows found for target plasmid 4 nM and linear 8 nM conditions.")
    kinetic = (
        df.groupby(["Time_h", "DNA_Type", "DNA_Actual_nM"], as_index=False)
        .agg(
            N=("RFE_uM", "count"),
            Mean_RFE_uM=("RFE_uM", "mean"),
            SD_RFE_uM=("RFE_uM", "std"),
            Mean_MgAptamer_RFU=("Raw_610_650_RFU", "mean"),
            SD_MgAptamer_RFU=("Raw_610_650_RFU", "std"),
        )
        .sort_values(["DNA_Type", "DNA_Actual_nM", "Time_h"])
    )
    return kinetic.fillna(0.0)


def make_figure(input_dir: Path, out_svg: Path, out_png: Path | None = None) -> None:
    endpoint_wells = pd.read_csv(input_dir / "260514_fig2_endpoint_well_table.csv")
    long_df = pd.read_csv(input_dir / "260514_fig2_processed_long.csv")
    endpoint = summarize_endpoint(endpoint_wells)
    kinetic = summarize_kinetics(long_df)

    # Write computed summaries next to the figure so the plotted values can be inspected.
    summary_dir = out_svg.parent
    summary_dir.mkdir(parents=True, exist_ok=True)
    endpoint.to_csv(summary_dir / "sepia_fig1_endpoint_summary_used.csv", index=False)
    kinetic.to_csv(summary_dir / "sepia_fig1_kinetics_summary_used.csv", index=False)

    fig = plt.figure(figsize=(FIG_W_IN, FIG_H_IN), facecolor="none")

    # ------------------------------------------------------------------
    # Panel A: endpoint DNA concentration response, linear x axis.
    # ------------------------------------------------------------------
    ax_a = fig.add_axes(LEFT_AX)
    style_axis(ax_a)

    order = ["Plasmid", "Linear"]
    for dna_type in order:
        sub = endpoint.loc[endpoint["DNA_Type"].eq(dna_type)].sort_values("x_nM")
        color = PLASMID_GREEN if dna_type == "Plasmid" else LINEAR_GREEN
        marker = "o" if dna_type == "Plasmid" else "s"
        linestyle = "-"
        face = color if dna_type == "Plasmid" else "white"
        ax_a.errorbar(
            sub["x_nM"],
            sub["Mean_RFE_uM"],
            yerr=sub["SD_RFE_uM"],
            color=color,
            linestyle=linestyle,
            linewidth=3.4,
            marker=marker,
            markersize=9.5,
            markerfacecolor=face,
            markeredgecolor=color,
            markeredgewidth=1.55,
            capsize=4.5,
            elinewidth=1.8,
            capthick=1.8,
            zorder=3,
        )

    ax_a.set_xticks([0, 4, 8, 16, 32])
    ax_a.set_xlim(-0.8, 33.5)
    ax_a.set_ylim(0, 50)
    ax_a.set_yticks([0, 25, 50])
    ax_a.set_xlabel("DNA concentration (nM)", labelpad=8)
    ax_a.set_ylabel("Relative Fluorescein Equivalents (µM)", labelpad=8)

    ax_a.text(
        7.2,
        44.8,
        "Plasmid DNA",
        color=TEXT,
        fontsize=17,
        fontweight="bold",
        va="center",
        bbox=LABEL_BOX,
    )
    ax_a.text(
        19.2,
        27.0,
        "Linear DNA",
        color=TEXT,
        fontsize=17,
        fontweight="bold",
        va="center",
        bbox=LABEL_BOX,
    )

    # ------------------------------------------------------------------
    # Panel B: kinetic traces, two selected template conditions.
    # TL = GFP/RFE on left axis; TX = MG aptamer fluorescence on right axis.
    # ------------------------------------------------------------------
    ax_b = fig.add_axes(RIGHT_AX)
    style_axis(ax_b)
    ax_b2 = ax_b.twinx()
    style_axis(ax_b2, show_right_spine=True)
    ax_b2.grid(False)
    ax_b2.tick_params(axis="y", color=TEXT, labelcolor=TEXT, length=5.5, width=1.35)

    for (dna_type, dna_nM), spec in TARGETS.items():
        sub = kinetic.loc[
            kinetic["DNA_Type"].eq(dna_type) & np.isclose(kinetic["DNA_Actual_nM"], dna_nM)
        ].sort_values("Time_h")
        t = sub["Time_h"].to_numpy(float)
        tl = sub["Mean_RFE_uM"].to_numpy(float)
        tl_sd = sub["SD_RFE_uM"].to_numpy(float)
        tx = sub["Mean_MgAptamer_RFU"].to_numpy(float)
        tx_sd = sub["SD_MgAptamer_RFU"].to_numpy(float)
        color = spec["color"]
        marker = spec["marker"]
        marker_face = color if dna_type == "Plasmid" else "white"
        # TL / translation / GFP output.
        ax_b.plot(
            t,
            tl,
            color=color,
            linestyle="-",
            linewidth=3.4,
            marker=marker,
            markersize=6.2,
            markevery=3,
            markerfacecolor=marker_face,
            markeredgecolor=color,
            markeredgewidth=1.35,
            zorder=4,
        )
        # Subtle SD bands show replicate variation without extra chart furniture.
        tl_band_color = color
        ax_b.fill_between(t, tl - tl_sd, tl + tl_sd, color=tl_band_color, alpha=0.10, linewidth=0, zorder=2)
        # TX / transcription / MG aptamer fluorescence output.
        tx_color = color
        ax_b2.plot(
            t,
            tx,
            color=tx_color,
            linestyle=(0, (1.2, 2.2)),
            linewidth=2.85,
            marker=marker,
            markersize=5.2,
            markevery=3,
            markerfacecolor=marker_face,
            markeredgecolor=tx_color,
            markeredgewidth=1.2,
            zorder=3,
        )
        ax_b2.fill_between(t, tx - tx_sd, tx + tx_sd, color=tx_color, alpha=0.10, linewidth=0, zorder=1)

    x_max = float(kinetic["Time_h"].max())
    ax_b.set_xlim(0, math.ceil(x_max))
    ax_b.set_xticks([0, 6, 12])
    ax_b.set_ylim(0, 50)
    ax_b.set_yticks([0, 25, 50])
    ax_b2.set_ylim(0, 15000)
    ax_b2.set_yticks([0, 7500, 15000])
    ax_b2.set_yticklabels([f"{int(v):,}" for v in [0, 7500, 15000]])
    ax_b.set_xlabel("Time (h)", labelpad=8)
    ax_b.set_ylabel("Relative Fluorescein Equivalents (µM)", labelpad=8)
    ax_b2.set_ylabel("MG aptamer fluorescence (RFU)", labelpad=8)

    ax_b.text(
        7.25,
        52.2,
        "Plasmid TL",
        color=TEXT,
        fontsize=17,
        fontweight="bold",
        va="center",
        bbox=LABEL_BOX,
    )
    ax_b.text(
        8.35,
        27.2,
        "Linear TL",
        color=TEXT,
        fontsize=17,
        fontweight="bold",
        va="center",
        bbox=LABEL_BOX,
    )
    ax_b2.text(
        0.70,
        12750,
        "Plasmid TX",
        color=TEXT,
        fontsize=17,
        fontweight="bold",
        va="center",
        bbox=LABEL_BOX,
    )
    ax_b2.text(
        0.58,
        9200,
        "Linear TX",
        color=TEXT,
        fontsize=17,
        fontweight="bold",
        va="center",
        bbox=LABEL_BOX,
    )

    for axis in (ax_a, ax_b, ax_b2):
        axis.xaxis.label.set_fontweight("bold")
        axis.yaxis.label.set_fontweight("bold")
        for tick_label in axis.get_xticklabels() + axis.get_yticklabels():
            tick_label.set_fontweight("bold")

    out_svg.parent.mkdir(parents=True, exist_ok=True)
    fig.savefig(out_svg, format="svg", transparent=True, bbox_inches=None, metadata={"Date": None})
    if out_png is not None:
        fig.savefig(out_png, dpi=220, transparent=True, bbox_inches=None)

    manifest = {
        "figure": out_svg.name,
        "script": Path(__file__).name,
        "script_version": SCRIPT_VERSION,
        "width_pt": FIG_W_PT,
        "height_pt": FIG_H_PT,
        "site_slot": "OpenCFPS product template Fig. 1 asset; template image width=1100 height=432",
        "endpoint_source": "260514_fig2_endpoint_well_table.csv",
        "kinetic_source": "260514_fig2_processed_long.csv",
        "endpoint_summary": "sepia_fig1_endpoint_summary_used.csv",
        "kinetics_summary": "sepia_fig1_kinetics_summary_used.csv",
        "endpoint_error_bars": "sample SD across replicate wells",
        "kinetic_traces": "mean across replicate wells; shaded bands are sample SD for TL and TX traces",
        "kinetic_conditions": {"plasmid_nM": 4.0, "linear_nM": 8.0},
        "endpoint_x_axis": "linear DNA concentration in nM",
        "right_y_axis": "MG aptamer fluorescence (RFU), computed from Raw_610_650_RFU",
        "visual_encoding": "DNA type controls trace color and marker in both panels: plasmid is dark green with filled circle markers, linear is light green with open square markers; all direct labels are black; TX traces use dotted lines.",
        "theme_notes": "Sepia site style: Inter/Lexend-like sans-serif, transparent background, black label text, coordinated green DNA-type palette."
    }
    import json
    (summary_dir / "sepia_fig1_reproducibility_manifest.json").write_text(json.dumps(manifest, indent=2))
    plt.close(fig)


def main() -> None:
    input_dir = find_input_dir()
    out_dir = Path(__file__).resolve().parent / "outputs"
    make_figure(
        input_dir=input_dir,
        out_svg=out_dir / "sepia_fig1_dna_endpoints_kinetics.svg",
        out_png=out_dir / "sepia_fig1_dna_endpoints_kinetics.png",
    )
    print(f"Wrote {out_dir / 'sepia_fig1_dna_endpoints_kinetics.svg'}")


if __name__ == "__main__":
    main()
