__all__ = ["plot_latent_space_with_cat", "plot_latent_space_with_con"]

from typing import Any

import matplotlib.cm as cm
import matplotlib.figure
import matplotlib.pyplot as plt
import matplotlib.style
import numpy as np
from matplotlib.colors import Normalize, TwoSlopeNorm

from move.core.typing import BoolArray, FloatArray
from move.visualization.style import (
    DEFAULT_DIVERGING_PALETTE,
    DEFAULT_PLOT_STYLE,
    DEFAULT_QUALITATIVE_PALETTE,
    color_cycle,
    style_settings,
)


def plot_latent_space_with_cat(
    latent_space: FloatArray,
    feature_name: str,
    feature_values: FloatArray,
    feature_mapping: dict[str, Any],
    is_nan: BoolArray,
    style: str = DEFAULT_PLOT_STYLE,
    colormap: str = DEFAULT_QUALITATIVE_PALETTE,
) -> matplotlib.figure.Figure:
    """Plot a 2D latent space together with a legend mapping the latent
    space to the values of a discrete feature.

    Args:
        latent_space:
            Embedding, a ND array with at least two dimensions.
        feature_name:
            Name of categorical feature
        feature_values:
            Values of categorical feature
        feature_mapping:
            Mapping of codes to categories for the categorical feature
        is_nan:
            Array of bool values indicating which feature values are NaNs
        style:
            Name of style to apply to the plot
        colormap:
            Name of qualitative colormap to use for each category

    Raises:
        ValueError: If latent space does not have at least two dimensions.

    Returns:
        Figure
    """
    if latent_space.ndim < 2:
        raise ValueError("Expected at least two dimensions in latent space.")
    with style_settings(style), color_cycle(colormap):
        fig, ax = plt.subplots()
        codes = np.unique(feature_values)
        for code in codes:
            category = feature_mapping[str(code)]
            is_category = (feature_values == code) & ~is_nan
            dims = np.take(latent_space.compress(is_category, axis=0), [0, 1], axis=1).T
            ax.scatter(*dims, label=category)
        dims = np.take(latent_space.compress(is_nan, axis=0), [0, 1], axis=1).T
        ax.scatter(*dims, label="NaN")
        ax.set(xlabel="dim 0", ylabel="dim 1")
        legend = ax.legend()
        legend.set_title(feature_name)
    return fig


def plot_latent_space_with_con(
    latent_space: FloatArray,
    feature_name: str,
    feature_values: FloatArray,
    style: str = DEFAULT_PLOT_STYLE,
    colormap: str = DEFAULT_DIVERGING_PALETTE,
) -> matplotlib.figure.Figure:
    """Plot a 2D latent space together with a colorbar mapping the latent
    space to the values of a continuous feature.

    Args:
        latent_space: Embedding, a ND array with at least two dimensions.
        feature_name: Name of continuous feature
        feature_values: Values of continuous feature
        style: Name of style to apply to the plot
        colormap: Name of colormap to use for the colorbar

    Raises:
        ValueError: If latent space does not have at least two dimensions.

    Returns:
        Figure
    """
    if latent_space.ndim < 2:
        raise ValueError("Expected at least two dimensions in latent space.")
    norm = TwoSlopeNorm(0.0, min(feature_values), max(feature_values))
    with style_settings(style):
        fig, ax = plt.subplots()
        dims = latent_space[:, 0], latent_space[:, 1]
        pts = ax.scatter(*dims, c=feature_values, cmap=colormap, norm=norm)
        cbar = fig.colorbar(pts, ax=ax)
        cbar.ax.set(ylabel=feature_name)
        ax.set(xlabel="dim 0", ylabel="dim 1")
    return fig


def plot_3D_latent_and_displacement(
    mu_baseline: FloatArray,
    mu_perturbed: FloatArray,
    feature_values: FloatArray,
    feature_name: str,
    show_baseline: bool = True,
    show_perturbed: bool = True,
    show_arrows: bool = True,
    step: int = 1,
    altitude: int = 30,
    azimuth: int = 45,
) -> matplotlib.figure.Figure:
    """
    Plot the movement of the samples in the 3D latent space after perturbing one
    input variable.

    Args:
        mu_baseline:
            ND array with dimensions n_samples x n_latent_nodes containing
            the latent representation of each sample
        mu_perturbed:
            ND array with dimensions n_samples x n_latent_nodes containing
            the latent representation of each sample after perturbing the input
        feature_values:
            1D array with feature values to map to a colormap ("bwr"). Each sample is
            colored according to its value for the feature of interest.
        feature_name:
            name of the feature mapped to a colormap
        show_baseline:
            plot orginal location of the samples in the latent space
        show_perturbed:
            plot final location (after perturbation) of the samples in latent space
        show_arrows:
            plot arrows from original to final location of each sample
        angle:
            elevation from dim1-dim2 plane for the visualization of latent space.

    Raises:
        ValueError: If latent space is not 3-dimensional (3 hidden nodes).
    Returns:
        Figure
    """
    if [np.shape(mu_baseline)[1], np.shape(mu_perturbed)[1]] != [3, 3]:
        raise ValueError(
            " The latent space must be 3-dimensional. Redefine num_latent to 3."
        )

    fig = plt.figure(layout="constrained", figsize=(10, 10))
    ax = fig.add_subplot(projection="3d")
    ax.view_init(altitude, azimuth)

    if show_baseline:
        # vmin, vmax = np.min(feature_values[::step]), np.max(feature_values[::step])
        # abs_max = np.max([abs(vmin), abs(vmax)])
        ax.scatter(
            mu_baseline[::step, 0],
            mu_baseline[::step, 1],
            mu_baseline[::step, 2],
            marker=".",
            c=feature_values[::step],
            s=10,
            lw=0,
            cmap="seismic",
            vmin=-2,
            vmax=2,
        )
        ax.set_title(feature_name)
        fig.colorbar(
            cm.ScalarMappable(cmap="seismic", norm=Normalize(-2, 2)), ax=ax
        )  # Normalize(min(feature_values[::step]),max(feature_values[::step]))), ax=ax)
    if show_perturbed:
        ax.scatter(
            mu_perturbed[::step, 0],
            mu_perturbed[::step, 1],
            mu_perturbed[::step, 2],
            marker=".",
            color="lightblue",
            label="perturbed",
            lw=0.5,
        )
    if show_arrows:
        u = mu_perturbed[::step, 0] - mu_baseline[::step, 0]
        v = mu_perturbed[::step, 1] - mu_baseline[::step, 1]
        w = mu_perturbed[::step, 2] - mu_baseline[::step, 2]

        # module = np.sqrt(u * u + v * v + w * w)

        max_u, max_v, max_w = np.max(abs(u)), np.max(abs(v)), np.max(abs(w))
        # Arrow colors will be weighted contributions of
        # red -> dim1,
        # green -> dim2,
        # and blue-> dim3.
        # I.e. purple arrow means movement in dims 1 and 3
        colors = [
            (abs(du) / max_u, abs(dv) / max_v, abs(dw) / max_w, 0.7)
            for du, dv, dw in zip(u, v, w)
        ]
        ax.quiver(
            mu_baseline[::step, 0],
            mu_baseline[::step, 1],
            mu_baseline[::step, 2],
            u,
            v,
            w,
            color=colors,
            lw=0.8,
        )  # alpha=(1-module/np.max(module))**6, arrow_length_ratio=0)
        # help(ax.quiver)
    ax.set_xlabel("Dim 1")
    ax.set_ylabel("Dim 2")
    ax.set_zlabel("Dim 3")
    # ax.set_axis_off()

    return fig