import matplotlib as mpl
import matplotlib.cm as cm
import matplotlib.colorbar as cbar
import matplotlib.pyplot as plt
import numpy as np
from mpl_toolkits.axes_grid1 import make_axes_locatable

from extreme_estimator.margin_fits.extreme_params import ExtremeParams
from extreme_estimator.margin_fits.plot.shifted_color_map import shiftedColorMap


def plot_extreme_param(ax, label: str, values: np.ndarray):
    # Load the shifted cmap to center on a middle point
    vmin, vmax = np.min(values), np.max(values)
    if vmin < 0 < vmax:
        midpoint = 1 - vmax / (vmax + abs(vmin))
    elif vmin < 0 and vmax < 0:
        midpoint = 1.0
    elif vmin > 0 and vmax > 0:
        midpoint = 0.0
    else:
        raise ValueError('Unexpected values: vmin={}, vmax={}'.format(vmin, vmax))
    cmap = [plt.cm.coolwarm, plt.cm.bwr, plt.cm.seismic][1]
    shifted_cmap = shiftedColorMap(cmap, midpoint=midpoint, name='shifted')
    norm = mpl.colors.Normalize(vmin=vmin, vmax=vmax)
    divider = make_axes_locatable(ax)
    cax = divider.append_axes('right', size='5%', pad=0.03)
    cb = cbar.ColorbarBase(cax, cmap=shifted_cmap, norm=norm)
    cb.set_label(label)
    return norm, shifted_cmap


def get_color_rbga_shifted(ax, replace_blue_by_white: bool, values: np.ndarray, label=None):
    """
    For some display it was necessary to transform dark blue values into white values
    """
    norm, shifted_cmap = plot_extreme_param(ax, label, values)
    m = cm.ScalarMappable(norm=norm, cmap=shifted_cmap)
    colors = [m.to_rgba(value) for value in values]
    # We do not want any blue values for parameters other than the Shape
    # So when the value corresponding to the blue color is 1, then we set the color to white, i.e. (1,1,1,1)
    if replace_blue_by_white:
        colors = [color if color[2] != 1 else (1, 1, 1, 1) for color in colors]
    return colors


def imshow_shifted(ax, gev_param_name, values, x, y):
    masked_array = np.ma.masked_where(np.isnan(values), values)
    norm, shifted_cmap = plot_extreme_param(ax, gev_param_name, masked_array)
    shifted_cmap.set_bad(color='white')
    if gev_param_name != ExtremeParams.SHAPE:
        epsilon = 1e-2 * (np.max(values) - np.min(values))
        value = np.min(values)
        # The right blue corner will be blue (but most of the time, another display will be on top)
        masked_array[-1, -1] = value - epsilon
    ax.imshow(masked_array, extent=(x.min(), x.max(), y.min(), y.max()), cmap=shifted_cmap)