import os
import pickle
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt

# =========================
# small generic helpers
# =========================

CLASS_LABELS = {0: "NC", 1: "AE"}
PANEL_LABELS = list("ABCDEFGHIJKLMNOPQRSTUVWXYZ")

def _iter_groups(df):
    """Iterate over (scale, dataset), grouped dataframe."""
    yield from df.groupby(["scale", "dataset"])


def _make_points_df(g, scale, dataset):
    """Return raw plotting points for one (scale, dataset) group."""
    out = g[["X", "Y"]].copy()
    out["scale"] = scale
    out["dataset"] = dataset
    out["class_label"] = out["Y"].map(CLASS_LABELS)
    return out[["scale", "dataset", "X", "Y", "class_label"]]


def _make_fit_df(x, y_fit, scale, dataset):
    """Return fitted curve dataframe for one (scale, dataset) group."""
    return pd.DataFrame({
        "scale": scale,
        "dataset": dataset,
        "X": x,
        "y_fit": y_fit,
    })


def _panel_labels(keys):
    """Map keys to panel labels A, B, C, ..."""
    return dict(zip(keys, PANEL_LABELS[:len(keys)]))


def _savefig(results_path, filename):
    if results_path is not None:
        plt.savefig(os.path.join(results_path, filename), bbox_inches="tight")


def _dedup_legend(ax):
    handles, labels = ax.get_legend_handles_labels()
    by_label = dict(zip(labels, handles))
    return by_label.values(), by_label.keys()


def _confidence_str(confidence):
    return f"{confidence:.0f}" if float(confidence).is_integer() else f"{confidence:.1f}"


def _default_basic_colors():
    return {"NC": "blue", "AE": "orange"}


def _default_ci_colors():
    return {
        "data_NC": "blue",
        "data_AE": "orange",
        "fit": "black",
        "normal": "red",
        "delta": "green",
        "nonparam_boots": "blue",
        "nonparam_stratified_boots": "purple",
        "parametric_boots": "cyan",
    }


def _default_ci_linestyles():
    return {
        "fit": "-",
        "normal": "--",
        "delta": "-.",
        "nonparam_boots": ":",
        "nonparam_stratified_boots": (0, (3, 1, 1, 1)),
        "parametric_boots": (0, (5, 2)),
    }


# =========================
# builders
# =========================

def build_plot_data_basic(df_data, df_fit_index, lg, n_grid=100, verbose=False):
    """
    Build raw points and fitted curves.
    """
    points, fits = [], []

    for (scale, dataset), g in _iter_groups(df_data):
        if verbose:
            print(f"Calculating: {scale}, {dataset}")

        x = g["X"].to_numpy()
        pars = df_fit_index.loc[(scale, dataset), "pars"]

        x_fit = np.linspace(x.min(), x.max(), n_grid)
        y_fit = lg.model(x_fit, pars)

        points.append(_make_points_df(g, scale, dataset))
        fits.append(_make_fit_df(x_fit, y_fit, scale, dataset))

    return pd.concat(points, ignore_index=True), pd.concat(fits, ignore_index=True)


def build_cost_perturbation_data(
    df_data,
    df_fit,
    lg,
    scale="log",
    perturb_min=-200,
    perturb_max=200,
    n_grid=400,
    verbose=False,
):
    """
    Build cost profiles under one-parameter perturbations.
    """
    rows = []
    dpar_grid = np.linspace(perturb_min, perturb_max, n_grid)

    datasets = df_data.loc[df_data["scale"] == scale, "dataset"].drop_duplicates()
    for dataset in datasets:
        if verbose:
            print(f"Calculating: {scale}, {dataset}")

        g = df_data[(df_data["scale"] == scale) & (df_data["dataset"] == dataset)]
        x, y = g["X"].to_numpy(), g["Y"].to_numpy()

        pars = np.array(
            df_fit.loc[(df_fit["scale"] == scale) & (df_fit["dataset"] == dataset), "pars"].iloc[0],
            copy=True,
        )

        for j in range(len(pars)):
            if verbose:
                print(f"\tparameter {j}")

            cost = np.empty_like(dpar_grid, dtype=float)
            for i, dp in enumerate(dpar_grid):
                pars_tmp = pars.copy()
                pars_tmp[j] += dp
                cost[i] = lg.get_cost(x, y, pars_tmp, jac=False)

            rows.append(pd.DataFrame({
                "scale": scale,
                "dataset": dataset,
                "parameter": j,
                "dpar": dpar_grid,
                "cost": cost,
            }))

    return pd.concat(rows, ignore_index=True)


def build_plot_data(df_data, df_fit_index, alpha, lg, n_grid=100, verbose=False):
    """
    Build raw points, fitted curves, and analytic confidence intervals.
    """
    probs = [alpha / 2, 1 - alpha / 2]
    confidence = 100 * (1 - alpha)

    points, fits, cis = [], [], []

    for (scale, dataset), g in _iter_groups(df_data):
        if verbose:
            print(f"Calculating: {scale}, {dataset}")

        x = g["X"].to_numpy()
        pars = df_fit_index.loc[(scale, dataset), "pars"]
        cov = df_fit_index.loc[(scale, dataset), "cov"]

        x_fit = np.linspace(x.min(), x.max(), n_grid)
        y_fit = lg.model(x_fit, pars)

        points.append(_make_points_df(g, scale, dataset))
        fits.append(_make_fit_df(x_fit, y_fit, scale, dataset))

        for method in ("normal", "delta"):
            qfun = getattr(lg, f"get_model_quantiles_{method}")
            y_low, y_high = qfun(x_fit, probs, pars, cov)

            cis.append(pd.DataFrame({
                "scale": scale,
                "dataset": dataset,
                "method": method,
                "confidence": confidence,
                "X": x_fit,
                "y_low": y_low,
                "y_high": y_high,
            }))

    return (
        pd.concat(points, ignore_index=True),
        pd.concat(fits, ignore_index=True),
        pd.concat(cis, ignore_index=True),
    )


def load_or_build_bootstrap_pars(
    df_data,
    lg,
    results_path,
    filename="boots_pars_results.pkl",
    n_boots=10000,
    methods=("normal", "nonparam_boots", "nonparam_stratified_boots", "parametric_boots"),
    verbose=False,
):
    """
    Load bootstrap parameter draws from disk or build them.
    """
    result_file = os.path.join(results_path, filename)

    if os.path.isfile(result_file):
        if verbose:
            print("Loading existing file.")
        with open(result_file, "rb") as f:
            return pickle.load(f)

    out = {}
    for (scale, dataset), g in _iter_groups(df_data):
        if verbose:
            print(f"Bootstrap parameters: {scale}, {dataset}")

        x, y = g["X"].to_numpy(), g["Y"].to_numpy()

        for method in methods:
            if verbose:
                print(f"\tmethod: {method}")

            par_fun = getattr(lg, f"get_{method}_pars")
            try:
                out[(scale, dataset, method)] = par_fun(x, y, m=n_boots)
            except Exception:
                if verbose:
                    print("\tfailed")
                out[(scale, dataset, method)] = None

    with open(result_file, "wb") as f:
        pickle.dump(out, f)

    return out


def build_plot_data_full_ci(
    df_data,
    df_fit_index,
    lg,
    alpha,
    boots_pars_results,
    n_grid=100,
    analytic_methods=("normal", "delta"),
    bootstrap_methods=("nonparam_boots", "parametric_boots"),
    verbose=False,
):
    """
    Build raw points, fitted curves, analytic CI, and bootstrap CI.
    """
    probs = [alpha / 2, 1 - alpha / 2]
    confidence = 100 * (1 - alpha)

    points, fits, cis = [], [], []

    for (scale, dataset), g in _iter_groups(df_data):
        if verbose:
            print(f"Calculating plot data: {scale}, {dataset}")

        x = g["X"].to_numpy()
        pars = df_fit_index.loc[(scale, dataset), "pars"]
        cov = df_fit_index.loc[(scale, dataset), "cov"]

        x_fit = np.linspace(x.min(), x.max(), n_grid)
        y_fit = lg.model(x_fit, pars)

        points.append(_make_points_df(g, scale, dataset))
        fits.append(_make_fit_df(x_fit, y_fit, scale, dataset))

        for method in analytic_methods:
            if verbose:
                print(f"\tanalytic CI: {method}")

            qfun = getattr(lg, f"get_model_quantiles_{method}")
            y_low, y_high = qfun(x_fit, probs, pars, cov)

            cis.append(pd.DataFrame({
                "scale": scale,
                "dataset": dataset,
                "method": method,
                "ci_source": "analytic",
                "confidence": confidence,
                "X": x_fit,
                "y_low": y_low,
                "y_high": y_high,
            }))

        for method in bootstrap_methods:
            if verbose:
                print(f"\tbootstrap CI: {method}")

            bpars = boots_pars_results[(scale, dataset, method)]
            if bpars is None:
                continue

            quant = np.quantile([lg.model(x_fit, p) for p in bpars], probs, axis=0)
            y_low, y_high = quant

            cis.append(pd.DataFrame({
                "scale": scale,
                "dataset": dataset,
                "method": method,
                "ci_source": "bootstrap",
                "confidence": confidence,
                "X": x_fit,
                "y_low": y_low,
                "y_high": y_high,
            }))

    return (
        pd.concat(points, ignore_index=True),
        pd.concat(fits, ignore_index=True),
        pd.concat(cis, ignore_index=True),
    )


# =========================
# plotting
# =========================
def plot_from_dataframes_basic(
    df_points,
    df_fit_plot,
    xlabs,
    scales=None,
    results_path=None,
    filename="logit_fit_paper.pdf",
    verbose=False,
    width=5,
    height=4,
    class_color_map=None,
    fit_color_map=None,
    point_size_full=28,
    point_size_trim=42,
    alpha_full=0.22,
    alpha_trim=0.85,
    jitter_y=0.0,
    jitter_y_full=None,
    jitter_y_trim=None,
    random_state=123,
    plot_labs = None,
):
    """
    Plot raw data and fitted curves. One panel per scale.

    Visual encoding:
    - class (NC / AE) -> color
    - FULL vs TRIM    -> background vs highlighted overlay
    - optional horizontal jitter to reduce overlap
    """
    if scales is None:
        scales = list(df_points["scale"].drop_duplicates())

    datasets = list(df_points["dataset"].drop_duplicates())

    if class_color_map is None:
        class_color_map = {
            "NC": "tab:blue",
            "AE": "tab:orange",
        }

    if fit_color_map is None:
        fit_color_map = {dataset: None for dataset in datasets}

    if jitter_y_full is None:
        jitter_y_full = jitter_y
    if jitter_y_trim is None:
        jitter_y_trim = jitter_y

    rng = np.random.default_rng(random_state)

    fig, axs = plt.subplots(
        ncols=len(scales),
        figsize=(width * len(scales), height),
        layout="constrained",
        squeeze=False,
    )
    axs = axs[0]
    
    if plot_labs is None:
        plot_labs = dict(zip(scales, list("ABCDEFGHIJKLMNOPQRSTUVWXYZ")[:len(scales)]))
    
    if verbose:
        print(plot_labs)

    dataset_order = [d for d in ["FULL", "TRIM", "TRIMMED"] if d in datasets]
    dataset_order += [d for d in datasets if d not in dataset_order]

    for ax, scale in zip(axs, scales):
        ax.set_xlabel(xlabs[scale])
        ax.set_ylabel("P(AE|X = x)")
        ax.text(
            0.05,
            0.9,
            plot_labs[scale],
            fontsize=14,
            transform=ax.transAxes,
        )

        g_points = df_points[df_points["scale"] == scale]

        for dataset in dataset_order:
            g_dataset = g_points[g_points["dataset"] == dataset]

            for lab in ["NC", "AE"]:
                sub = g_dataset[g_dataset["class_label"] == lab]

                x_plot = sub["X"].to_numpy().copy()
                y_plot = sub["Y"].to_numpy()

                jitter_sd = jitter_y_full if dataset == "FULL" else jitter_y_trim
                if jitter_sd > 0:
                    #y_plot = y_plot + rng.normal(0.0, jitter_sd, size=len(y_plot))
                    y_plot = y_plot + (jitter_sd if dataset == "FULL" else -jitter_sd)

                if dataset == "FULL":
                    ax.scatter(
                        x_plot,
                        y_plot,
                        label=f"data: {dataset} {lab}",
                        s=point_size_full,
                        alpha=alpha_full,
                        color=class_color_map[lab],
                        edgecolors="none",
                        zorder=1,
                    )
                else:
                    ax.scatter(
                        x_plot,
                        y_plot,
                        label=f"data: {dataset} {lab}",
                        s=point_size_trim,
                        alpha=alpha_trim,
                        color=class_color_map[lab],
                        edgecolors="black",
                        linewidths=0.6,
                        zorder=3,
                    )

        g_fit = df_fit_plot[df_fit_plot["scale"] == scale]

        for dataset in dataset_order:
            sub = g_fit[g_fit["dataset"] == dataset].sort_values("X")
            ax.plot(
                sub["X"],
                sub["y_fit"],
                label=f"fit {dataset}",
                color=fit_color_map.get(dataset),
                linewidth=2,
                zorder=2 if dataset == "FULL" else 4,
            )

    for ax in axs:
        ax.label_outer()

    handles, labels = axs[0].get_legend_handles_labels()
    by_label = dict(zip(labels, handles))
    fig.legend(
        by_label.values(),
        by_label.keys(),
        bbox_to_anchor=(0.975, 0.4),
        loc="center right",
    )

    if results_path is not None:
        plt.savefig(
            os.path.join(results_path, filename),
            bbox_inches="tight",
        )

    plt.show()

def plot_cost_perturbation_from_dataframe(
    df_cost_plot,
    results_path=None,
    filename="logit_cost_min_paper.pdf",
    verbose=False,
    width=5,
    height=4,
    plot_labs=None,
    yscale="log",
):
    """
    Plot cost perturbation curves. One panel per dataset.
    """
    datasets = list(df_cost_plot["dataset"].drop_duplicates())
    plot_labs = _panel_labels(datasets) if plot_labs is None else plot_labs

    if verbose:
        print(plot_labs)

    fig, axs = plt.subplots(
        ncols=len(datasets),
        nrows=1,
        figsize=(width * len(datasets), height),
        layout="constrained",
        sharey=True,
        squeeze=False,
    )
    axs = axs[0]

    for ax, dataset in zip(axs, datasets):
        if verbose:
            print(f"Plotting: {dataset}")

        ax.text(0.05, 0.05, plot_labs[dataset], fontsize=14, transform=ax.transAxes)
        sub_dataset = df_cost_plot[df_cost_plot["dataset"] == dataset]

        for parameter, sub in sub_dataset.groupby("parameter", sort=True):
            sub = sub.sort_values("dpar")
            ax.plot(sub["dpar"], sub["cost"], label=rf"$\alpha_{{{parameter}}}$")

        ax.set_yscale(yscale)
        ax.set_xlabel(r"parameter ($\alpha_i$) perturbation")
        ax.set_ylabel("cost")
        ax.label_outer()

    handles, labels = _dedup_legend(axs[0])
    fig.legend(handles, labels, bbox_to_anchor=(0.975, 0.3), loc="center right")

    _savefig(results_path, filename)
    plt.show()


def plot_from_dataframes(
    df_points,
    df_fit_plot,
    df_ci_plot,
    xlabs,
    scales=None,
    datasets=None,
    results_path=None,
    filename="logit_fit_CI_simple_paper.pdf",
    verbose=False,
    color_map=None,
    ls_map=None,
    ci_alpha=0.15,
    ci_linewidth=1.5,
    fit_linewidth=2.0,
    width=5,
    height=3.5,
):
    """
    Plot raw data, fitted curve, and confidence intervals.
    """
    scales = sorted(df_points["scale"].unique()) if scales is None else list(scales)
    datasets = sorted(df_points["dataset"].unique()) if datasets is None else list(datasets)

    color_map = _default_ci_colors() if color_map is None else color_map
    ls_map = _default_ci_linestyles() if ls_map is None else ls_map

    fig, axs = plt.subplots(
        ncols=len(scales),
        nrows=len(datasets),
        figsize=(width * len(scales), height * len(datasets)),
        layout="constrained",
        squeeze=False,
    )

    panel_keys = [(scale, dataset) for dataset in datasets for scale in scales]
    plot_labs = _panel_labels(panel_keys)

    if verbose:
        print(plot_labs)

    for i, dataset in enumerate(datasets):
        for j, scale in enumerate(scales):
            ax = axs[i, j]
            ax.set_xlabel(xlabs[scale])
            ax.set_ylabel("P(AE|X = x)")
            ax.text(0.05, 0.9, plot_labs[(scale, dataset)], fontsize=14, transform=ax.transAxes)

            g_points = df_points[(df_points["scale"] == scale) & (df_points["dataset"] == dataset)]
            for lab in ("NC", "AE"):
                sub = g_points[g_points["class_label"] == lab]
                ax.scatter(
                    sub["X"], sub["Y"],
                    label=f"data: {lab}",
                    alpha=0.5,
                    color=color_map[f"data_{lab}"],
                )

            g_fit = df_fit_plot[(df_fit_plot["scale"] == scale) & (df_fit_plot["dataset"] == dataset)].sort_values("X")
            ax.plot(
                g_fit["X"], g_fit["y_fit"],
                label="fit",
                color=color_map["fit"],
                linestyle=ls_map["fit"],
                linewidth=fit_linewidth,
            )

            g_ci = df_ci_plot[(df_ci_plot["scale"] == scale) & (df_ci_plot["dataset"] == dataset)]
            for method, sub in g_ci.groupby("method", sort=False):
                sub = sub.sort_values("X")
                conf_str = _confidence_str(sub["confidence"].iloc[0])

                ax.fill_between(
                    sub["X"], sub["y_low"], sub["y_high"],
                    color=color_map[method],
                    alpha=ci_alpha,
                    linewidth=0,
                )
                ax.plot(
                    sub["X"], sub["y_low"],
                    color=color_map[method],
                    linestyle=ls_map[method],
                    linewidth=ci_linewidth,
                    label=f"CI: {method} {conf_str}%",
                )
                ax.plot(
                    sub["X"], sub["y_high"],
                    color=color_map[method],
                    linestyle=ls_map[method],
                    linewidth=ci_linewidth,
                )

            ax.label_outer()

    handles, labels = _dedup_legend(axs[0, 0])
    fig.legend(
        handles, labels,
        bbox_to_anchor=(0.975, 0.2),
        framealpha=0.5,
        loc="center right",
    )

    _savefig(results_path, filename)
    plt.show()


# alias, if you want the previous name to remain available
plot_from_dataframes_full_ci = plot_from_dataframes