#!/usr/bin/env python3
"""
Creates a notched boxplot of Median_cfos 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:
- Fig6B_data.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 Median_cfos 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.

Outputs written to the same folder:
- downsampled_p_d_summary.csv
- Fig6B_downsampled_notched_boxplot.pdf
- Fig6B_downsampled_notched_boxplot.jpg
- downsampled_p_boxplot.pdf
- downsampled_cohensd_boxplot.pdf
"""

from pathlib import Path

# ---------- USER CONFIG ----------
INPUT_CSV_NAME = "Fig6B_data.csv"
OUTPUT_SUMMARY_NAME = "downsampled_p_d_summary.csv"
OUTPUT_FIG_PDF_NAME = "Fig6B_downsampled_notched_boxplot.pdf"
OUTPUT_FIG_JPG_NAME = "Fig6B_downsampled_notched_boxplot.jpg"
OUTPUT_P_PLOT_NAME = "downsampled_p_boxplot.pdf"
OUTPUT_D_PLOT_NAME = "downsampled_cohensd_boxplot.pdf"

EXPOSURES = [80, 160]
TREATMENT_ORDER = ["CTRL", "KA"]
AGE_ORDER = ["young", "earlyOld"]
AGE_LABELS = {
    "young": "young",
    "earlyOld": "early old",
}

N_ITER = 200
N_PER_GROUP = 100
RANDOM_SEED = 98765
JPG_DPI = 200
PDF_DPI = 300
# ---------------------------------

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from matplotlib.patches import Rectangle
from scipy.stats import mannwhitneyu


SCRIPT_DIR = Path(__file__).resolve().parent
INPUT_CSV = SCRIPT_DIR / INPUT_CSV_NAME
OUTPUT_SUMMARY_CSV = SCRIPT_DIR / OUTPUT_SUMMARY_NAME
OUTPUT_FIG_PDF = SCRIPT_DIR / OUTPUT_FIG_PDF_NAME
OUTPUT_FIG_JPG = SCRIPT_DIR / OUTPUT_FIG_JPG_NAME
OUTPUT_P_PLOT = SCRIPT_DIR / OUTPUT_P_PLOT_NAME
OUTPUT_D_PLOT = SCRIPT_DIR / OUTPUT_D_PLOT_NAME


# Figure style
PANEL_BG = "#ffffff"
STRIP_BG = "#d9d9d9"
EDGE = "#3f3f3f"
CTRL_FILL = "#bdd7b3"
KA_FILL = "#c9bdd4"
TXT = "#000000"

PLOT_POSITIONS = {
    (80, "CTRL"): 1.0,
    (80, "KA"): 2.0,
    (160, "CTRL"): 4.0,
    (160, "KA"): 5.0,
}

TREATMENT_FILL = {
    "CTRL": CTRL_FILL,
    "KA": KA_FILL,
}

ANNOTATION_X = {
    80: 1.5,
    160: 4.5,
}


def validate_input_file(path: Path) -> None:
    if not path.exists():
        raise FileNotFoundError(
            f"Input CSV not found: {path}\n"
            f"Place '{INPUT_CSV_NAME}' in the same folder as this script."
        )
    if not path.is_file():
        raise FileNotFoundError(f"Input path exists but is not a file: {path}")


def validate_required_columns(df: pd.DataFrame) -> None:
    required_columns = {"age_grp", "exposure", "treatment", "Median_cfos"}
    missing_columns = required_columns - set(df.columns)
    if missing_columns:
        raise ValueError(f"Missing required columns: {sorted(missing_columns)}")


def cohens_d_two_sample(x, y):
    """Cohen's d for two independent samples (x - y)."""
    x = np.asarray(x)
    y = np.asarray(y)

    n_x = len(x)
    n_y = len(y)
    if n_x < 2 or n_y < 2:
        return np.nan

    mean_x = x.mean()
    mean_y = y.mean()
    std_x = x.std(ddof=1)
    std_y = y.std(ddof=1)

    pooled_sd = np.sqrt(((n_x - 1) * std_x**2 + (n_y - 1) * std_y**2) / (n_x + n_y - 2))
    if not np.isfinite(pooled_sd) or pooled_sd <= 0:
        return np.nan

    return (mean_x - mean_y) / pooled_sd


def values_for_group(df: pd.DataFrame, exposure: int, treatment: str) -> np.ndarray:
    mask = (df["exposure"] == exposure) & (df["treatment"] == treatment)
    return df.loc[mask, "Median_cfos"].dropna().to_numpy()


def stratum_has_both_treatments(df: pd.DataFrame) -> bool:
    treatments_present = set(df["treatment"].dropna().unique())
    return set(TREATMENT_ORDER).issubset(treatments_present)


def summarize_downsampling(subsample_df: pd.DataFrame) -> pd.DataFrame:
    summary = (
        subsample_df.groupby(["age_grp", "exposure"], as_index=False)
        .agg(
            n_subsamples=("iter", "nunique"),
            n_used_per_group=("n_used_per_group", "first"),
            n_ctrl_all=("n_ctrl_all", "first"),
            n_ka_all=("n_ka_all", "first"),
            p_min=("p_value", "min"),
            p_25=("p_value", lambda x: np.quantile(x, 0.25)),
            p_median=("p_value", "median"),
            p_75=("p_value", lambda x: np.quantile(x, 0.75)),
            p_max=("p_value", "max"),
            d_min=("cohens_d", "min"),
            d_25=("cohens_d", lambda x: np.quantile(x, 0.25)),
            d_median=("cohens_d", "median"),
            d_75=("cohens_d", lambda x: np.quantile(x, 0.75)),
            d_max=("cohens_d", "max"),
        )
        .sort_values(["age_grp", "exposure"])
        .reset_index(drop=True)
    )
    return summary


def run_downsampling(df: pd.DataFrame) -> tuple[pd.DataFrame, pd.DataFrame]:
    """Return per-iteration stats and the per-stratum summary table."""
    rng = np.random.default_rng(RANDOM_SEED)
    records = []

    for age_group in AGE_ORDER:
        age_group_df = df[df["age_grp"] == age_group]
        if age_group_df.empty:
            continue

        for exposure in EXPOSURES:
            stratum_df = age_group_df[age_group_df["exposure"] == exposure]
            if stratum_df.empty:
                continue
            if not stratum_has_both_treatments(stratum_df):
                continue

            ctrl_values = values_for_group(stratum_df, exposure, "CTRL")
            ka_values = values_for_group(stratum_df, exposure, "KA")

            n_ctrl_all = len(ctrl_values)
            n_ka_all = len(ka_values)
            sample_size = min(N_PER_GROUP, n_ctrl_all, n_ka_all)
            if sample_size < 10:
                continue

            for iteration in range(N_ITER):
                ctrl_sample = ctrl_values[rng.choice(n_ctrl_all, size=sample_size, replace=False)]
                ka_sample = ka_values[rng.choice(n_ka_all, size=sample_size, replace=False)]

                _, p_value = mannwhitneyu(ctrl_sample, ka_sample, alternative="two-sided")
                cohens_d = cohens_d_two_sample(ka_sample, ctrl_sample)  # KA - CTRL

                records.append(
                    {
                        "age_grp": age_group,
                        "exposure": exposure,
                        "iter": iteration,
                        "p_value": float(p_value),
                        "cohens_d": float(cohens_d),
                        "n_used_per_group": int(sample_size),
                        "n_ctrl_all": int(n_ctrl_all),
                        "n_ka_all": int(n_ka_all),
                    }
                )

    if not records:
        raise RuntimeError("No strata produced subsamples. Check the input data and EXPOSURES.")

    subsample_df = pd.DataFrame(records)
    summary_df = summarize_downsampling(subsample_df)
    return subsample_df, summary_df


def format_p_value(p: float) -> str:
    if not np.isfinite(p):
        return "p=NA"
    if p < 2e-16:
        return "<2e−16"
    if p < 1e-3:
        return f"p={p:.1e}"
    return f"p={p:.3f}"


def format_d_value(d: float) -> str:
    if not np.isfinite(d):
        return "Cohen’s d = NA"
    if abs(d) >= 1:
        return f"Cohen’s d= {d:.2f}"
    return f"Cohen’s d= {d:.1f}"


def make_summary_lookup(summary_df: pd.DataFrame) -> dict:
    return {
        (row.age_grp, int(row.exposure)): row
        for row in summary_df.itertuples(index=False)
    }


def build_boxplot_arrays(df_panel: pd.DataFrame):
    arrays = []
    positions = []
    fills = []

    for exposure in EXPOSURES:
        for treatment in TREATMENT_ORDER:
            values = values_for_group(df_panel, exposure, treatment)
            arrays.append(values)
            positions.append(PLOT_POSITIONS[(exposure, treatment)])
            fills.append(TREATMENT_FILL[treatment])

    return arrays, positions, fills


def style_panel_axes(ax, show_ylabel: bool) -> None:
    ax.set_xlim(0.0, 7.0)
    ax.set_ylim(-0.02, 0.515)
    ax.set_xticks([1.5, 4.5])
    ax.set_xticklabels([str(exposure) for exposure in EXPOSURES], fontsize=16)
    ax.set_yticks(np.arange(0.0, 0.51, 0.1))
    ax.tick_params(axis="y", labelsize=15, width=1.6, length=6, colors=EDGE)
    ax.tick_params(axis="x", width=1.6, length=6, colors=EDGE)

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

    ax.set_xlabel("exposure (mins)", fontsize=18)
    if show_ylabel:
        ax.set_ylabel("nuclear cFOS median intensity (A.U.)", fontsize=18)
    else:
        ax.set_ylabel("")


def draw_median_bars(ax, df_panel: pd.DataFrame) -> None:
    for exposure in EXPOSURES:
        for treatment in TREATMENT_ORDER:
            values = values_for_group(df_panel, exposure, treatment)
            if len(values) == 0:
                continue

            x_position = PLOT_POSITIONS[(exposure, treatment)]
            median_value = float(np.median(values))
            ax.plot(
                [x_position - 0.22, x_position + 0.22],
                [median_value, median_value],
                color=EDGE,
                linewidth=4.0,
                solid_capstyle="butt",
                zorder=4,
            )


def add_age_strip(ax, age_key: str) -> None:
    strip_x0 = 6.05
    strip_width = 0.95
    ax.add_patch(
        Rectangle(
            (strip_x0, -0.03),
            strip_width,
            0.57,
            facecolor=STRIP_BG,
            edgecolor=EDGE,
            linewidth=1.5,
            zorder=0,
        )
    )
    ax.text(
        strip_x0 + strip_width / 2,
        0.245,
        AGE_LABELS.get(age_key, str(age_key)),
        rotation=270,
        ha="center",
        va="center",
        fontsize=16,
        color=TXT,
    )


def add_summary_annotations(ax, age_key: str, summary_lookup: dict) -> None:
    p_y = 0.488
    d_y = 0.462

    for exposure in EXPOSURES:
        row = summary_lookup.get((age_key, exposure))
        if row is None:
            continue

        x_center = ANNOTATION_X[exposure]
        ax.text(
            x_center,
            p_y,
            format_p_value(float(row.p_median)),
            ha="center",
            va="bottom",
            fontsize=12,
            fontweight="bold",
            color=TXT,
        )
        ax.text(
            x_center,
            d_y,
            format_d_value(float(row.d_median)),
            ha="center",
            va="bottom",
            fontsize=10,
            fontweight="bold",
            color=TXT,
        )


def draw_panel(
    ax,
    df_panel: pd.DataFrame,
    age_key: str,
    summary_lookup: dict,
    show_ylabel: bool = False,
) -> None:
    ax.set_facecolor(PANEL_BG)
    add_age_strip(ax, age_key)

    boxplot_arrays, boxplot_positions, boxplot_fills = build_boxplot_arrays(df_panel)
    boxplot = ax.boxplot(
        boxplot_arrays,
        positions=boxplot_positions,
        widths=0.9,
        notch=True,
        patch_artist=True,
        showfliers=False,
        showcaps=False,
        whis=1.5,
        medianprops=dict(color=EDGE, linewidth=1.8),
        whiskerprops=dict(color=EDGE, linewidth=1.5),
        boxprops=dict(color=EDGE, linewidth=1.5),
    )

    for patch, fill in zip(boxplot["boxes"], boxplot_fills):
        patch.set_facecolor(fill)
        patch.set_alpha(1.0)
        patch.set_edgecolor(EDGE)
        patch.set_linewidth(1.5)

    draw_median_bars(ax, df_panel)
    style_panel_axes(ax, show_ylabel=show_ylabel)
    add_summary_annotations(ax, age_key, summary_lookup)


def draw_figure(df: pd.DataFrame, summary_df: pd.DataFrame) -> None:
    summary_lookup = make_summary_lookup(summary_df)

    fig, axes = plt.subplots(1, 2, figsize=(9.29, 6.37), sharey=False)
    plt.subplots_adjust(left=0.11, right=0.995, bottom=0.13, top=0.985, wspace=0.24)

    for ax, age_key in zip(axes, AGE_ORDER):
        panel_df = df[df["age_grp"] == age_key].copy()
        show_ylabel = ax is axes[0]
        draw_panel(ax, panel_df, age_key, summary_lookup, show_ylabel=show_ylabel)

    fig.savefig(OUTPUT_FIG_PDF, dpi=PDF_DPI)
    fig.savefig(OUTPUT_FIG_JPG, dpi=JPG_DPI)
    plt.close(fig)


def build_downsampled_plot_data(subsample_df: pd.DataFrame, summary_df: pd.DataFrame):
    ordered_pairs = [(row.age_grp, int(row.exposure)) for row in summary_df.itertuples(index=False)]
    ordered_labels = [f"{AGE_LABELS.get(age_grp, age_grp)}\n{exposure}" for age_grp, exposure in ordered_pairs]

    p_value_data = []
    cohens_d_data = []

    for age_group, exposure in ordered_pairs:
        stratum_mask = (subsample_df["age_grp"] == age_group) & (subsample_df["exposure"] == exposure)
        p_value_data.append(subsample_df.loc[stratum_mask, "p_value"].dropna().to_numpy())
        cohens_d_data.append(subsample_df.loc[stratum_mask, "cohens_d"].dropna().to_numpy())

    return ordered_labels, p_value_data, cohens_d_data


def draw_p_value_distribution_plot(order_labels, p_value_data, summary_df: pd.DataFrame) -> None:
    fig, ax = plt.subplots(figsize=(6.3, 5.2))
    ax.boxplot(
        p_value_data,
        tick_labels=order_labels,
        showfliers=True,
        notch=False,
        patch_artist=True,
        medianprops=dict(color=EDGE, linewidth=1.5),
        whiskerprops=dict(color=EDGE, linewidth=1.2),
        boxprops=dict(color=EDGE, linewidth=1.2, facecolor="#dddddd"),
        capprops=dict(color=EDGE, linewidth=1.2),
    )
    ax.set_yscale("log")
    ax.set_ylabel("Mann–Whitney p-value (log scale)")
    ax.set_xlabel("age group and exposure")
    ax.set_title(f"Downsampled p-value distributions\n({N_ITER} subsamples, {N_PER_GROUP} cells/group)")
    ax.tick_params(axis="x", labelsize=10)
    ax.tick_params(axis="y", labelsize=10)

    for index, row in enumerate(summary_df.itertuples(index=False), start=1):
        if np.isfinite(row.p_median) and row.p_median > 0:
            ax.text(index, row.p_median, f"median={row.p_median:.2g}", ha="center", va="bottom", fontsize=8)

    fig.tight_layout()
    fig.savefig(OUTPUT_P_PLOT, dpi=PDF_DPI)
    plt.close(fig)


def draw_cohens_d_distribution_plot(order_labels, cohens_d_data, summary_df: pd.DataFrame) -> None:
    fig, ax = plt.subplots(figsize=(6.3, 5.2))
    ax.boxplot(
        cohens_d_data,
        tick_labels=order_labels,
        showfliers=True,
        notch=False,
        patch_artist=True,
        medianprops=dict(color=EDGE, linewidth=1.5),
        whiskerprops=dict(color=EDGE, linewidth=1.2),
        boxprops=dict(color=EDGE, linewidth=1.2, facecolor="#dddddd"),
        capprops=dict(color=EDGE, linewidth=1.2),
    )
    ax.axhline(0, linewidth=0.9, linestyle="--", color=EDGE)
    ax.set_ylabel("Cohen's d (KA - CTRL)")
    ax.set_xlabel("age group and exposure")
    ax.set_title(f"Downsampled Cohen's d distributions\n({N_ITER} subsamples, {N_PER_GROUP} cells/group)")
    ax.tick_params(axis="x", labelsize=10)
    ax.tick_params(axis="y", labelsize=10)

    for index, row in enumerate(summary_df.itertuples(index=False), start=1):
        if np.isfinite(row.d_median):
            ax.text(index, row.d_median, f"median={row.d_median:.2f}", ha="center", va="bottom", fontsize=8)

    fig.tight_layout()
    fig.savefig(OUTPUT_D_PLOT, dpi=PDF_DPI)
    plt.close(fig)


def draw_downsampled_distribution_boxplots(subsample_df: pd.DataFrame, summary_df: pd.DataFrame) -> None:
    order_labels, p_value_data, cohens_d_data = build_downsampled_plot_data(subsample_df, summary_df)
    draw_p_value_distribution_plot(order_labels, p_value_data, summary_df)
    draw_cohens_d_distribution_plot(order_labels, cohens_d_data, summary_df)


def main() -> None:
    validate_input_file(INPUT_CSV)
    df = pd.read_csv(INPUT_CSV)
    validate_required_columns(df)

    subsample_df, summary_df = run_downsampling(df)
    summary_df.to_csv(OUTPUT_SUMMARY_CSV, index=False)

    draw_figure(df, summary_df)
    draw_downsampled_distribution_boxplots(subsample_df, summary_df)

    print(f"Saved summary table:      {OUTPUT_SUMMARY_CSV}")
    print(f"Saved main figure PDF:    {OUTPUT_FIG_PDF}")
    print(f"Saved main figure JPG:    {OUTPUT_FIG_JPG}")
    print(f"Saved p-value PDF:        {OUTPUT_P_PLOT}")
    print(f"Saved Cohen's d PDF:      {OUTPUT_D_PLOT}")


if __name__ == "__main__":
    main()
