__all__ = ["plot_value_distributions"]

import matplotlib

import matplotlib.figure

import matplotlib.pyplot as plt

import networkx as nx

import numpy as np

from move.core.typing import FloatArray

from move.visualization.style import (

    DEFAULT_DIVERGING_PALETTE,

    DEFAULT_PLOT_STYLE,

    style_settings,

)

def plot_value_distributions(

    feature_values: FloatArray,

    style: str = "fast",

    nbins: int = 100,

    colormap: str = DEFAULT_DIVERGING_PALETTE,

) -> matplotlib.figure.Figure:

    """

    Given a certain dataset, plot its distribution of values.

    Args:

        feature_values:

            Values of the features, a 2D array (`num_samples` x `num_features`).

        style:

            Name of style to apply to the plot.

        colormap:

            Name of colormap to apply to the colorbar.

    Returns:

        Figure

    """

    vmin, vmax = np.nanmin(feature_values), np.nanmax(feature_values)

    with style_settings(style):

        fig = plt.figure(layout="constrained", figsize=(7, 7))

        ax = fig.add_subplot(projection="3d")

        x_val = np.linspace(vmin, vmax, nbins)

        y_val = np.arange(np.shape(feature_values)[1])

        x_val, y_val = np.meshgrid(x_val, y_val)

        histogram = []

        for i in range(np.shape(feature_values)[1]):

            feat_i_list = feature_values[:, i]

            feat_hist, feat_bin_edges = np.histogram(

                feat_i_list, bins=nbins, range=(vmin, vmax)

            )

            histogram.append(feat_hist)

        ax.plot_surface(x_val, y_val, np.array(histogram), cmap=colormap)

        ax.set_xlabel("Feature value")

        ax.set_ylabel("Feature ID number")

        ax.set_zlabel("Frequency")

        # ax.legend()

    return fig

def plot_reconstruction_diff(

    diff_array: FloatArray,

    vmin=None,

    vmax=None,

    style: str = DEFAULT_PLOT_STYLE,

    colormap: str = DEFAULT_DIVERGING_PALETTE,

) -> matplotlib.figure.Figure:

    """

    Plot the reconstruction differences as a heatmap.

    """

    with style_settings(style):

        if vmin is None:

            vmin = np.min(diff_array)

        elif vmax is None:

            vmax = np.max(diff_array)

        fig = plt.figure(layout="constrained", figsize=(7, 7))

        plt.imshow(diff_array, cmap=colormap, vmin=vmin, vmax=vmax)

        plt.xlabel("Feature")

        plt.ylabel("Sample")

        plt.colorbar()

    return fig

def plot_feature_association_graph(

    association_df, output_path, layout="circular", style: str = DEFAULT_PLOT_STYLE

) -> matplotlib.figure.Figure:

    """

    This function plots a graph where each node corresponds to a feature and the edges

    represent the associations between features. Edge width represents the probability

    of said association, not the association's effect size.

    Input:

        association_df: pandas dataframe containing the following columns:

                            - feature_a: source node

                            - feature_b: target node

                            - p_value/bayes_score: edge weight

        output_path: Path object where the picture will be stored.

    Output:

        Feature_association_graph.png: picture of the graph

    """

    if "p_value" in association_df.columns:

        association_df["weight"] = 1 - association_df["p_value"]

    elif "proba" in association_df.columns:

        association_df["weight"] = association_df["proba"]

    elif "ks_distance" in association_df.columns:

        association_df["weight"] = association_df["ks_distance"]

    with style_settings(style):

        fig = plt.figure(figsize=(45, 45))

        G = nx.from_pandas_edgelist(

            association_df,

            source="feature_a_name",

            target="feature_b_name",

            edge_attr="weight",

        )

        nodes = list(G.nodes)

        datasets = association_df["feature_b_dataset"].unique()

        color_map = {

            dataset: (np.random.uniform(), np.random.uniform(), np.random.uniform())

            for dataset in datasets

        }

        node_dataset_map = {

            target_feature: dataset

            for (target_feature, dataset) in zip(

                association_df["feature_b_name"], association_df["feature_b_dataset"]

            )

        }

        if layout == "spring":

            pos = nx.spring_layout(G)

            with_labels = True

        elif layout == "circular":

            pos = nx.circular_layout(G)

            _ = [

                plt.text(

                    pos[node][0],

                    pos[node][1],

                    nodes[i],

                    rotation=(i / float(len(nodes))) * 360,

                    fontsize=10,

                    horizontalalignment="center",

                    verticalalignment="center",

                )

                for i, node in enumerate(nodes)

            ]

            with_labels = False

        else:

            raise ValueError(

                "Graph layout (layout argument) must be either 'circular' or 'spring'."

            )

        nx.draw(

            G,

            pos=pos,

            with_labels=with_labels,

            node_size=2000,

            node_color=[

                (

                    color_map[node_dataset_map[feature]]

                    if feature in node_dataset_map.keys()

                    else "white"

                )

                for feature in G.nodes

            ],

            edge_color=list(nx.get_edge_attributes(G, "weight").values()),

            font_color="black",

            font_size=10,

            edge_cmap=matplotlib.colormaps["Purples"],

            connectionstyle="arc3, rad=1",

        )

        plt.tight_layout()

        fig.savefig(

            output_path / f"Feature_association_graph_{layout}.png", format="png"

        )

    return fig

def plot_feature_mean_median(

    array: FloatArray, axis=0, style: str = DEFAULT_PLOT_STYLE

) -> matplotlib.figure.Figure:

    """

    Plot feature values together with the mean, median, min and max values

    at each array position.

    """

    with style_settings(style):

        fig = plt.figure(figsize=(15, 3))

        y = np.mean(array, axis=axis)

        y_2 = np.median(array, axis=axis)

        y_3 = np.max(array, axis=axis)

        y_4 = np.min(array, axis=axis)

        plt.plot(np.arange(len(y)), y, "bo", label="mean")

        plt.plot(np.arange(len(y_2)), y_2, "ro", label="median")

        plt.plot(np.arange(len(y_3)), y_3, "go", label="max")

        plt.plot(np.arange(len(y_4)), y_4, "yo", label="min")

        plt.legend()

        plt.xlabel("feature")

        plt.ylabel("mean/median/min/max")

    return fig

def plot_reconstruction_movement(

    baseline_recon: FloatArray,

    perturb_recon: FloatArray,

    k: int,

    style: str = DEFAULT_PLOT_STYLE,

) -> matplotlib.figure.Figure:

    """

    Plot, for each sample, the change in value from the unperturbed reconstruction to

    the perturbed reconstruction. Blue lines are left/negative shifts,

    red lines are right/positive shifts.

    Args:

        baseline_recon: baseline reconstruction array with s samples

                        and k features (s,k).

        perturb_recon:  perturbed

        k: feature index. The shift (movement) of this feature's reconstruction

                          will be plotted for all samples s.

    """

    with style_settings(style):

        # Feature changes

        fig = plt.figure(figsize=(25, 25))

        for s in range(np.shape(baseline_recon)[0]):

            plt.arrow(

                baseline_recon[s, k],

                s / 100,

                perturb_recon[s, k],

                0,

                length_includes_head=True,

                color=["r" if baseline_recon[s, k] < perturb_recon[s, k] else "b"][0],

            )

        plt.ylabel("Sample (e2)", size=40)

        plt.xlabel("Feature_value", size=40)

    return fig

def plot_cumulative_distributions(

    edges: FloatArray,

    hist_base: FloatArray,

    hist_pert: FloatArray,

    title: str,

    style: str = DEFAULT_PLOT_STYLE,

) -> matplotlib.figure.Figure:

    """

    Plot the cumulative distribution of the histograms for the baseline

    and perturbed reconstructions. This is useful to visually assess the

    magnitude of the shift corresponding to the KS score.

    """

    with style_settings(style):

        # Cumulative distribution:

        fig = plt.figure(figsize=(7, 7))

        plt.plot(

            (edges[:-1] + edges[1:]) / 2,

            np.cumsum(hist_base),

            color="blue",

            label="baseline",

            alpha=0.5,

        )

        plt.plot(

            (edges[:-1] + edges[1:]) / 2,

            np.cumsum(hist_pert),

            color="red",

            label="Perturbed",

            alpha=0.5,

        )

        plt.title(f"{title}.png")

        plt.xlabel("Feature value")

        plt.ylabel("Cumulative distribution")

        plt.legend()

    return fig

def plot_correlations(

    x: FloatArray,

    y: FloatArray,

    x_pol: FloatArray,

    y_pol: FloatArray,

    a2: float,

    a1: float,

    a: float,

    k: int,

    style: str = DEFAULT_PLOT_STYLE,

) -> matplotlib.figure.Figure:

    """

    Plot y vs x and the corresponding polynomial fit.

    """

    with style_settings(style):

        # Plot correlations

        fig = plt.figure(figsize=(7, 7))

        plt.plot(x, y, marker=".", lw=0, markersize=1, color="red")

        plt.plot(

            x_pol,

            y_pol,

            color="blue",

            label="{0:.2f}x^2 {1:.2f}x {2:.2f}".format(a2, a1, a),

            lw=1,

        )

        plt.plot(x_pol, x_pol, lw=1, color="k")

        plt.xlabel(f"Feature {k} baseline values ")

        plt.ylabel(f"Feature {k} baseline  value reconstruction")

        plt.legend()

    return fig

def get_2nd_order_polynomial(x_array, y_array, n_points=100):

    """

    Given a set of x an y values, find the 2nd oder polynomial fitting best the data.

    Returns:

        x_pol: x coordinates for the polynomial function evaluation.

        y_pol: y coordinates for the polynomial function evaluation.

    """

    a2, a1, a = np.polyfit(x_array, y_array, deg=2)

    x_pol = np.linspace(np.min(x_array), np.max(x_array), n_points)

    y_pol = a2 * x_pol**2 + a1 * x_pol + a

    return x_pol, y_pol, (a2, a1, a)