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

import pickle

def build_plot_data(df_data, df_fit_index, alpha, lg, n_grid=100, verbose=False):
    """
    Compute all quantities needed for plotting and return them as dataframes.

    Parameters
    ----------
    df_data : pd.DataFrame
        Input data with columns including scale, dataset, X, Y.
    df_fit_index : pd.DataFrame
        Fit results indexed by (scale, dataset), with columns 'pars' and 'cov'.
    alpha : float
        Two-sided significance level for confidence intervals.
        Example: alpha = 0.05 gives a 95% CI.
    lg : object
        Model object with methods:
            - model(x, pars)
            - get_model_quantiles_normal(x, probs, pars, cov)
            - get_model_quantiles_delta(x, probs, pars, cov)
    n_grid : int, default 100
        Number of x-grid points for plotting.
    verbose : bool, default False
        Print progress information.

    Returns
    -------
    df_points : pd.DataFrame
        Original observed data, one row per observation.
    df_fit_plot : pd.DataFrame
        Fitted curve on a grid, one row per x-grid value.
    df_ci_plot : pd.DataFrame
        Confidence interval bands on a grid, one row per x-grid value and method.
    """
    probs = [alpha / 2, 1 - alpha / 2]
    confidence = 100 * (1 - alpha)

    points_list = []
    fit_list = []
    ci_list = []

    for (scale, dataset), g in df_data.groupby(["scale", "dataset"]):
        if verbose:
            print(f"Calculating: {scale}, {dataset}")

        X = g["X"].to_numpy()
        Y = g["Y"].to_numpy()

        # store raw points
        tmp_points = g[["X", "Y"]].copy()
        tmp_points["scale"] = scale
        tmp_points["dataset"] = dataset
        tmp_points["class_label"] = tmp_points["Y"].map({0: "NC", 1: "AE"})
        points_list.append(tmp_points[["scale", "dataset", "X", "Y", "class_label"]])

        # get fit objects
        fit_pars = df_fit_index.loc[(scale, dataset), "pars"]
        fit_cov = df_fit_index.loc[(scale, dataset), "cov"]

        # grid for smooth curve / bands
        x_fit = np.linspace(X.min(), X.max(), n_grid)

        # fitted curve
        y_fit = lg.model(x_fit, fit_pars)
        fit_list.append(
            pd.DataFrame({
                "scale": scale,
                "dataset": dataset,
                "X": x_fit,
                "y_fit": y_fit,
            })
        )

        # confidence bands
        for method in ["normal", "delta"]:
            qfun = getattr(lg, f"get_model_quantiles_{method}")
            y_low, y_high = qfun(x_fit, probs, fit_pars, fit_cov)

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

    df_points = pd.concat(points_list, ignore_index=True)
    df_fit_plot = pd.concat(fit_list, ignore_index=True)
    df_ci_plot = pd.concat(ci_list, ignore_index=True)

    return df_points, df_fit_plot, df_ci_plot


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={
        "data_NC": "blue",
        "data_AE": "orange",
        "fit": "black",
        "normal": "red",
        "delta": "green",
    },
    ls_map={
        "fit": "-",
        "normal": "--",
        "delta": "-.",
    },
    ci_alpha=0.15,
    ci_linewidth=1.5,
    fit_linewidth=2.0,
    width=5,
    height=4,
):
    

    if scales is None:
        scales = sorted(df_points["scale"].unique())
    if datasets is None:
        datasets = sorted(df_points["dataset"].unique())

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

    plot_labs = dict(
        zip(
            [(scale, dataset) for dataset in datasets for scale in scales],
            ["A", "B", "C", "D"][: len(scales) * len(datasets)],
        )
    )
    if verbose:
        print(plot_labs)

    for dataset_i, dataset in enumerate(datasets):
        for scale_i, scale in enumerate(scales):
            ax = axs[dataset_i, scale_i]

            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,
            )

            # raw points
            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}"],
                )

            # fitted curve
            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,
            )

            # CI bands + border lines
            g_ci = df_ci_plot[
                (df_ci_plot["scale"] == scale) & (df_ci_plot["dataset"] == dataset)
            ]

            for method, sub in g_ci.groupby("method"):
                sub = sub.sort_values("X")

                color = color_map[method]
                ls = ls_map[method]
                confidence = sub["confidence"].iloc[0]

                # format nicely: 95 instead of 95.0 when possible
                conf_str = f"{confidence:.0f}" if float(confidence).is_integer() else f"{confidence:.1f}"

                # light filled band
                ax.fill_between(
                    sub["X"],
                    sub["y_low"],
                    sub["y_high"],
                    color=color,
                    alpha=ci_alpha,
                    linewidth=0,
                )

                # lower border
                ax.plot(
                    sub["X"],
                    sub["y_low"],
                    color=color,
                    linestyle=ls,
                    linewidth=ci_linewidth,
                    label=f"CI: {method} {conf_str}%",
                )

                # upper border
                ax.plot(
                    sub["X"],
                    sub["y_high"],
                    color=color,
                    linestyle=ls,
                    linewidth=ci_linewidth,
                )

    for ax in axs.flat:
        ax.label_outer()

    handles, labels = axs[0, 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.2),
        framealpha=0.5,
        loc="center right",
    )

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

    plt.show()


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 results from disk if available, otherwise compute and save.

    Returns
    -------
    boots_pars_results : dict
        Dictionary with keys (scale, dataset, method) and values bootstrap parameter arrays.
    """
    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:
            boots_pars_results = pickle.load(f)
        
        return boots_pars_results

    boots_pars_results = {}

    for (scale, dataset), g in df_data.groupby(["scale", "dataset"]):
        if verbose:
            print(f"Bootstrap parameters: {scale}, {dataset}")

        X = g["X"].to_numpy()
        Y = g["Y"].to_numpy()

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

            par_fun = getattr(lg, f"get_{method}_pars")
            
            try:
                pars_res = par_fun(X, Y, m=n_boots)
            except:
                if verbose: print("Fatal problem!")
                pars_res = None

            boots_pars_results[(scale, dataset, method)] = pars_res

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

    return boots_pars_results


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,
):
    """
    Compute all quantities needed for plotting and return them as dataframes.

    Parameters
    ----------
    df_data : pd.DataFrame
        Input data with columns including scale, dataset, X, Y.
    df_fit_index : pd.DataFrame
        Fit results indexed by (scale, dataset), with columns 'pars' and 'cov'.
    lg : object
        Model object with methods:
            - model(x, pars)
            - get_model_quantiles_<method>(x, probs, pars, cov)
    alpha : float
        Two-sided significance level. Example alpha=0.05 gives 95% CI.
    boots_pars_results : dict
        Dictionary with keys (scale, dataset, method) and bootstrap parameter samples as values.
    n_grid : int
        Number of x-grid points.
    analytic_methods : tuple
        CI methods based on fitted covariance.
    bootstrap_methods : tuple
        CI methods based on bootstrap parameter samples.
    verbose : bool
        Print progress.

    Returns
    -------
    df_points : pd.DataFrame
    df_fit_plot : pd.DataFrame
    df_ci_plot : pd.DataFrame
        Long dataframe with both analytic and bootstrap confidence intervals.
    """
    probs = [alpha / 2, 1 - alpha / 2]
    confidence = 100 * (1 - alpha)

    points_list = []
    fit_list = []
    ci_list = []

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

        X = g["X"].to_numpy()
        Y = g["Y"].to_numpy()

        # raw points
        tmp_points = g[["X", "Y"]].copy()
        tmp_points["scale"] = scale
        tmp_points["dataset"] = dataset
        tmp_points["class_label"] = tmp_points["Y"].map({0: "NC", 1: "AE"})
        points_list.append(tmp_points[["scale", "dataset", "X", "Y", "class_label"]])

        # fitted objects
        fit_pars = df_fit_index.loc[(scale, dataset), "pars"]
        fit_cov = df_fit_index.loc[(scale, dataset), "cov"]

        # x-grid
        x_fit = np.linspace(X.min(), X.max(), n_grid)

        # fitted curve
        y_fit = lg.model(x_fit, fit_pars)
        fit_list.append(
            pd.DataFrame({
                "scale": scale,
                "dataset": dataset,
                "X": x_fit,
                "y_fit": y_fit,
            })
        )

        # analytic CIs
        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, fit_pars, fit_cov)

            ci_list.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,
                })
            )

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

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

            ci_list.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,
                })
            )

    df_points = pd.concat(points_list, ignore_index=True)
    df_fit_plot = pd.concat(fit_list, ignore_index=True)
    df_ci_plot = pd.concat(ci_list, ignore_index=True)

    return df_points, df_fit_plot, df_ci_plot


def plot_from_dataframes_full_ci(
    df_points,
    df_fit_plot,
    df_ci_plot,
    xlabs,
    scales=None,
    datasets=None,
    results_path=None,
    filename="logit_fit_CI_paper.pdf",
    verbose=False,
    color_map=None,
    ls_map=None,
    ci_alpha=0.12,
    ci_linewidth=1.5,
    fit_linewidth=2.0,
    width=5,
    height=4,
):
    """
    Plot data, fitted curve, and confidence intervals from prepared dataframes.
    """
    if scales is None:
        scales = list(df_points["scale"].drop_duplicates())
    if datasets is None:
        datasets = list(df_points["dataset"].drop_duplicates())

    if color_map is None:
        color_map = {
            "data_NC": "blue",
            "data_AE": "orange",
            "fit": "black",
            "normal": "red",
            "delta": "green",
            "nonparam_boots": "blue",
            "nonparam_stratified_boots": "purple",
            "parametric_boots": "cyan",
        }

    if ls_map is None:
        ls_map = {
            "fit": "-",
            "normal": "--",
            "delta": "-.",
            "nonparam_boots": ":",
            "nonparam_stratified_boots": (0, (3, 1, 1, 1)),
            "parametric_boots": (0, (5, 2)),
        }

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

    plot_labs = dict(
        zip(
            [(scale, dataset) for dataset in datasets for scale in scales],
            list("ABCDEFGHIJKLMNOPQRSTUVWXYZ")[: len(scales) * len(datasets)],
        )
    )
    if verbose:
        print(plot_labs)

    for dataset_i, dataset in enumerate(datasets):
        for scale_i, scale in enumerate(scales):
            ax = axs[dataset_i, scale_i]

            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,
            )

            # raw points
            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}"],
                )

            # fitted curve
            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,
            )

            # CI bands + border lines
            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")

                color = color_map[method]
                ls = ls_map[method]
                confidence = sub["confidence"].iloc[0]
                conf_str = (
                    f"{confidence:.0f}"
                    if float(confidence).is_integer()
                    else f"{confidence:.1f}"
                )

                ax.fill_between(
                    sub["X"],
                    sub["y_low"],
                    sub["y_high"],
                    color=color,
                    alpha=ci_alpha,
                    linewidth=0,
                )

                ax.plot(
                    sub["X"],
                    sub["y_low"],
                    color=color,
                    linestyle=ls,
                    linewidth=ci_linewidth,
                    label=f"CI: {method} {conf_str}%",
                )

                ax.plot(
                    sub["X"],
                    sub["y_high"],
                    color=color,
                    linestyle=ls,
                    linewidth=ci_linewidth,
                )

    for ax in axs.flat:
        ax.label_outer()

    handles, labels = axs[0, 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.2),
        framealpha=0.5,
        loc="center right",
    )

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

    plt.show()