from contextlib import contextmanager
from pathlib import Path
from typing import Dict, Literal, Union, List

import pandas as pd
from matplotlib import pyplot as plt
import numpy as np
import seaborn as sns
from rtg_score import compute_RTG_contribution_matrix
from joblib import Parallel, delayed

rc_params = {
    'font.family': 'sans-serif',
    'font.sans-serif': 'Arial',
    'font.size': 18,
    'xtick.labelsize': 14,
    'ytick.labelsize': 14,
    'figure.autolayout': True,
}


def display_mean_and_std_of_dataframes(
        dfs,
        ax,
        n_sigmas,
        fontsize=12
):
    dfs = [df.applymap(lambda x: [x]) for df in dfs]
    joint = dfs[0]
    for df in dfs[1:]:
        joint += df

    def compress_to_mean_and_std(seq):
        if np.isfinite(seq).sum() < 3:
            return ''
        else:
            return f' {np.nanmean(seq):.3f}\n±{np.nanstd(seq) * n_sigmas:.3f}'

    avg, labels = joint.applymap(np.mean), joint.applymap(compress_to_mean_and_std)
    rectangle = plt.Rectangle((0., 0.), avg.shape[1], avg.shape[0], fill=False, edgecolor='k', alpha=0.15, hatch='//',
                              zorder=-42)
    ax.add_patch(rectangle)
    return sns.heatmap(
        avg,
        annot=labels,
        vmin=0.5,
        vmax=1.0,
        ax=ax,
        fmt='s',
        annot_kws={"size": fontsize},
    )


def plot_rtg_contribution_with_errors_in_seans_style(
        metadata,
        embeddings,
        *,
        include_confounders_dict: Dict[str, list],
        exclude_confounders_dict: Dict[str, list],
        subsample: Union[float, Literal['bootstrap']],
        ax,
        fontsize,
        n_folds=5,
        recursion_data: bool,
        metric='euclidean',
        n_sigmas=2,
):
    def compute_rtg_for_random_state(random_state):
        if recursion_data:
            mask = metadata['treatment'].eq('EMPTY').values
        else:
            mask = np.ones(len(metadata), dtype='bool')

        if subsample == 'bootstrap':
            mask = np.where(mask)[0]
            mask = np.random.RandomState(random_state).choice(mask, size=len(mask), replace=True)
            meta_subample = metadata.iloc[mask]
            embd_subample = embeddings[mask]
        else:
            mask = mask & (np.random.RandomState(random_state).uniform(0, 1, size=len(mask)) < subsample)
            meta_subample = metadata[mask]
            embd_subample = embeddings[mask]

        rtg = compute_RTG_contribution_matrix(
            metadata=meta_subample,
            embeddings=embd_subample,
            include_confounders_dict=include_confounders_dict,
            exclude_confounders_dict=exclude_confounders_dict,
            use_fast_computations=True,
            metric=metric,
            minimal_n_samples=10,
        )
        if recursion_data:
            rtg = rtg.T
            rtg.insert(1, r'$ \mathbf{ Exclude\;\;same: \;\;\; ~}$ ', np.nan)
            rtg.columns = [x.replace('Exclude\nsame', '') for x in rtg.columns]
            rtg = rtg.T
        return rtg

    with Parallel(n_jobs=n_folds) as p:
        rtgs = p(delayed(compute_rtg_for_random_state)(state) for state in range(n_folds))

    display_mean_and_std_of_dataframes(rtgs, n_sigmas=n_sigmas, ax=ax, fontsize=fontsize)
    return rtgs


def anonymize_dataframe(df, column_name2prefix: Dict[str, str], unchanged_columns: List[str]) -> pd.DataFrame:
    result = {}
    for column, prefix in column_name2prefix.items():
        col = df[column].values
        _mapping, col_values = np.unique(col, return_inverse=True)
        result[column] = [f'{prefix}_{i}' for i in col_values]

    for col in unchanged_columns:
        result[col] = df[col].values

    return pd.DataFrame(result)

@contextmanager
def mpl_subplots(
    filename: str,
    show=True,
    nrows=1,
    ncols=1,
    sharex=False,
    sharey=False,
    nb_dpi: float = 100,
    save_dpi=200,
    facecolor="white",
    **kwargs,
):
    assert " " not in filename, f"Better choose a name without spaces {filename}"

    fig, axes = plt.subplots(
        nrows=nrows, ncols=ncols, sharex=sharex, sharey=sharey, dpi=nb_dpi, facecolor=facecolor, **kwargs
    )
    fig: plt.Figure = fig
    yield fig, axes
    fig.tight_layout()

    Path(filename + ".png").parent.mkdir(exist_ok=True)
    fig.savefig(filename + ".png", bbox_inches="tight", dpi=save_dpi)
    fig.savefig(filename + ".svg", bbox_inches="tight")
    if show:
        from IPython.display import display
        display(fig)

    plt.close(fig)