from __future__ import annotations

import argparse
from pathlib import Path
from typing import Iterable

try:
    import matplotlib

    matplotlib.use("Agg")
    import matplotlib.pyplot as plt
except ModuleNotFoundError as exc:
    raise ModuleNotFoundError(
        "plot_burden_index_trajectories.py requires matplotlib. "
        "Install it in the server environment before running this script."
    ) from exc
import numpy as np
import pandas as pd


REQUIRED_COLUMNS = {
    "patient_id",
    "sex",
    "landmark_age",
    "burden_type",
    "burden_dimension",
    "bi_historical",
    "bi_future",
    "bi_total",
}


def _sex_label(value: object) -> str:
    text = str(value).strip().lower()
    if text in {"0", "0.0", "female", "f", "woman"}:
        return "female"
    if text in {"1", "1.0", "male", "m", "man"}:
        return "male"
    return text or "unknown"


def _load_bi_csv(path: Path) -> pd.DataFrame:
    header = pd.read_csv(path, nrows=0)
    missing = sorted(REQUIRED_COLUMNS - set(header.columns))
    if missing:
        raise ValueError(f"{path} is missing required columns: {missing}")

    df = pd.read_csv(
        path,
        usecols=[
            "patient_id",
            "sex",
            "landmark_age",
            "burden_type",
            "burden_dimension",
            "bi_historical",
            "bi_future",
            "bi_total",
        ],
    )
    df["sex_label"] = df["sex"].map(_sex_label)
    df["landmark_age"] = pd.to_numeric(df["landmark_age"], errors="raise")
    for col in ["bi_historical", "bi_future", "bi_total"]:
        df[col] = pd.to_numeric(df[col], errors="raise")
    return df


def _sample_patients(
    df: pd.DataFrame,
    *,
    n_per_sex: int,
    seed: int,
) -> dict[str, np.ndarray]:
    rng = np.random.default_rng(int(seed))
    samples: dict[str, np.ndarray] = {}
    for sex_label in ["female", "male"]:
        ids = np.asarray(sorted(df.loc[df["sex_label"] == sex_label, "patient_id"].unique()))
        if ids.size == 0:
            samples[sex_label] = np.asarray([], dtype=np.int64)
            continue
        take = min(int(n_per_sex), int(ids.size))
        samples[sex_label] = np.asarray(rng.choice(ids, size=take, replace=False))
    return samples


def _aggregate_total(df: pd.DataFrame) -> pd.DataFrame:
    return (
        df.groupby(["patient_id", "sex_label", "burden_type", "landmark_age"], as_index=False)
        [["bi_historical", "bi_future", "bi_total"]]
        .sum()
    )


def _plot_selected_trajectories(
    total_df: pd.DataFrame,
    *,
    burden_type: str,
    selected: dict[str, np.ndarray],
    output_dir: Path,
) -> None:
    sub = total_df[total_df["burden_type"] == burden_type].copy()
    mean_df = (
        sub.groupby(["sex_label", "landmark_age"], as_index=False)["bi_total"]
        .mean()
        .sort_values("landmark_age")
    )

    fig, ax = plt.subplots(figsize=(9.5, 5.5), dpi=160)
    colors = {"female": "#b83280", "male": "#2563eb"}
    for sex_label in ["female", "male"]:
        patient_ids = selected.get(sex_label, np.asarray([], dtype=np.int64))
        for pid in patient_ids:
            one = sub[sub["patient_id"] == pid].sort_values("landmark_age")
            if one.empty:
                continue
            ax.plot(
                one["landmark_age"],
                one["bi_total"],
                color=colors.get(sex_label, "0.4"),
                alpha=0.22,
                linewidth=1.0,
            )
        mean_one = mean_df[mean_df["sex_label"] == sex_label]
        if not mean_one.empty:
            ax.plot(
                mean_one["landmark_age"],
                mean_one["bi_total"],
                color=colors.get(sex_label, "0.4"),
                linewidth=2.8,
                label=f"{sex_label} mean",
            )

    ax.set_title(f"{burden_type}: sampled individual trajectories and sex-specific means")
    ax.set_xlabel("Landmark age")
    ax.set_ylabel("Total burden index")
    ax.grid(True, alpha=0.25)
    ax.legend(frameon=False)
    fig.tight_layout()
    fig.savefig(output_dir / f"{burden_type}_sampled_trajectories_by_sex.png")
    plt.close(fig)


def _top_dimensions(df: pd.DataFrame, *, burden_type: str, top_n: int) -> list[str]:
    sub = df[df["burden_type"] == burden_type]
    if sub.empty:
        return []
    ranked = (
        sub.groupby("burden_dimension")["bi_total"]
        .mean()
        .sort_values(ascending=False)
        .head(int(top_n))
    )
    return [str(x) for x in ranked.index.tolist()]


def _plot_dimension_mean_panels(
    df: pd.DataFrame,
    *,
    burden_type: str,
    dimensions: Iterable[str],
    output_dir: Path,
) -> None:
    dims = list(dimensions)
    if not dims:
        return

    mean_df = (
        df[(df["burden_type"] == burden_type) & (df["burden_dimension"].isin(dims))]
        .groupby(["sex_label", "burden_dimension", "landmark_age"], as_index=False)["bi_total"]
        .mean()
        .sort_values("landmark_age")
    )
    n_cols = min(3, len(dims))
    n_rows = int(np.ceil(len(dims) / n_cols))
    fig, axes = plt.subplots(
        n_rows,
        n_cols,
        figsize=(4.2 * n_cols, 3.2 * n_rows),
        dpi=160,
        squeeze=False,
    )
    colors = {"female": "#b83280", "male": "#2563eb"}
    for ax, dim in zip(axes.ravel(), dims):
        panel = mean_df[mean_df["burden_dimension"] == dim]
        for sex_label in ["female", "male"]:
            one = panel[panel["sex_label"] == sex_label]
            if one.empty:
                continue
            ax.plot(
                one["landmark_age"],
                one["bi_total"],
                color=colors.get(sex_label, "0.4"),
                linewidth=2.0,
                label=sex_label,
            )
        ax.set_title(str(dim), fontsize=9)
        ax.set_xlabel("Age")
        ax.set_ylabel("Mean BI")
        ax.grid(True, alpha=0.25)

    for ax in axes.ravel()[len(dims):]:
        ax.axis("off")
    handles, labels = axes.ravel()[0].get_legend_handles_labels()
    if handles:
        fig.legend(handles, labels, loc="upper right", frameon=False)
    fig.suptitle(f"{burden_type}: sex-specific mean trajectories for top dimensions")
    fig.tight_layout(rect=(0, 0, 0.98, 0.95))
    fig.savefig(output_dir / f"{burden_type}_top_dimension_means_by_sex.png")
    plt.close(fig)


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Plot DeepHealth burden-index trajectories by sex."
    )
    parser.add_argument("--input_csv", type=str, required=True)
    parser.add_argument("--output_dir", type=str, default=None)
    parser.add_argument("--n_per_sex", type=int, default=10)
    parser.add_argument("--seed", type=int, default=2026)
    parser.add_argument("--top_dimensions", type=int, default=12)
    args = parser.parse_args()

    input_csv = Path(args.input_csv)
    output_dir = Path(args.output_dir) if args.output_dir else input_csv.with_suffix("")
    output_dir.mkdir(parents=True, exist_ok=True)

    df = _load_bi_csv(input_csv)
    total_df = _aggregate_total(df)
    selected = _sample_patients(total_df, n_per_sex=int(args.n_per_sex), seed=int(args.seed))

    for burden_type in sorted(df["burden_type"].dropna().unique().tolist()):
        burden_type = str(burden_type)
        _plot_selected_trajectories(
            total_df,
            burden_type=burden_type,
            selected=selected,
            output_dir=output_dir,
        )
        dims = _top_dimensions(df, burden_type=burden_type, top_n=int(args.top_dimensions))
        _plot_dimension_mean_panels(
            df,
            burden_type=burden_type,
            dimensions=dims,
            output_dir=output_dir,
        )

    selected_rows = []
    for sex_label, patient_ids in selected.items():
        for pid in patient_ids.tolist():
            selected_rows.append({"sex": sex_label, "patient_id": int(pid)})
    pd.DataFrame(selected_rows).to_csv(output_dir / "sampled_patients_by_sex.csv", index=False)
    print(f"Input: {input_csv}")
    print(f"Output directory: {output_dir}")
    print(f"Sampled patients per sex: {int(args.n_per_sex)}")


if __name__ == "__main__":
    main()