import argparse
from pathlib import Path

import numpy as np
import pandas as pd

import matplotlib as mpl

mpl.use("Agg")

import matplotlib.pyplot as plt
import seaborn as sns


def _coerce_numeric(df: pd.DataFrame, cols: list[str]) -> pd.DataFrame:
    out = df.copy()
    for c in cols:
        if c in out.columns:
            out[c] = pd.to_numeric(out[c], errors="coerce")
    return out


def _coerce_bool(x: pd.Series) -> pd.Series:
    if x.dtype == bool:
        return x
    s = x.astype(str).str.strip().str.lower()
    return s.map({"true": True, "false": False})


def _save(fig: plt.Figure, out_dir: Path, stem: str, write_pdf: bool, write_png: bool, dpi: int) -> None:
    out_dir.mkdir(parents=True, exist_ok=True)
    if write_pdf:
        fig.savefig(out_dir / f"{stem}.pdf", bbox_inches="tight")
    if write_png:
        fig.savefig(out_dir / f"{stem}.png", dpi=dpi, bbox_inches="tight")


def plot_obs_grad(obs: pd.DataFrame, out_dir: Path, write_pdf: bool, write_png: bool, dpi: int) -> None:
    obs = _coerce_numeric(
        obs,
        [
            "q_edge_sample_prob",
            "noise_sigma",
            "mu_disp_max",
            "time_in_e",
            "det_brier",
            "det_log_loss",
        ],
    )

    obs = obs.copy()
    obs["mu_disp_active"] = obs["mu_disp_max"].notna()

    noise_type = obs.get("noise_type", pd.Series(["none"] * len(obs))).fillna("none").astype(str)
    sigma = obs.get("noise_sigma", pd.Series([np.nan] * len(obs)))

    obs["noise_key"] = noise_type
    m = sigma.notna()
    if m.any():
        obs.loc[m, "noise_key"] = noise_type[m] + " sigma=" + sigma[m].map(lambda x: f"{x:g}")

    metrics = [
        ("time_in_e", "time in $\\mathcal{E}$"),
        ("det_brier", "Brier score"),
        ("det_log_loss", "log loss"),
    ]

    grp = ["noise_key", "mu_disp_active", "q_edge_sample_prob"]
    agg = obs.groupby(grp)[[m for m, _ in metrics]].agg(["mean", "std", "count"]).reset_index()
    agg.columns = ["_".join([c for c in col if c]) for col in agg.columns.to_flat_index()]

    noise_keys = sorted([x for x in agg["noise_key"].dropna().unique()])

    sns.set_theme(context="paper", style="whitegrid")

    nrows = max(1, len(noise_keys))
    fig, axes = plt.subplots(nrows=nrows, ncols=len(metrics), figsize=(12, 3.6 * nrows), sharex="col")
    if nrows == 1:
        axes = np.array([axes])

    palette = sns.color_palette("deep", 2)

    for i, noise_key in enumerate(noise_keys):
        sub = agg[agg["noise_key"] == noise_key].copy()
        sub = sub.sort_values("q_edge_sample_prob")

        for j, (metric, ylabel) in enumerate(metrics):
            ax = axes[i, j]
            for k, mu_disp_active in enumerate([False, True]):
                s = sub[sub["mu_disp_active"] == mu_disp_active].copy()
                s = s.sort_values("q_edge_sample_prob")
                if len(s) == 0:
                    continue

                x = s["q_edge_sample_prob"].to_numpy(dtype=float)
                y = s[f"{metric}_mean"].to_numpy(dtype=float)
                yerr = s[f"{metric}_std"].to_numpy(dtype=float)

                label = "meso off" if not mu_disp_active else "meso on"

                ax.errorbar(
                    x,
                    y,
                    yerr=yerr,
                    color=palette[k],
                    marker="o",
                    linewidth=2,
                    capsize=3,
                    label=label,
                )

            ax.set_ylabel(ylabel)
            ax.set_xlabel("edge sample prob $q$")
            if metric == "time_in_e":
                ax.set_ylim(-0.02, 1.02)

        axes[i, 0].set_title(noise_key)

    handles, labels = axes[0, 0].get_legend_handles_labels()
    if handles:
        axes[0, 0].legend(handles, labels, loc="best", frameon=True)

    fig.tight_layout()
    _save(fig, out_dir, "fig_obs_grad", write_pdf=write_pdf, write_png=write_png, dpi=dpi)
    plt.close(fig)


def _bar_panel(
    ax: plt.Axes,
    data: pd.DataFrame,
    order: list[str],
    key_col: str,
    metric: str,
    ylabel: str,
) -> None:
    s = data.groupby(key_col)[metric].agg(["mean", "std", "count"]).reindex(order)
    means = s["mean"].to_numpy(dtype=float)
    stds = s["std"].to_numpy(dtype=float)
    counts = s["count"].to_numpy(dtype=int)

    x = np.arange(len(order))
    bars = ax.bar(x, np.nan_to_num(means, nan=0.0), yerr=np.nan_to_num(stds, nan=0.0), capsize=3)

    for i, b in enumerate(bars):
        if counts[i] == 0 or np.isnan(means[i]):
            ax.text(b.get_x() + b.get_width() / 2, 0.0, "NA", ha="center", va="bottom", fontsize=9)

    ax.set_xticks(x)
    ax.set_xticklabels(order)
    ax.set_ylabel(ylabel)


def plot_top_act(top: pd.DataFrame, out_dir: Path, write_pdf: bool, write_png: bool, dpi: int) -> None:
    top = _coerce_numeric(top, ["time_in_e", "u_mean", "u_saturation_fraction", "top_fraction_steps_active"])

    if "top_act_enabled" in top.columns:
        top = top.copy()
        top["top_act_enabled"] = _coerce_bool(top["top_act_enabled"])

    top["enabled_label"] = top["top_act_enabled"].map({False: "disabled", True: "enabled"})

    sns.set_theme(context="paper", style="whitegrid")

    metrics = [
        ("time_in_e", "time in $\\mathcal{E}$"),
        ("u_mean", "mean heat $u$")
        if "u_mean" in top.columns
        else None,
        ("u_saturation_fraction", "saturation fraction")
        if "u_saturation_fraction" in top.columns
        else None,
        ("top_fraction_steps_active", "topology-active steps")
        if "top_fraction_steps_active" in top.columns
        else None,
    ]
    metrics = [m for m in metrics if m is not None]

    ncols = 2
    nrows = int(np.ceil(len(metrics) / ncols))
    fig, axes = plt.subplots(nrows=nrows, ncols=ncols, figsize=(10, 3.4 * nrows))
    axes = np.array(axes).reshape(nrows, ncols)

    order = ["disabled", "enabled"]
    for i, (metric, ylabel) in enumerate(metrics):
        r, c = divmod(i, ncols)
        ax = axes[r, c]
        _bar_panel(ax, top, order=order, key_col="enabled_label", metric=metric, ylabel=ylabel)

    for j in range(len(metrics), nrows * ncols):
        r, c = divmod(j, ncols)
        axes[r, c].axis("off")

    fig.tight_layout()
    _save(fig, out_dir, "fig_top_act_abl", write_pdf=write_pdf, write_png=write_png, dpi=dpi)
    plt.close(fig)


def plot_stress(stress: pd.DataFrame, out_dir: Path, write_pdf: bool, write_png: bool, dpi: int) -> None:
    stress = _coerce_numeric(
        stress,
        [
            "time_in_e",
            "u_mean",
            "shock_count",
            "rewire_count",
            "shock_recovery_first",
            "rewire_recovery_first",
        ],
    )

    stress = stress.copy()
    stress["has_shock"] = stress["shock_count"].fillna(0) > 0
    stress["has_rewire"] = stress["rewire_count"].fillna(0) > 0

    def _label_row(r: pd.Series) -> str:
        if bool(r["has_shock"]) and bool(r["has_rewire"]):
            return "shock+rewire"
        if bool(r["has_shock"]):
            return "shock"
        if bool(r["has_rewire"]):
            return "rewire"
        return "none"

    stress["stress_label"] = stress.apply(_label_row, axis=1)

    sns.set_theme(context="paper", style="whitegrid")

    metrics = [
        ("time_in_e", "time in $\\mathcal{E}$"),
        ("u_mean", "mean heat $u$") if "u_mean" in stress.columns else None,
        ("rewire_recovery_first", "rewire recovery (steps)")
        if "rewire_recovery_first" in stress.columns
        else None,
        ("shock_recovery_first", "shock recovery (steps)")
        if "shock_recovery_first" in stress.columns
        else None,
    ]
    metrics = [m for m in metrics if m is not None]

    ncols = 2
    nrows = int(np.ceil(len(metrics) / ncols))
    fig, axes = plt.subplots(nrows=nrows, ncols=ncols, figsize=(10, 3.4 * nrows))
    axes = np.array(axes).reshape(nrows, ncols)

    order = ["none", "shock", "rewire", "shock+rewire"]
    for i, (metric, ylabel) in enumerate(metrics):
        r, c = divmod(i, ncols)
        ax = axes[r, c]
        _bar_panel(ax, stress, order=order, key_col="stress_label", metric=metric, ylabel=ylabel)

    for j in range(len(metrics), nrows * ncols):
        r, c = divmod(j, ncols)
        axes[r, c].axis("off")

    fig.tight_layout()
    _save(fig, out_dir, "fig_stress_abl", write_pdf=write_pdf, write_png=write_png, dpi=dpi)
    plt.close(fig)


def main() -> None:
    p = argparse.ArgumentParser()
    p.add_argument("--obs_csv", required=True)
    p.add_argument("--stress_csv", required=True)
    p.add_argument("--top_csv", required=True)
    p.add_argument("--out_dir", default="figures")
    p.add_argument("--no_pdf", action="store_true")
    p.add_argument("--no_png", action="store_true")
    p.add_argument("--dpi", type=int, default=300)
    args = p.parse_args()

    out_dir = Path(args.out_dir)

    obs = pd.read_csv(Path(args.obs_csv))
    stress = pd.read_csv(Path(args.stress_csv))
    top = pd.read_csv(Path(args.top_csv))

    write_pdf = not args.no_pdf
    write_png = not args.no_png

    plot_obs_grad(obs, out_dir=out_dir, write_pdf=write_pdf, write_png=write_png, dpi=args.dpi)
    plot_stress(stress, out_dir=out_dir, write_pdf=write_pdf, write_png=write_png, dpi=args.dpi)
    plot_top_act(top, out_dir=out_dir, write_pdf=write_pdf, write_png=write_png, dpi=args.dpi)


if __name__ == "__main__":
    main()