############################################
# imports
############################################

import numpy as np
import pandas as pd
import seaborn as sns
import matplotlib
import matplotlib.pyplot as plt
from matplotlib.patches import Patch

import fgclustering.utils as utils


############################################
# functions
############################################


def _plot_global_feature_importance(p_value_of_features, save):
    """Plot global feature importance based on p-values given as input.

    :param p_value_of_features: dictionary where keys are names of features and values are p-values of these features
    :type p_value_of_features: dict
    :param save: Filename to save plot.
    :type save: str
    """
    p_value_of_features = p_value_of_features.copy()
    p_value_of_features.pop("target")
    p_value_of_features.pop("cluster")
    importance = pd.DataFrame(p_value_of_features, index=[0])
    importance = pd.melt(importance)
    importance.sort_values(by="value", ascending=True, inplace=True)
    importance.value = 1 - importance.value

    n_features = importance.shape[0]
    figure_size = max(6.5, int(np.ceil(5 * n_features / 25)))

    sns.set_theme(style="whitegrid")
    plt.figure(
        figsize=(6.5, figure_size)
    )  # keep width default, change height depending on the number of features
    plot = sns.barplot(data=importance, x="value", y="variable", color="#3470a3")
    plot.set_xlabel("importance")
    plot.set_ylabel("feature")
    plt.title("Global Feature Importance")
    plt.tight_layout()

    if save is not None:
        plt.savefig(
            "{}_global_feature_importance.png".format(save),
            bbox_inches="tight",
            dpi=300,
        )
    plt.show()


def _plot_local_feature_importance(
    p_value_of_features_per_cluster, thr_pvalue, num_cols, save
):
    """Plot local feature importance to show the importance of each feature for each cluster.

    :param p_value_of_features_per_cluster: p-value matrix of all features per cluster.
    :type p_value_of_features_per_cluster: pandas.DataFrame
    :param thr_pvalue: P-value threshold used for feature filtering
    :type thr_pvalue: float, optional
    :param num_cols: Number of plots in one row.
    :type num_cols: int
    :param save: Filename to save plot.
    :type save: str
    """
    importance = 1 - p_value_of_features_per_cluster

    X_barplot = pd.melt(importance, ignore_index=False)
    X_barplot = X_barplot.rename_axis("feature").reset_index(level=0, inplace=False)
    X_barplot = X_barplot.sort_values("value", ascending=False)

    n_features = importance.shape[0]
    figure_size = max(6.5, int(np.ceil(5 * n_features / 25)))
    num_cols = min(num_cols, len(importance.columns))

    sns.set_theme(style="whitegrid")
    plot = sns.FacetGrid(
        X_barplot, col="variable", sharey=False, col_wrap=num_cols, height=figure_size
    )
    plot.map(sns.barplot, "value", "feature", color="#3470a3")
    plot.set_axis_labels("importance", "feature")
    plot.set_titles(col_template="Cluster {col_name}")
    plt.suptitle(
        f"Local Feature Importance - Showing features with p-value < {thr_pvalue}"
    )
    plt.tight_layout()

    if save is not None:
        plt.savefig(
            "{}_local_feature_importance.png".format(save), bbox_inches="tight", dpi=300
        )
    plt.show()


def _plot_heatmap(data_clustering_ranked, thr_pvalue, model_type, save):
    """Plot feature heatmap sorted by clusters, where features are filtered and ranked
    with statistical tests (ANOVA for continuous featres, chi square for categorical features).

    :param data_clustering_ranked: Filtered and ranked data frame incl features, target and cluster numbers.
    :type data_clustering_ranked: pandas.DataFrame
    :param thr_pvalue: P-value threshold used for feature filtering
    :type thr_pvalue: float, optional
    :param model_type: Model type of Random Forest model: classifier or regression.
    :type model_type: str
    :param save: Filename to save plot.
    :type save: str
    """
    data_clustering_ranked = data_clustering_ranked.copy()
    target_values_original = data_clustering_ranked["target"]

    for feature in data_clustering_ranked.columns:
        if pd.api.types.is_string_dtype(data_clustering_ranked[feature]):
            data_clustering_ranked[feature] = data_clustering_ranked[feature].astype(
                "category"
            )
        if (
            pd.api.types.is_categorical_dtype(data_clustering_ranked[feature])
            and feature != "cluster"
        ):
            data_clustering_ranked[feature] = data_clustering_ranked[feature].cat.codes

    data_scaled = utils.scale_minmax(data_clustering_ranked)
    data_heatmap = pd.DataFrame(columns=data_scaled.columns)
    target_values_scaled = data_scaled["target"]

    one_percent_of_number_of_samples = int(np.ceil(0.01 * len(data_clustering_ranked)))

    for cluster in data_scaled.cluster.unique():
        data_heatmap = pd.concat(
            [data_heatmap, data_scaled[data_scaled.cluster == cluster]],
            ignore_index=True,
        )
        data_heatmap = pd.concat(
            [
                data_heatmap,
                pd.DataFrame(
                    np.nan,
                    index=np.arange(one_percent_of_number_of_samples),
                    # blank lines which are 1% of num samples
                    columns=data_scaled.columns,
                ),
            ],
            ignore_index=True,
        )
    data_heatmap = data_heatmap[:-5]
    data_heatmap.drop("cluster", axis=1, inplace=True)

    n_samples, n_features = data_heatmap.shape
    heatmap_ = np.zeros((n_features, n_samples, 4))
    cmap_features = matplotlib.cm.get_cmap("coolwarm").copy()
    cmap_target = matplotlib.cm.get_cmap("viridis")  # gist_ncar

    for feature in range(n_features):
        for sample in range(n_samples):
            if feature == 0:
                heatmap_[feature, sample, :] = cmap_target(
                    data_heatmap.iloc[sample, feature]
                )
            else:
                heatmap_[feature, sample, :] = cmap_features(
                    data_heatmap.iloc[sample, feature]
                )

    figure_size = max(6.5, int(np.ceil(5 * n_features / 25)))

    sns.set_theme(style="white")
    fig = plt.figure(figsize=(figure_size, figure_size))
    img = plt.imshow(heatmap_, interpolation="none", aspect="auto")

    plt.suptitle(
        f"Subgroups of instances that follow similar decision paths in the RF model \n Showing features with p-value < {thr_pvalue}"
    )
    plt.xticks([], [])
    plt.yticks(range(n_features), data_heatmap.columns)

    # remove bounding box
    for spine in plt.gca().spines.values():
        spine.set_visible(False)

    if model_type == "regression":
        norm = matplotlib.colors.Normalize(
            vmin=target_values_original.min(), vmax=target_values_original.max()
        )
        cbar_target = plt.colorbar(
            matplotlib.cm.ScalarMappable(norm=norm, cmap=cmap_target)
        )
        cbar_target.set_label("target")
    else:
        legend_elements = [
            Patch(
                facecolor=cmap_target(tv_n),
                edgecolor=cmap_target(tv_n),
                label=f"{tv_o}",
            )
            for tv_n, tv_o in zip(
                target_values_scaled.unique(), target_values_original.unique()
            )
        ]
        ll = plt.legend(
            handles=legend_elements,
            bbox_to_anchor=(0, 0),
            loc="upper left",
            ncol=min(6, len(legend_elements)),
            title="target",
        )

    norm = matplotlib.colors.Normalize(vmin=0, vmax=1)
    cbar_features = plt.colorbar(
        matplotlib.cm.ScalarMappable(norm=norm, cmap=cmap_features)
    )
    cbar_features.set_label("standardized feature values")

    if save is not None:
        plt.savefig("{}_heatmap.png".format(save), bbox_inches="tight", dpi=300)
    plt.show()


def _plot_distributions(data_clustering_ranked, thr_pvalue, num_cols, save):
    """Plot feature boxplots (for continuous features) or barplots (for categorical features) divided by clusters,
    where features are filtered and ranked by p-value of a statistical test (ANOVA for continuous features,
    chi square for categorical features).

    :param data_clustering_ranked: Filtered and ranked data frame incl features, target and cluster numbers.
    :type data_clustering_ranked: pandas.DataFrame
    :param thr_pvalue: P-value threshold used for feature filtering
    :type thr_pvalue: float, optional
    :param num_cols: Number of plots in one row.
    :type num_cols: int
    :param save: Filename to save plot.
    :type save: str
    """
    data_clustering_ranked = data_clustering_ranked.copy()

    variables_to_plot = data_clustering_ranked.drop(
        ["target", "cluster"], axis=1, inplace=False
    ).columns.to_list()
    variables_to_plot = [
        "target"
    ] + variables_to_plot  # adding target, to plot it first

    num_rows = int(np.ceil(len(variables_to_plot) / num_cols))
    plt.figure(figsize=(num_cols * 4.5, num_rows * 4.5))
    plt.tight_layout()
    plt.subplots_adjust(top=0.95, hspace=0.8, wspace=0.8)
    plt.suptitle(
        f"Distribution of feature values across subgroups - Showing features with p-value < {thr_pvalue}",
        fontsize=14,
    )

    for n, feature in enumerate(variables_to_plot):
        # add a new subplot iteratively
        ax = plt.subplot(num_rows, num_cols, n + 1)
        if data_clustering_ranked[
            feature
        ].nunique() < 5 or pd.api.types.is_categorical_dtype(
            data_clustering_ranked[feature]
        ):
            sns.countplot(
                x="cluster",
                hue=feature,
                data=data_clustering_ranked,
                ax=ax,
                palette=sns.color_palette(
                    "Blues_r", n_colors=len(np.unique(data_clustering_ranked[feature]))
                ),
            )
            ax.set_title("Feature: {}".format(feature))
            ax.legend(bbox_to_anchor=(1, 1), loc=2)
        else:
            sns.boxplot(
                x="cluster",
                y=feature,
                data=data_clustering_ranked,
                ax=ax,
                color="#3470a3",
            )
            ax.set_title("Feature: {}".format(feature))

    if save is not None:
        plt.savefig("{}_boxplots.png".format(save), bbox_inches="tight", dpi=300)
    plt.show()
