from functools import partial
from typing import Callable, Dict, List, Optional

import numpy as np
import pandas as pd
import scipy
from scipy.stats import wasserstein_distance as dist_wasserstein_1d
import sklearn
from sklearn import metrics as skm
from sklearn.ensemble import BaseEnsemble
from sklearn.preprocessing import StandardScaler
import dcor

from qolmat.utils.exceptions import NotEnoughSamples

EPS = np.finfo(float).eps

###########################
# Column-wise metrics     #
###########################


def columnwise_metric(
    df1: pd.DataFrame, df2: pd.DataFrame, df_mask: pd.DataFrame, metric: Callable, **kwargs
) -> pd.Series:
    """For each column, compute a metric score based on the true dataframe
    and the predicted dataframe

    Parameters
    ----------
    df1 : pd.DataFrame
        True dataframe
    df2 : pd.DataFrame
        Predicted dataframe
    df_mask : pd.DataFrame
        Elements of the dataframes to compute on
    metric : Callable
        metric function

    Returns
    -------
    pd.Series
        Series of scores for all columns
    """
    values = {}
    for col in df1.columns:
        df1_col = df1.loc[df_mask[col], col]
        df2_col = df2.loc[df_mask[col], col]
        assert df1_col.notna().all()
        assert df2_col.notna().all()
        values[col] = metric(df1_col, df2_col, **kwargs)
    return pd.Series(values)


def mean_squared_error(df1: pd.DataFrame, df2: pd.DataFrame, df_mask: pd.DataFrame) -> pd.Series:
    """Mean squared error between two dataframes.

    Parameters
    ----------
    df1 : pd.DataFrame
        True dataframe
    df2 : pd.DataFrame
        Predicted dataframe
    df_mask : pd.DataFrame
        Elements of the dataframes to compute on

    Returns
    -------
    pd.Series
    """
    return columnwise_metric(df1, df2, df_mask, skm.mean_squared_error)


def root_mean_squared_error(
    df1: pd.DataFrame, df2: pd.DataFrame, df_mask: pd.DataFrame
) -> pd.Series:
    """Root mean squared error between two dataframes.

    Parameters
    ----------
    df1 : pd.DataFrame
        True dataframe
    df2 : pd.DataFrame
        Predicted dataframe
    df_mask : pd.DataFrame
        Elements of the dataframes to compute on

    Returns
    -------
    pd.Series
    """
    return columnwise_metric(df1, df2, df_mask, skm.mean_squared_error, squared=False)


def mean_absolute_error(df1: pd.DataFrame, df2: pd.DataFrame, df_mask: pd.DataFrame) -> pd.Series:
    """Mean absolute error between two dataframes.

    Parameters
    ----------
    df1 : pd.DataFrame
        True dataframe
    df2 : pd.DataFrame
        Predicted dataframe
    df_mask : pd.DataFrame
        Elements of the dataframes to compute on

    Returns
    -------
    pd.Series
    """
    return columnwise_metric(df1, df2, df_mask, skm.mean_absolute_error)


def mean_absolute_percentage_error(
    df1: pd.DataFrame, df2: pd.DataFrame, df_mask: pd.DataFrame
) -> pd.Series:
    """Mean absolute percentage error between two dataframes.

    Parameters
    ----------
    df1 : pd.DataFrame
        True dataframe
    df2 : pd.DataFrame
        Predicted dataframe
    df_mask : pd.DataFrame
        Elements of the dataframes to compute on

    Returns
    -------
    pd.Series
    """
    return columnwise_metric(df1, df2, df_mask, skm.mean_absolute_percentage_error)


def _weighted_mean_absolute_percentage_error_1D(values1: pd.Series, values2: pd.Series) -> float:
    """Weighted mean absolute percentage error between two series.
    Based on https://en.wikipedia.org/wiki/Mean_absolute_percentage_error

    Parameters
    ----------
    values1 : pd.Series
        true series
    values2 : pd.Series
        predicted series

    Returns
    -------
    float
        Weighted mean absolute percentage error
    """
    return (values1 - values2).abs().sum() / values1.abs().sum()


def weighted_mean_absolute_percentage_error(
    df1: pd.DataFrame, df2: pd.DataFrame, df_mask: pd.DataFrame
) -> pd.Series:
    """Weighted mean absolute percentage error between two dataframes.

    Parameters
    ----------
    df1 : pd.DataFrame
        True dataframe
    df2 : pd.DataFrame
        Predicted dataframe
    df_mask : pd.DataFrame
        Elements of the dataframes to compute on

    Returns
    -------
    pd.Series
    """
    return columnwise_metric(df1, df2, df_mask, _weighted_mean_absolute_percentage_error_1D)


def dist_wasserstein(
    df1: pd.DataFrame, df2: pd.DataFrame, df_mask: pd.DataFrame, method: str = "columnwise"
) -> pd.Series:
    """Wasserstein distances between columns of 2 dataframes.
    Wasserstein distance can only be computed columnwise

    Parameters
    ----------
    df1 : pd.DataFrame
        True dataframe
    df2 : pd.DataFrame
        Predicted dataframe
    df_mask : pd.DataFrame
        Elements of the dataframes to compute on

    Returns
    -------
    pd.Series
        wasserstein distances
    """
    if method == "columnwise":
        return columnwise_metric(df1, df2, df_mask, scipy.stats.wasserstein_distance)
    else:
        raise AssertionError(
            f"The parameter of the function wasserstein_distance should be one of"
            f"the following: [`columnwise`], not `{method}`!"
        )


def _get_numerical_features(df1: pd.DataFrame) -> List[str]:
    """Get numerical features from dataframe

    Parameters
    ----------
    df1 : pd.DataFrame

    Returns
    -------
    List[str]
        List of numerical features

    Raises
    ------
    Exception
        No numerical feature is found
    """
    cols_numerical = df1.select_dtypes(include=np.number).columns.tolist()
    if len(cols_numerical) == 0:
        raise Exception("No numerical feature is found.")
    else:
        return cols_numerical


def _get_categorical_features(df1: pd.DataFrame) -> List[str]:
    """Get categorical features from dataframe

    Parameters
    ----------
    df1 : pd.DataFrame

    Returns
    -------
    List[str]
        List of categorical features

    Raises
    ------
    Exception
        No categorical feature is found
    """

    cols_numerical = df1.select_dtypes(include=np.number).columns.tolist()
    cols_categorical = [col for col in df1.columns.to_list() if col not in cols_numerical]
    if len(cols_categorical) == 0:
        raise Exception("No categorical feature is found.")
    else:
        return cols_categorical


def kolmogorov_smirnov_test_1D(df1: pd.Series, df2: pd.Series) -> float:
    """Compute KS test statistic of the two-sample Kolmogorov-Smirnov test for goodness of fit.
    See more in https://docs.scipy.org/doc/scipy/reference/generated/scipy.stats.ks_2samp.html.

    Parameters
    ----------
    df1 : pd.Series
        true series
    df2 : pd.Series
        predicted series

    Returns
    -------
    float
        KS test statistic
    """
    return scipy.stats.ks_2samp(df1, df2)[0]


def kolmogorov_smirnov_test(
    df1: pd.DataFrame, df2: pd.DataFrame, df_mask: pd.DataFrame
) -> pd.Series:
    """Kolmogorov Smirnov Test for numerical features.
    Lower score means better performance.

    Parameters
    ----------
    df1 : pd.DataFrame
        true dataframe
    df2 : pd.DataFrame
        predicted dataframe
    df_mask : pd.DataFrame
        Elements of the dataframes to compute on

    Returns
    -------
    pd.Series
        KS test statistic
    """
    cols_numerical = _get_numerical_features(df1)
    return columnwise_metric(
        df1[cols_numerical],
        df2[cols_numerical],
        df_mask[cols_numerical],
        kolmogorov_smirnov_test_1D,
    )


def _total_variance_distance_1D(df1: pd.Series, df2: pd.Series) -> float:
    """Compute Total Variance Distance for a categorical feature
    It is based on TVComplement in https://github.com/sdv-dev/SDMetrics

    Parameters
    ----------
    df1 : pd.Series
        true series
    df2 : pd.Series
        predicted series

    Returns
    -------
    float
        Total variance distance
    """
    list_categories = list(set(df1.unique()).union(set(df2.unique())))
    freqs1 = df1.value_counts() / len(df1)
    freqs1 = freqs1.reindex(list_categories, fill_value=0.0)
    freqs2 = df2.value_counts() / len(df2)
    freqs2 = freqs2.reindex(list_categories, fill_value=0.0)
    return (freqs1 - freqs2).abs().sum()


def total_variance_distance(
    df1: pd.DataFrame, df2: pd.DataFrame, df_mask: pd.DataFrame
) -> pd.Series:
    """Total variance distance for categorical features
    It is based on TVComplement in https://github.com/sdv-dev/SDMetrics

    Parameters
    ----------
    df1 : pd.DataFrame
        true dataframe
    df2 : pd.DataFrame
        predicted dataframe
    df_mask : pd.DataFrame
        Elements of the dataframes to compute on

    Returns
    -------
    pd.Series
        Total variance distance
    """
    cols_categorical = _get_categorical_features(df1)
    return columnwise_metric(
        df1[cols_categorical],
        df2[cols_categorical],
        df_mask[cols_categorical],
        _total_variance_distance_1D,
    )


def _check_same_number_columns(df1: pd.DataFrame, df2: pd.DataFrame):
    if len(df1.columns) != len(df2.columns):
        raise Exception("inputs have to have the same number of columns.")


def _get_correlation_pearson_matrix(df: pd.DataFrame, use_p_value: bool = True) -> pd.DataFrame:
    """Get matrix of correlation values for numerical features
    based on Pearson correlation coefficient or p-value for testing non-correlation.

    Parameters
    ----------
    df : pd.DataFrame
        dataframe
    use_p_value : bool, optional
        use the p-value instead of the correlation coefficient, by default True

    Returns
    -------
    pd.DataFrame
        Correlation matrix
    """
    cols = df.columns.tolist()
    matrix = np.zeros((len(df.columns), len(df.columns)))
    for idx_1, col_1 in enumerate(cols):
        for idx_2, col_2 in enumerate(cols):
            res = scipy.stats.mstats.pearsonr(
                df[col_1].array.reshape(-1, 1), df[col_2].array.reshape(-1, 1)
            )
            if use_p_value:
                matrix[idx_1, idx_2] = res[1]
            else:
                matrix[idx_1, idx_2] = res[0]

    return pd.DataFrame(matrix, index=cols, columns=cols)


def mean_difference_correlation_matrix_numerical_features(
    df1: pd.DataFrame,
    df2: pd.DataFrame,
    df_mask: pd.DataFrame,
    use_p_value: bool = True,
) -> pd.Series:
    """Mean absolute of differences between the correlation matrices of df1 and df2.
    based on Pearson correlation coefficient or p-value for testing non-correlation.

    Parameters
    ----------
    df1 : pd.DataFrame
        true dataframe
    df2 : pd.DataFrame
        predicted dataframe
    df_mask : pd.DataFrame
        Elements of the dataframes to compute on
    use_p_value : bool, optional
        use the p-value instead of the correlation coefficient, by default True

    Returns
    -------
    pd.Series
        Mean absolute of differences for each feature
    """
    df1 = df1[df_mask].dropna(axis=0)
    df2 = df2[df_mask].dropna(axis=0)

    _check_same_number_columns(df1, df2)

    cols_numerical = _get_numerical_features(df1)
    df_corr1 = _get_correlation_pearson_matrix(df1[cols_numerical], use_p_value=use_p_value)
    df_corr2 = _get_correlation_pearson_matrix(df2[cols_numerical], use_p_value=use_p_value)

    diff_corr = (df_corr1 - df_corr2).abs().mean(axis=1)
    return pd.Series(diff_corr, index=cols_numerical)


def _get_correlation_chi2_matrix(data: pd.DataFrame, use_p_value: bool = True) -> pd.DataFrame:
    """Get matrix of correlation values for categorical features
    based on Chi-square test of independence of variables (the test statistic or the p-value).

    Parameters
    ----------
    df : pd.DataFrame
        dataframe
    use_p_value : bool, optional
        use the p-value of the test instead of the test statistic, by default True

    Returns
    -------
    pd.DataFrame
        Correlation matrix
    """
    cols = data.columns.tolist()
    matrix = np.zeros((len(data.columns), len(data.columns)))
    for idx_1, col_1 in enumerate(cols):
        for idx_2, col_2 in enumerate(cols):
            freq = data.pivot_table(
                index=col_1, columns=col_2, aggfunc="size", fill_value=0
            ).to_numpy()
            res = scipy.stats.chi2_contingency(freq)
            if use_p_value:
                matrix[idx_1, idx_2] = res[1]
            else:
                matrix[idx_1, idx_2] = res[0]
    return pd.DataFrame(matrix, index=cols, columns=cols)


def mean_difference_correlation_matrix_categorical_features(
    df1: pd.DataFrame,
    df2: pd.DataFrame,
    df_mask: pd.DataFrame,
    use_p_value: bool = True,
) -> pd.Series:
    """Mean absolute of differences between the correlation matrix of df1 and df2
    based on Chi-square test of independence of variables (the test statistic or the p-value)

    Parameters
    ----------
    df1 : pd.DataFrame
        true dataframe
    df2 : pd.DataFrame
        predicted dataframe
    df_mask : pd.DataFrame
        Elements of the dataframes to compute on
    use_p_value : bool, optional
        use the p-value of the test instead of the test statistic, by default True

    Returns
    -------
    pd.Series
        Mean absolute of differences for each feature
    """
    df1 = df1[df_mask].dropna(axis=0)
    df2 = df2[df_mask].dropna(axis=0)

    _check_same_number_columns(df1, df2)

    cols_categorical = _get_categorical_features(df1)
    df_corr1 = _get_correlation_chi2_matrix(df1[cols_categorical], use_p_value=use_p_value)
    df_corr2 = _get_correlation_chi2_matrix(df2[cols_categorical], use_p_value=use_p_value)

    diff_corr = (df_corr1 - df_corr2).abs().mean(axis=1)
    return pd.Series(diff_corr, index=cols_categorical)


def _get_correlation_f_oneway_matrix(
    df: pd.DataFrame,
    cols_categorical: List[str],
    cols_numerical: List[str],
    use_p_value: bool = True,
) -> pd.DataFrame:
    """Get matrix of correlation values between categorical and numerical features
    based on the one-way ANOVA.

    Parameters
    ----------
    df : pd.DataFrame
        dataframe
    cols_categorical : List[str]
        list categorical columns
    cols_numerical : List[str]
        list numerical columns
    use_p_value : bool, optional
        use the p-value of the test instead of the test statistic, by default True

    Returns
    -------
    pd.DataFrame
        Correlation matrix
    """
    matrix = np.zeros((len(cols_categorical), len(cols_numerical)))
    for idx_cat, col_cat in enumerate(cols_categorical):
        for idx_num, col_num in enumerate(cols_numerical):
            category_group_lists = df.groupby(col_cat)[col_num].apply(list)
            res = scipy.stats.f_oneway(*category_group_lists)
            if use_p_value:
                matrix[idx_cat, idx_num] = res[1]
            else:
                matrix[idx_cat, idx_num] = res[0]
    return pd.DataFrame(matrix, index=cols_categorical, columns=cols_numerical)


def mean_diff_corr_matrix_categorical_vs_numerical_features(
    df1: pd.DataFrame,
    df2: pd.DataFrame,
    df_mask: pd.DataFrame,
    use_p_value: bool = True,
) -> pd.Series:
    """Mean absolute of differences between the correlation matrix of df1 and df2
    based on the one-way ANOVA.

    Parameters
    ----------
    df1 : pd.DataFrame
        true dataframe
    df2 : pd.DataFrame
        predicted dataframe
    df_mask : pd.DataFrame
        Elements of the dataframes to compute on
    use_p_value : bool, optional
        use the p-value of the test instead of the test statistic, by default True

    Returns
    -------
    pd.Series
        Mean absolute of differences for each feature
    """
    df1 = df1[df_mask].dropna(axis=0)
    df2 = df2[df_mask].dropna(axis=0)

    _check_same_number_columns(df1, df2)

    cols_categorical = _get_categorical_features(df1)
    cols_numerical = _get_numerical_features(df1)
    df_corr1 = _get_correlation_f_oneway_matrix(
        df1, cols_categorical, cols_numerical, use_p_value=use_p_value
    )
    df_corr2 = _get_correlation_f_oneway_matrix(
        df2, cols_categorical, cols_numerical, use_p_value=use_p_value
    )
    diff_corr = (df_corr1 - df_corr2).abs().mean(axis=1)
    return pd.Series(diff_corr, index=cols_categorical)


###########################
# Row-wise metrics        #
###########################


def _sum_manhattan_distances_1D(values: pd.Series) -> float:
    """Sum of Manhattan distances computed for one column
    It is based on https://www.geeksforgeeks.org/sum-manhattan-distances-pairs-points/

    Parameters
    ----------
    values : pd.Series
        Values of a column

    Returns
    -------
    float
        Sum of Manhattan distances
    """
    values = values.sort_values(ascending=True)
    sums_partial = values.shift().fillna(0.0).cumsum()
    differences_partial = values * np.arange(len(values)) - sums_partial
    res = differences_partial.sum()
    return res


def _sum_manhattan_distances(df1: pd.DataFrame) -> float:
    """Sum Manhattan distances between all pairs of rows.
    It is based on https://www.geeksforgeeks.org/sum-manhattan-distances-pairs-points/

    Parameters
    ----------
    df1 : pd.DataFrame

    Returns
    -------
    float
        Sum of Manhattan distances for all pairs of rows.
    """
    cols = df1.columns.tolist()
    result = sum([_sum_manhattan_distances_1D(df1[col]) for col in cols])
    return result


def sum_energy_distances(df1: pd.DataFrame, df2: pd.DataFrame, df_mask: pd.DataFrame) -> pd.Series:
    """Sum of energy distances between df1 and df2.
    It is based on https://dcor.readthedocs.io/en/latest/theory.html#

    Parameters
    ----------
    df1 : pd.DataFrame
        true dataframe
    df2 : pd.DataFrame
        predicted dataframe
    df_mask : pd.DataFrame
        Elements of the dataframes to compute on

    Returns
    -------
    pd.Series
        Sum of energy distances between df1 and df2.
    """

    # Replace nan in dataframe
    df1 = df1[df_mask].fillna(0.0)
    df2 = df2[df_mask].fillna(0.0)

    # sum of (len_df1 * (len_df1 - 1) / 2) distances for df1
    sum_distances_df1 = _sum_manhattan_distances(df1)
    sum_distances_df2 = _sum_manhattan_distances(df2)

    df = pd.concat([df1, df2])
    sum_distances_df1_df2 = _sum_manhattan_distances(df)
    sum_distance = 2 * sum_distances_df1_df2 - 4 * sum_distances_df1 - 4 * sum_distances_df2

    return pd.Series(sum_distance, index=["All"])


def sum_pairwise_distances(
    df1: pd.DataFrame, df2: pd.DataFrame, df_mask: pd.DataFrame, metric: str = "cityblock"
) -> float:
    """Sum of pairwise distances based on a predefined metric.
    Metrics are found in this link
    https://docs.scipy.org/doc/scipy/reference/generated/scipy.spatial.distance.cdist.html

    Parameters
    ----------
    df1 : pd.DataFrame
        First empirical distribution without nans
    df2 : pd.DataFrame
        Second empirical distribution without nans
    df_mask : pd.DataFrame
        Elements of the dataframes to compute on
    metric : str, optional
        distance metric, by default 'cityblock'

    Returns
    -------
    float
        Sum of pairwise distances based on a predefined metric
    """
    df1 = df1[df_mask.any(axis=1)]
    df2 = df2[df_mask.any(axis=1)]
    distances = np.sum(scipy.spatial.distance.cdist(df1, df2, metric=metric))

    return distances


###########################
# Dataframe-wise metrics  #
###########################


def frechet_distance(
    df1: pd.DataFrame,
    df2: pd.DataFrame,
    df_mask: pd.DataFrame,
) -> float:
    """Compute the Fréchet distance between two dataframes df1 and df2
    frechet_distance = || mu_1 - mu_2 ||_2^2 + Tr(Sigma_1 + Sigma_2 - 2(Sigma_1 . Sigma_2)^(1/2))
    It is normalized, df1 and df_ are first scaled by a factor
        (std(df1) + std(df2)) / 2
    and then centered around
        (mean(df1) + mean(df2)) / 2

    Dowson, D. C., and BV666017 Landau. "The Fréchet distance between multivariate normal
    distributions."
    Journal of multivariate analysis 12.3 (1982): 450-455.

    Parameters
    ----------
    df1 : pd.DataFrame
        true dataframe
    df2 : pd.DataFrame
        predicted dataframe
    df_mask : pd.DataFrame
        Mask indicating on which values the distance has to computed on

    Returns
    -------
    float
        frechet_distance
    """

    if df1.shape != df2.shape:
        raise Exception("inputs have to be of same dimensions.")

    df_true = df1[df_mask.any(axis=1)]
    df_pred = df2[df_mask.any(axis=1)]

    std = (np.std(df_true) + np.std(df_pred) + EPS) / 2
    mu = (np.nanmean(df_true, axis=0) + np.nanmean(df_pred, axis=0)) / 2
    df_true = (df_true - mu) / std
    df_pred = (df_pred - mu) / std

    mu_true = np.nanmean(df_true, axis=0)
    sigma_true = np.ma.cov(np.ma.masked_invalid(df_true), rowvar=False).data
    mu_pred = np.nanmean(df_pred, axis=0)
    sigma_pred = np.ma.cov(np.ma.masked_invalid(df_pred), rowvar=False).data

    ssdiff = np.sum((mu_true - mu_pred) ** 2.0)
    product = np.array(sigma_true @ sigma_pred)
    if product.ndim < 2:
        product = product.reshape(-1, 1)
    covmean = scipy.linalg.sqrtm(product)
    if np.iscomplexobj(covmean):
        covmean = covmean.real
    frechet_dist = ssdiff + np.trace(sigma_true + sigma_pred - 2.0 * covmean)

    return frechet_dist / df_true.shape[0]


def frechet_distance_pattern(
    df1: pd.DataFrame,
    df2: pd.DataFrame,
    df_mask: pd.DataFrame,
) -> pd.Series:
    """Frechet distance computed using a pattern decomposition

    Parameters
    ----------
    df1 : pd.DataFrame
        First empirical ditribution
    df2 : pd.DataFrame
        Second empirical ditribution
    df_mask : pd.DataFrame
        Mask indicating on which values the distance has to computed on

    Returns
    -------
    pd.Series
        Series of computed metrics
    """
    cols_numerical = _get_numerical_features(df1)
    distance = frechet_distance(df1[cols_numerical], df2[cols_numerical], df_mask[cols_numerical])
    return pd.Series(distance, index=["All"])


def density_from_rf(
    df: pd.DataFrame, estimator: BaseEnsemble, df_est: Optional[pd.DataFrame] = None
):
    """Estimates the density of the empirical distribution given by df at the sample points given
    by df_est. The estimation uses an random forest estimator and relies on the average number of
    samples in the leaf corresponding to each estimation point.

    Disclaimer: this method is experimental and has no known theoretical grounds

    Parameters
    ----------
    df : pd.DataFrame
        Empirical distribution which density should be estimated
    estimator : BaseEnsemble
        Estimator defining the forest upon which is based the density counting.
    df_est : pd.DataFrame, optional
        Sample points of the estimation, by default None
        If None, the density is estimated at the points given by `df`.

    Returns
    -------
    pd.Series
        Series of floats providing the normalized density
    """
    if df_est is None:
        df_est = df.copy()
    if df_est.index.names == [None]:
        cols_index = ["index"]
    else:
        cols_index = df_est.index.names
    counts = pd.Series(0, index=df_est.index)
    df_leafs = pd.DataFrame(estimator.apply(df), index=df.index)
    df_leafs_est = pd.DataFrame(estimator.apply(df_est), index=df_est.index)
    for i_tree in range(estimator.n_estimators):
        leafs = df_leafs[i_tree].rename("id_leaf")
        leafs_est = df_leafs_est[i_tree].rename("id_leaf")
        counts_leafs = leafs.value_counts().rename("count")
        df_merge = pd.merge(leafs_est.reset_index(), counts_leafs.reset_index(), on="id_leaf")
        df_merge = df_merge.set_index(cols_index)
        counts += df_merge["count"]
    counts /= counts.sum()
    return counts


def kl_divergence_1D(df1: pd.Series, df2: pd.Series) -> float:
    """Estimation of the Kullback-Leibler divergence between the two 1D empirical distributions
    given by `df1`and `df2`. The samples are binarized using a uniform spacing with 20 bins from
    the smallest to the largest value. Not that this may be a coarse estimation.

    Parameters
    ----------
    df1 : pd.Series
        First empirical distribution
    df2 : pd.Series
        Second empirical distribution

    Returns
    -------
    float
        Kullback-Leibler divergence between the two empirical distributions.
    """
    min_val = min(df1.min(), df2.min())
    max_val = max(df1.max(), df2.max())
    bins = np.linspace(min_val, max_val, 20)
    p = np.histogram(df1, bins=bins, density=True)[0]
    q = np.histogram(df2, bins=bins, density=True)[0]
    return scipy.stats.entropy(p + EPS, q + EPS)


def kl_divergence_gaussian_exact(
    mean1: pd.Series, cov1: pd.DataFrame, mean2: pd.Series, cov2: pd.DataFrame
) -> float:
    """Exact Kullback-Leibler divergence computed between two multivariate normal distributions

    Parameters
    ----------
    mean1: pd.Series
        Mean of the first distribution
    cov1: pd.DataFrame
        Covariance matrx of the first distribution
    mean2: pd.Series
        Mean of the second distribution
    cov2: pd.DataFrame
        Covariance matrx of the second distribution
    Returns
    -------
    float
        Kulback-Leibler divergence
    """
    n_variables = len(mean1)
    L1, lower1 = scipy.linalg.cho_factor(cov1)
    L2, lower2 = scipy.linalg.cho_factor(cov2)
    M = scipy.linalg.solve(L2, L1)
    y = scipy.linalg.solve(L2, mean2 - mean1)
    norm_M = (M**2).sum().sum()
    norm_y = (y**2).sum()
    term_diag_L = 2 * np.sum(np.log(np.diagonal(L2) / np.diagonal(L1)))
    print(norm_M, "-", n_variables, "+", norm_y, "+", term_diag_L)
    div_kl = 0.5 * (norm_M - n_variables + norm_y + term_diag_L)
    return div_kl


def kl_divergence_gaussian(df1: pd.DataFrame, df2: pd.DataFrame, df_mask: pd.Series) -> float:
    """Kullback-Leibler divergence estimation based on a Gaussian approximation of both empirical
    distributions

    Parameters
    ----------
    df1 : pd.DataFrame
        First empirical distribution
    df2 : pd.DataFrame
        Second empirical distribution
    df_mask: pd.DataFrame
        Mask indicating on what values the divergence should be computed

    Returns
    -------
    pd.Series
        Series of estimated metrics
    """
    df1 = df1[df_mask.any(axis=1)]
    df2 = df2[df_mask.any(axis=1)]
    cov1 = df1.cov()
    cov2 = df2.cov()
    mean1 = df1.mean()
    mean2 = df2.mean()

    div_kl = kl_divergence_gaussian_exact(mean1, cov1, mean2, cov2)
    return div_kl


def kl_divergence_forest(df1: pd.DataFrame, df2: pd.DataFrame, df_mask: pd.DataFrame) -> float:
    """Kullback-Leibler divergence estimation based on a random forest fitted on the first
    empirical distribution

    Disclaimer: this method is experimental and has no known theoretical grounds

    Parameters
    ----------
    df1 : pd.DataFrame
        First empirical distribution
    df2 : pd.DataFrame
        Second empirical distribution
    df_mask: pd.DataFrame
        Mask indicating on what values the divergence should be computed

    Returns
    -------
    pd.Series
        Series of estimated metrics
    """
    df1 = df1[df_mask.any(axis=1)]
    df2 = df2[df_mask.any(axis=1)]
    # df_1 = StandardScaler().fit_transform(df1[df_mask.any(axis=1)])
    # df_2 = StandardScaler().fit_transform(df2[df_mask.any(axis=1)])
    n_estimators = 100
    # estimator = sklearn.ensemble.RandomForestClassifier(
    #     n_estimators=n_estimators, max_depth=10
    # )
    # X = pd.concat([df1, df2])
    # y = pd.concat([pd.Series([False] * len(df1)), pd.Series([True] * len(df2))])
    # estimator.fit(X, y)
    estimator = sklearn.ensemble.RandomTreesEmbedding(n_estimators=n_estimators, random_state=123)
    estimator.fit(df1)
    counts1 = density_from_rf(df1, estimator, df_est=df2)
    counts2 = density_from_rf(df2, estimator, df_est=df2)
    div_kl = np.mean(np.log(counts1 / counts2) * counts1 / counts2)
    return div_kl


def kl_divergence(
    df1: pd.DataFrame,
    df2: pd.DataFrame,
    df_mask: pd.Series,
    method: str = "columnwise",
    min_n_rows: int = 10,
) -> pd.Series:
    """
    Estimation of the Kullback-Leibler divergence between too empirical distributions. Three
    methods are implemented:
    - columnwise, relying on a uniform binarization and only taking marginals into account
    (https://en.wikipedia.org/wiki/Kullback%E2%80%93Leibler_divergence),
    - gaussian, relying on a Gaussian approximation,
    - random_forest, experimental

    Parameters
    ----------
    df1 : pd.DataFrame
        First empirical distribution
    df2 : pd.DataFrame
        Second empirical distribution
    df_mask: pd.DataFrame
        Mask indicating on what values the divergence should be computed
    method: str
        Method used
    min_n_rows: int
        Minimum number of rows for a KL estimation

    Returns
    -------
    pd.Series
        Kullback-Leibler divergence

    Raise
    -----
    AssertionError
        If the empirical distributions do not have enough samples to estimate a KL divergence.
        Consider using a larger dataset of lowering the parameter `min_n_rows`.
    """
    if method == "columnwise":
        cols_numerical = _get_numerical_features(df1)
        return columnwise_metric(
            df1[cols_numerical],
            df2[cols_numerical],
            df_mask[cols_numerical],
            kl_divergence_1D,
        )
    elif method == "gaussian":
        cols_numerical = _get_numerical_features(df1)
        return pattern_based_weighted_mean_metric(
            df1[cols_numerical],
            df2[cols_numerical],
            df_mask[cols_numerical],
            kl_divergence_gaussian,
            min_n_rows=min_n_rows,
        )
    elif method == "random_forest":
        return pattern_based_weighted_mean_metric(
            df1, df2, df_mask, kl_divergence_forest, min_n_rows=min_n_rows
        )
    else:
        raise AssertionError(
            f"The parameter of the function wasserstein_distance should be one of"
            f"the following: [`columnwise`, `gaussian`, `random_forest`], not `{method}`!"
        )


def distance_anticorr(df1: pd.DataFrame, df2: pd.DataFrame, df_mask: pd.DataFrame) -> float:
    """Score based on the distance anticorrelation between two empirical distributions.
    The theoretical basis can be found on dcor documentation:
    https://dcor.readthedocs.io/en/latest/theory.html

    Parameters
    ----------
    df1 : pd.DataFrame
        Dataframe representing the first empirical distribution
    df2 : pd.DataFrame
        Dataframe representing the second empirical distribution
    df_mask: pd.DataFrame
        Mask indicating on what values the divergence should be computed

    Returns
    -------
    float
        Distance correlation score
    """
    df1 = df1[df_mask.any(axis=1)]
    df2 = df2[df_mask.any(axis=1)]
    return (1 - dcor.distance_correlation(df1.values, df2.values)) / 2


def pattern_based_weighted_mean_metric(
    df1: pd.DataFrame,
    df2: pd.DataFrame,
    df_mask: pd.DataFrame,
    metric: Callable,
    min_n_rows: int = 10,
    **kwargs,
) -> pd.Series:
    """Compute a mean score based on missing patterns.
    Note that for each pattern, a score is returned by the function metric.
    This code is based on https://www.statsmodels.org/

    Parameters
    ----------
    df1 : pd.DataFrame
        Dataframe representing the first empirical distribution, with nans
    df2 : pd.DataFrame
        Dataframe representing the second empirical distribution
    df_mask : pd.DataFrame
        Elements of the dataframes to compute on
    metric : Callable
        metric function
    min_n_rows : int, optional
        minimum number of row allowed for a pattern without nan, by default 10

    Returns
    -------
    pd.Series
        _description_
    """
    scores = []
    weights = []
    df1 = df1[df_mask.any(axis=1)]
    df2 = df2[df_mask.any(axis=1)]
    df_nan = df1.notna()
    max_num_row = 0
    for tup_pattern, df_nan_pattern in df_nan.groupby(df_nan.columns.tolist()):
        ind_pattern = df_nan_pattern.index
        df1_pattern = df1.loc[ind_pattern, list(tup_pattern)]
        max_num_row = max(max_num_row, len(df1_pattern))
        if not any(tup_pattern) or len(df1_pattern) < min_n_rows:
            continue
        df2_pattern = df2.loc[ind_pattern, list(tup_pattern)]
        df_mask_pattern = df_mask.loc[ind_pattern, list(tup_pattern)]
        weights.append(len(df1_pattern) / len(df1))
        scores.append(metric(df1_pattern, df2_pattern, df_mask_pattern, **kwargs))
    if len(scores) == 0:
        raise NotEnoughSamples(max_num_row, min_n_rows)
    return pd.Series(sum([s * w for s, w in zip(scores, weights)]), index=["All"])


def get_metric(name: str) -> Callable:
    dict_metrics: Dict[str, Callable] = {
        "mse": mean_squared_error,
        "rmse": root_mean_squared_error,
        "mae": mean_absolute_error,
        "wmape": weighted_mean_absolute_percentage_error,
        "wasserstein_columnwise": dist_wasserstein,
        "KL_columnwise": partial(kl_divergence, method="columnwise"),
        "KL_gaussian": partial(kl_divergence, method="gaussian"),
        "KL_forest": partial(kl_divergence, method="random_forest"),
        "ks_test": kolmogorov_smirnov_test,
        "correlation_diff": mean_difference_correlation_matrix_numerical_features,
        "energy": sum_energy_distances,
        "frechet": frechet_distance,
        "dist_corr_pattern": partial(
            pattern_based_weighted_mean_metric,
            metric=distance_anticorr,
        ),
    }
    return dict_metrics[name]
