#!/usr/bin/env python3
"""Generate Sepia Fig. 3 scale yield figure for the Shopify product page.

The bar chart is rebuilt as vector graphics in the same style as Figs. 1-2.
The right-hand tube/gel visuals are cropped from the existing product figure
and embedded as a temporary raster column.
"""
from __future__ import annotations

import json
from decimal import Decimal, ROUND_HALF_UP
from pathlib import Path
from statistics import mean, stdev

import matplotlib
import openpyxl

matplotlib.use("Agg")
import matplotlib.pyplot as plt
from PIL import Image

from sepia_figure_style import TEXT, style_axis


SCRIPT_VERSION = "2026-07-10.2"
FIG_W_PT = 540.0
FIG_H_PT = 540.0
FIG_W_IN = FIG_W_PT / 72.0
FIG_H_IN = FIG_H_PT / 72.0

SOURCE_IMAGE = "fig3_existing_picture1.png"
PUBLIC_WORKBOOK = "sepia_fig3_data.xlsx"
RIGHT_CROP_BOX = (1010, 0, 1200, 1246)
LYSATE_TEAL = "#37bfa7"
PURIFIED_TEAL = "#087f75"
LADDER_LABELS = {"120": 520, "30": 820, "10": 1020}

def find_input_dir() -> Path:
    here = Path(__file__).resolve().parent
    candidates = [here / "data", here, Path.cwd() / "data", Path.cwd()]
    for candidate in candidates:
        if (candidate / SOURCE_IMAGE).exists():
            return candidate
    raise FileNotFoundError(f"Could not find {SOURCE_IMAGE} in ./data or next to this script.")


def display_round(value: float) -> float:
    return float(Decimal(str(value)).quantize(Decimal("0.01"), rounding=ROUND_HALF_UP))


def load_values() -> list[dict]:
    workbook_path = Path(__file__).resolve().parents[1] / PUBLIC_WORKBOOK
    workbook = openpyxl.load_workbook(workbook_path, data_only=True, read_only=True)

    quantification = workbook["Quantification_Results"]
    quant_rows = quantification.iter_rows(values_only=True)
    quant_headers = {value: index for index, value in enumerate(next(quant_rows))}
    lysate_values = []
    for row in quant_rows:
        if row[quant_headers["Result_Type"]] != "Sample":
            continue
        value = row[quant_headers["Linear_region_mg_per_mL"]]
        if value is not None:
            lysate_values.append(float(value))

    purification = workbook["A280_Purification"]
    purification_rows = purification.iter_rows(values_only=True)
    purification_headers = {value: index for index, value in enumerate(next(purification_rows))}
    purified_values = []
    for row in purification_rows:
        concentration = row[purification_headers["A280_concentration_mg_per_mL"]]
        pre_mass = row[purification_headers["Container_pre_g"]]
        post_mass = row[purification_headers["Container_post_g"]]
        density = row[purification_headers["Density_g_per_mL"]]
        input_lysate = row[purification_headers["Input_lysate_mL"]]
        if None in (concentration, pre_mass, post_mass, density, input_lysate):
            continue
        recovered_volume = (float(post_mass) - float(pre_mass)) / float(density)
        total_mg = float(concentration) * recovered_volume
        purified_values.append(total_mg / float(input_lysate))

    workbook.close()
    if len(lysate_values) != 3 or len(purified_values) != 3:
        raise ValueError(
            f"Expected three lysate and three purified replicates; found "
            f"{len(lysate_values)} and {len(purified_values)}."
        )

    rows = [
        ("In lysate", lysate_values, LYSATE_TEAL),
        ("Recovered after\npurification", purified_values, PURIFIED_TEAL),
    ]
    return [
        {
            "label": label,
            "mean": display_round(mean(replicates)),
            "sd": display_round(stdev(replicates)),
            "unrounded_mean": mean(replicates),
            "unrounded_sd": stdev(replicates),
            "replicates": replicates,
            "color": color,
        }
        for label, replicates, color in rows
    ]


def crop_right_column(input_dir: Path, output_dir: Path) -> Image.Image:
    source = Image.open(input_dir / SOURCE_IMAGE).convert("RGBA")
    crop = source.crop(RIGHT_CROP_BOX)
    output_dir.mkdir(parents=True, exist_ok=True)
    crop.save(output_dir / "sepia_fig3_right_column_crop.png")
    return crop


def make_figure(input_dir: Path, out_svg: Path, out_png: Path | None = None) -> None:
    output_dir = out_svg.parent
    output_dir.mkdir(parents=True, exist_ok=True)
    right_crop = crop_right_column(input_dir, output_dir)
    values = load_values()

    fig = plt.figure(figsize=(FIG_W_IN, FIG_H_IN), facecolor="none")
    ax = fig.add_axes([0.125, 0.170, 0.575, 0.730])
    style_axis(ax)

    labels = [row["label"] for row in values]
    means = [row["mean"] for row in values]
    sds = [row["sd"] for row in values]
    colors = [row["color"] for row in values]
    x = [0, 1]

    ax.bar(
        x,
        means,
        yerr=sds,
        width=0.48,
        color=colors,
        edgecolor=TEXT,
        linewidth=1.55,
        error_kw={"elinewidth": 1.8, "capthick": 1.8, "capsize": 4.5, "ecolor": TEXT},
        zorder=3,
    )
    for xi, mean, sd in zip(x, means, sds):
        ax.text(
            xi,
            mean + sd + 0.12,
            f"{mean:.2f} \u00b1 {sd:.2f}\nmg/mL input",
            color=TEXT,
            fontsize=17,
            fontweight="bold",
            ha="center",
            va="bottom",
            linespacing=0.95,
            zorder=5,
        )

    ax.set_xlim(-0.55, 1.55)
    ax.set_ylim(0, 2.5)
    ax.set_yticks([0, 0.5, 1.0, 1.5, 2.0, 2.5])
    ax.set_xticks(x)
    ax.set_xticklabels(labels)
    ax.set_xlabel("")
    ax.set_ylabel("sfGFP yield (mg/mL input)", labelpad=8)

    for tick_label in ax.get_xticklabels() + ax.get_yticklabels():
        tick_label.set_fontweight("bold")
    ax.yaxis.label.set_fontweight("bold")

    image_ax = fig.add_axes([0.770, 0.075, 0.185, 0.850])
    image_ax.set_axis_off()
    image_ax.imshow(right_crop)
    image_ax.text(
        -36,
        352,
        "kDa",
        color=TEXT,
        fontsize=8.5,
        ha="left",
        va="center",
        clip_on=False,
    )
    image_ax.text(
        39,
        332,
        "L",
        color=TEXT,
        fontsize=9.5,
        fontweight="bold",
        ha="center",
        va="bottom",
        clip_on=False,
    )
    image_ax.text(
        145,
        332,
        "sfGFP",
        color=TEXT,
        fontsize=9.5,
        fontweight="bold",
        ha="center",
        va="bottom",
        clip_on=False,
    )
    for label, y_pos in LADDER_LABELS.items():
        image_ax.plot([-7, 2], [y_pos, y_pos], color=TEXT, linewidth=1.05, solid_capstyle="butt", clip_on=False)
        image_ax.text(
            -12,
            y_pos,
            label,
            color=TEXT,
            fontsize=8.5,
            fontweight="bold",
            ha="right",
            va="center",
            clip_on=False,
        )

    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)

    summary = [
        {
            "Condition": row["label"],
            "Mean_sfGFP_mg_per_mL_input": row["mean"],
            "SD_sfGFP_mg_per_mL_input": row["sd"],
            "Color": row["color"],
        }
        for row in values
    ]
    (output_dir / "sepia_fig3_scale_yield_summary_used.json").write_text(json.dumps(summary, indent=2))

    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. 3 asset; template image width=540 height=540",
        "left_chart_values": summary,
        "y_axis": "sfGFP yield (mg/mL input lysate)",
        "x_axis": "Condition",
        "right_column_source": SOURCE_IMAGE,
        "quantification_source": f"{PUBLIC_WORKBOOK} / Quantification_Results and A280_Purification",
        "recovery_calculation": "A280 concentration × ((post mass - pre mass) ÷ 1.15 g/mL) ÷ 0.750 mL input lysate.",
        "display_rounding": "ROUND_HALF_UP to 2 decimals.",
        "right_column_crop_box_px": RIGHT_CROP_BOX,
        "right_column_ladder_labels_kDa": LADDER_LABELS,
        "theme_notes": "Left chart matches Figs. 1-2 typography/axes with transparent background; right column uses the pre-raw-image composite crop with three outside ladder labels and lane labels.",
    }
    (output_dir / "sepia_fig3_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_fig3_scale_yield.svg",
        out_png=out_dir / "sepia_fig3_scale_yield.png",
    )
    print(f"Wrote {out_dir / 'sepia_fig3_scale_yield.svg'}")


if __name__ == "__main__":
    main()