"""
Create a  notched boxplot of Cytoplasmic POLK size (area) with annotations from
repeated downsampling (median p-value and median Cohen's d), using files in the
same directory as this script.

Inputs expected in the same folder as this script:
- Fig6E_data.csv

Outputs written to the same folder:

- boxplot Fig6E
- CytoPolk_downsample_cohensd_boxplot.pdf
- CytoPolk_downsample_pvalues_boxplot.pdf
- CytoPolk_downsample_summary.csv

What the script does:
1. Repeats non-parametric downsampling within each age_grp × exposure stratum.
2. Saves the summary table with p_median and d_median.
3. Draws a two-panel notched boxplot using the full data.
4. Annotates each exposure group with p_median and d_median from the
   downsampled summary table.
5. Saves additional PDF boxplots showing the distributions of downsampled
   p-values and Cohen's d values by stratum.
"""

from pathlib import Path

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from matplotlib.gridspec import GridSpec
from scipy.stats import mannwhitneyu

# ---------------- USER SETTINGS ----------------
INPUT_STEM = "Fig6E_data"
VALUE_COL = "CytoPolk_Area"
AGE_COL = "age_grp"
TREATMENT_COL = "treatment"
EXPOSURE_COL = "exposure"

AGE_LEVELS = ["young", "old"]
DISPLAY_AGE = {"young": "young", "old": "early old"}
EXPOSURES = [80, 160]
TREATMENTS = ["KA", "CTRL"]

N_SUBSAMPLES = 200
N_PER_GROUP = 100
RANDOM_SEED = 12345

SUMMARY_CSV = "CytoPolk_downsample_summary.csv"
MAIN_JPG = "Fig6E.jpg"
MAIN_PDF = "Fig6E.pdf"
PVAL_PDF = "CytoPolk_downsample_pvalues_boxplot.pdf"
D_PDF = "CytoPolk_downsample_cohensd_boxplot.pdf"
# ------------------------------------------------

PANEL_BG = "#FFFFFF"
STRIP_BG = "#CFCFCF"
BOX_COLORS = {"KA": "#B8D1AE", "CTRL": "#C7B9D2"}
EDGE = "#3A3A3A"

BOXPLOT_ORDER = [(80, "KA"), (80, "CTRL"), (160, "KA"), (160, "CTRL")]
BOXPLOT_POSITIONS = [0.8, 1.2, 1.8, 2.2]
BOXPLOT_WIDTH = 0.34
MAIN_PANEL_X_CENTERS = [1.0, 2.0]
MAIN_PANEL_P_Y = 198
MAIN_PANEL_D_Y = 187
AUXILIARY_FILL_COLORS = ["#B8D1AE", "#C7B9D2", "#B8D1AE", "#C7B9D2"]


def script_dir() -> Path:
    return Path(__file__).resolve().parent



def find_input_csv(base_dir: Path) -> Path:
    exact = base_dir / f"{INPUT_STEM}.csv"
    if exact.exists():
        return exact

    matches = sorted(base_dir.glob(f"{INPUT_STEM}*.csv"))
    if matches:
        return matches[0]

    raise FileNotFoundError(
        f"Could not find '{INPUT_STEM}.csv' (or any '{INPUT_STEM}*.csv') in {base_dir}"
    )



def cohens_d(x: np.ndarray, y: np.ndarray) -> float:
    x = np.asarray(x, dtype=float)
    y = np.asarray(y, dtype=float)
    if len(x) < 2 or len(y) < 2:
        return np.nan

    sx = x.std(ddof=1)
    sy = y.std(ddof=1)
    pooled_var = ((len(x) - 1) * sx**2 + (len(y) - 1) * sy**2) / (len(x) + len(y) - 2)
    if pooled_var <= 0:
        return np.nan
    return (x.mean() - y.mean()) / np.sqrt(pooled_var)



def format_p(p: float) -> str:
    if pd.isna(p):
        return "NA"
    if p < 1e-3:
        return f"{p:.2e}"
    if p < 0.1:
        return f"{p:.3f}".rstrip("0").rstrip(".")
    return f"{p:.3g}"



def _validate_required_columns(df: pd.DataFrame) -> None:
    required_columns = {AGE_COL, TREATMENT_COL, EXPOSURE_COL, VALUE_COL}
    missing_columns = required_columns.difference(df.columns)
    if missing_columns:
        raise KeyError(f"Missing required columns: {sorted(missing_columns)}")



def load_and_filter_data(csv_path: Path) -> pd.DataFrame:
    df = pd.read_csv(csv_path)
    _validate_required_columns(df)

    df = df.dropna(subset=[VALUE_COL])

    if df.empty:
        raise RuntimeError("No rows remain after filtering.")
    return df



def _values_for_stratum(
    df: pd.DataFrame,
    age: str,
    exposure: int,
    treatment: str,
) -> np.ndarray:
    mask = (
        (df[AGE_COL] == age)
        & (df[EXPOSURE_COL] == exposure)
        & (df[TREATMENT_COL] == treatment)
    )
    return df.loc[mask, VALUE_COL].to_numpy()



def _subset_for_age_and_exposure(df: pd.DataFrame, age: str, exposure: int) -> pd.DataFrame:
    return df[(df[AGE_COL] == age) & (df[EXPOSURE_COL] == exposure)]



def _mannwhitney_pvalue(x: np.ndarray, y: np.ndarray) -> float:
    p_value = mannwhitneyu(x, y, alternative="two-sided").pvalue
    if p_value == 0:
        p_value = np.nextafter(0, 1)
    return p_value



def _comparison_label(age: str, exposure: int) -> str:
    return f"{age}_{exposure}"



def _build_subsampling_record(
    age: str,
    exposure: int,
    iteration: int,
    ka_sample: np.ndarray,
    ctrl_sample: np.ndarray,
) -> dict:
    return {
        "age_grp": age,
        "display_age": DISPLAY_AGE.get(age, age),
        "exposure": exposure,
        "comparison": _comparison_label(age, exposure),
        "iteration": iteration,
        "p_value": _mannwhitney_pvalue(ka_sample, ctrl_sample),
        "cohens_d": cohens_d(ka_sample, ctrl_sample),
    }



def run_subsampling(df: pd.DataFrame) -> pd.DataFrame:
    rng = np.random.default_rng(RANDOM_SEED)
    rows = []

    for age in AGE_LEVELS:
        for exposure in EXPOSURES:
            ka_values = _values_for_stratum(df, age, exposure, "KA")
            ctrl_values = _values_for_stratum(df, age, exposure, "CTRL")

            if len(ka_values) < N_PER_GROUP or len(ctrl_values) < N_PER_GROUP:
                raise ValueError(
                    f"Not enough rows for {age}, exposure {exposure}: "
                    f"KA={len(ka_values)}, CTRL={len(ctrl_values)}, need at least {N_PER_GROUP} each."
                )

            for iteration in range(1, N_SUBSAMPLES + 1):
                ka_sample = rng.choice(ka_values, size=N_PER_GROUP, replace=False)
                ctrl_sample = rng.choice(ctrl_values, size=N_PER_GROUP, replace=False)
                rows.append(
                    _build_subsampling_record(age, exposure, iteration, ka_sample, ctrl_sample)
                )

    return pd.DataFrame(rows)



def _build_summary_record(
    age: str,
    exposure: int,
    sub_res: pd.DataFrame,
    ka_values: np.ndarray,
    ctrl_values: np.ndarray,
) -> dict:
    return {
        "age_grp": age,
        "display_age": DISPLAY_AGE.get(age, age),
        "exposure": exposure,
        "comparison": _comparison_label(age, exposure),
        "n_ka_available": len(ka_values),
        "n_ctrl_available": len(ctrl_values),
        "n_per_group": N_PER_GROUP,
        "n_subsamples": len(sub_res),
        "p_full_data": _mannwhitney_pvalue(ka_values, ctrl_values),
        "d_full_data": cohens_d(ka_values, ctrl_values),
        "p_min": sub_res["p_value"].min(),
        "p_25": sub_res["p_value"].quantile(0.25),
        "p_median": sub_res["p_value"].median(),
        "p_75": sub_res["p_value"].quantile(0.75),
        "p_max": sub_res["p_value"].max(),
        "d_min": sub_res["cohens_d"].min(),
        "d_25": sub_res["cohens_d"].quantile(0.25),
        "d_median": sub_res["cohens_d"].median(),
        "d_75": sub_res["cohens_d"].quantile(0.75),
        "d_max": sub_res["cohens_d"].max(),
    }



def summarize_results(res_df: pd.DataFrame, df: pd.DataFrame) -> pd.DataFrame:
    summary_rows = []

    for age in AGE_LEVELS:
        for exposure in EXPOSURES:
            comparison = _comparison_label(age, exposure)
            sub_res = res_df[res_df["comparison"] == comparison]

            ka_values = _values_for_stratum(df, age, exposure, "KA")
            ctrl_values = _values_for_stratum(df, age, exposure, "CTRL")

            summary_rows.append(
                _build_summary_record(age, exposure, sub_res, ka_values, ctrl_values)
            )

    return pd.DataFrame(summary_rows)



def _style_axis(ax, show_ylabel: bool = False):
    ax.set_facecolor(PANEL_BG)
    for spine in ax.spines.values():
        spine.set_color(EDGE)
        spine.set_linewidth(1.8)

    ax.tick_params(axis="both", labelsize=14, width=1.8, length=6, colors=EDGE)
    ax.set_ylim(-10, 206)
    ax.set_yticks([0, 50, 100, 150, 200])
    ax.set_xlabel("exposure (mins)", fontsize=18, labelpad=10)
    if show_ylabel:
        ax.set_ylabel("Cytoplasmic POLK size (area)", fontsize=18, labelpad=2)
    else:
        ax.set_ylabel("")

    ax.grid(False)



def _draw_strip(ax, label: str):
    ax.set_facecolor(STRIP_BG)
    for spine in ax.spines.values():
        spine.set_color(EDGE)
        spine.set_linewidth(1.8)
    ax.set_xticks([])
    ax.set_yticks([])
    ax.text(
        0.5,
        0.5,
        label,
        rotation=-90,
        ha="center",
        va="center",
        fontsize=18,
        color=EDGE,
    )



def _boxplot_arrays_for_panel(df_panel: pd.DataFrame):
    data = []
    for exposure, treatment in BOXPLOT_ORDER:
        mask = (
            (df_panel[EXPOSURE_COL] == exposure)
            & (df_panel[TREATMENT_COL] == treatment)
        )
        values = df_panel.loc[mask, VALUE_COL].to_numpy()
        data.append(values)
    return data



def _draw_grouped_notched_boxplot(ax, df_panel: pd.DataFrame):
    data = _boxplot_arrays_for_panel(df_panel)

    bp = ax.boxplot(
        data,
        positions=BOXPLOT_POSITIONS,
        widths=BOXPLOT_WIDTH,
        notch=True,
        patch_artist=True,
        showfliers=False,
        whis=1.5,
        medianprops=dict(color=EDGE, linewidth=2.0),
        whiskerprops=dict(color=EDGE, linewidth=1.5),
        capprops=dict(color=EDGE, linewidth=0),
        boxprops=dict(edgecolor=EDGE, linewidth=1.5),
    )

    for patch, (_, treatment) in zip(bp["boxes"], BOXPLOT_ORDER):
        patch.set_facecolor(BOX_COLORS[treatment])

    for position, values in zip(BOXPLOT_POSITIONS, data):
        median_y = np.median(values)
        ax.hlines(median_y, position - 0.08, position + 0.08, color=EDGE, linewidth=3.6, zorder=4)

    ax.set_xlim(0.4, 2.6)
    ax.set_xticks([1.0, 2.0])
    ax.set_xticklabels(["80", "160"], fontsize=16)



def _panel_summary_for_age(summary_df: pd.DataFrame, age: str) -> pd.DataFrame:
    return summary_df[summary_df["age_grp"] == age].sort_values("exposure")



def _annotate_main_panel(ax, panel_summary: pd.DataFrame) -> None:
    for x_center, (_, row) in zip(MAIN_PANEL_X_CENTERS, panel_summary.iterrows()):
        ax.text(
            x_center,
            MAIN_PANEL_P_Y,
            f"p_median={format_p(row['p_median'])}",
            ha="center",
            va="bottom",
            fontsize=10,
            fontweight="bold",
            color="black",
        )
        ax.text(
            x_center,
            MAIN_PANEL_D_Y,
            f"d_median={row['d_median']:.2f}",
            ha="center",
            va="bottom",
            fontsize=9,
            fontweight="bold",
            color="black",
        )



def make_main_figure(df: pd.DataFrame, summary_df: pd.DataFrame, out_jpg: Path, out_pdf: Path):
    fig = plt.figure(figsize=(9.2, 6.63), facecolor="white")
    outer = GridSpec(1, 2, width_ratios=[1, 1], wspace=0.30, figure=fig)

    for idx, age in enumerate(AGE_LEVELS):
        inner = outer[idx].subgridspec(1, 2, width_ratios=[1, 0.15], wspace=0.0)
        ax = fig.add_subplot(inner[0, 0])
        strip_ax = fig.add_subplot(inner[0, 1])

        panel_df = df[df[AGE_COL] == age]
        panel_summary = _panel_summary_for_age(summary_df, age)

        _style_axis(ax, show_ylabel=(idx == 0))
        _draw_grouped_notched_boxplot(ax, panel_df)
        _draw_strip(strip_ax, DISPLAY_AGE.get(age, age))
        _annotate_main_panel(ax, panel_summary)

    plt.subplots_adjust(left=0.13, right=0.98, top=0.97, bottom=0.14)
    fig.savefig(out_jpg, dpi=600)
    fig.savefig(out_pdf)
    plt.close(fig)



def _styled_simple_boxplot(ax, data, labels, ylabel, title, log_scale=False, hline_zero=False):
    ax.set_facecolor(PANEL_BG)
    bp = ax.boxplot(
        data,
        notch=True,
        patch_artist=True,
        showfliers=False,
        widths=0.6,
        medianprops=dict(color=EDGE, linewidth=2.0),
        whiskerprops=dict(color=EDGE, linewidth=1.5),
        capprops=dict(color=EDGE, linewidth=1.5),
        boxprops=dict(edgecolor=EDGE, linewidth=1.5),
    )

    for patch, color in zip(bp["boxes"], AUXILIARY_FILL_COLORS):
        patch.set_facecolor(color)

    if log_scale:
        ax.set_yscale("log")
    if hline_zero:
        ax.axhline(0, color=EDGE, linewidth=1.4, linestyle="--")

    for spine in ax.spines.values():
        spine.set_color(EDGE)
        spine.set_linewidth(1.5)

    ax.set_xticklabels(labels, rotation=20, ha="right", fontsize=11)
    ax.set_ylabel(ylabel, fontsize=14)
    ax.set_title(title, fontsize=16, pad=10)
    ax.tick_params(axis="both", labelsize=11, width=1.5, colors=EDGE)
    ax.grid(False)



def _auxiliary_labels() -> list[str]:
    labels = []
    for age in AGE_LEVELS:
        for exposure in EXPOSURES:
            labels.append(f"{age} {exposure}")
    return labels



def _auxiliary_data(res_df: pd.DataFrame, column: str, dropna: bool = False) -> list[np.ndarray]:
    data = []
    for age in AGE_LEVELS:
        for exposure in EXPOSURES:
            mask = (res_df["age_grp"] == age) & (res_df["exposure"] == exposure)
            values = res_df.loc[mask, column]
            if dropna:
                values = values.dropna()
            data.append(values.to_numpy())
    return data



def _annotate_auxiliary_medians(ax, box, data: list[np.ndarray], fmt: str, prefix: str) -> None:
    for i, values in enumerate(data):
        median_value = np.median(values)
        median_line = box["medians"][i]
        x_pos = np.mean(median_line.get_xdata())
        y_pos = np.mean(median_line.get_ydata())
        ax.text(
            x_pos,
            y_pos,
            f"{prefix}{format(median_value, fmt)}",
            ha="center",
            va="bottom",
            fontsize=8,
        )



def make_pvalue_boxplot(res_df: pd.DataFrame, out_pdf: Path):
    """Match the original auxiliary p-value plot style."""
    labels = _auxiliary_labels()
    data = _auxiliary_data(res_df, column="p_value", dropna=False)

    fig, ax = plt.subplots(figsize=(6, 6))
    box = ax.boxplot(data, labels=labels, showfliers=False)
    ax.set_yscale("log")
    ax.set_ylabel("Mann–Whitney p-value (log scale)")
    ax.set_title(f"Distribution of p-values\n({N_SUBSAMPLES} subsamples, {N_PER_GROUP} cells/group)")

    _annotate_auxiliary_medians(ax, box, data, fmt=".2g", prefix="median p = ")

    plt.tight_layout()
    fig.savefig(out_pdf, bbox_inches="tight")
    plt.close(fig)



def make_cohensd_boxplot(res_df: pd.DataFrame, out_pdf: Path):
    """Match the original auxiliary Cohen's d plot style."""
    labels = _auxiliary_labels()
    data = _auxiliary_data(res_df, column="cohens_d", dropna=True)

    fig, ax = plt.subplots(figsize=(6, 6))
    box = ax.boxplot(data, labels=labels, showfliers=False)
    ax.axhline(0, linestyle="--")
    ax.set_ylabel("Cohen's d (KA − CTRL)")
    ax.set_title(f"Distribution of Cohen's d\n({N_SUBSAMPLES} subsamples, {N_PER_GROUP} cells/group)")

    _annotate_auxiliary_medians(ax, box, data, fmt=".2f", prefix="median d = ")

    plt.tight_layout()
    fig.savefig(out_pdf, bbox_inches="tight")
    plt.close(fig)



def _sorted_summary(summary_df: pd.DataFrame) -> pd.DataFrame:
    summary_df = summary_df.copy()
    summary_df["age_grp"] = pd.Categorical(summary_df["age_grp"], categories=AGE_LEVELS, ordered=True)
    return summary_df.sort_values(["age_grp", "exposure"]).reset_index(drop=True)



def main():
    base_dir = script_dir()
    csv_path = find_input_csv(base_dir)
    df = load_and_filter_data(csv_path)

    res_df = run_subsampling(df)
    summary_df = summarize_results(res_df, df)
    summary_df = _sorted_summary(summary_df)

    summary_df.to_csv(base_dir / SUMMARY_CSV, index=False)
    make_main_figure(df, summary_df, base_dir / MAIN_JPG, base_dir / MAIN_PDF)
    make_pvalue_boxplot(res_df, base_dir / PVAL_PDF)
    make_cohensd_boxplot(res_df, base_dir / D_PDF)

    print(f"Input CSV: {csv_path}")
    print(f"Saved: {base_dir / SUMMARY_CSV}")
    print(f"Saved: {base_dir / MAIN_JPG}")
    print(f"Saved: {base_dir / MAIN_PDF}")
    print(f"Saved: {base_dir / PVAL_PDF}")
    print(f"Saved: {base_dir / D_PDF}")



if __name__ == "__main__":
    main()
