import numpy as np
from . import evaluate
from skimage import color
from matplotlib import cm, pyplot as plt
import itertools as it
from math import ceil

###########################
# VISUALIZATION FUNCTIONS #
###########################

def imshow_grey(im, axis=None):
    """Show a segmentation using a gray colormap.

    Parameters
    ----------
    im : np.ndarray of int, shape (M, N)
        The segmentation to be displayed.

    Returns
    -------
    fig : plt.Figure
        The image shown.
    """
    if axis is None:
        fig, axis = plt.subplots()
    return axis.imshow(im, cmap='gray')


def imshow_magma(im, axis=None):
    """Show a segmentation using a magma colormap.

    Parameters
    ----------
    im : np.ndarray of int, shape (M, N)
        The segmentation to be displayed.

    Returns
    -------
    fig : plt.Figure
        The image shown.
    """
    if axis is None:
        fig, axis = plt.subplots()
    return axis.imshow(im, cmap='magma')


def imshow_rand(im, axis=None, labrandom=True):
    """Show a segmentation using a random colormap.

    Parameters
    ----------
    im : np.ndarray of int, shape (M, N)
        The segmentation to be displayed.
    labrandom : bool, optional
        Use random points in the Lab colorspace instead of RGB.

    Returns
    -------
    fig : plt.Figure
        The image shown.
    """
    if axis is None:
        fig, axis = plt.subplots()
    rand_colors = np.random.random(size=(ceil(np.max(im)), 3))
    if labrandom:
        rand_colors[:, 0] = rand_colors[:, 0] * 81 + 39
        rand_colors[:, 1] = rand_colors[:, 1] * 185 - 86
        rand_colors[:, 2] = rand_colors[:, 2] * 198 - 108
        rand_colors = color.lab2rgb(rand_colors[np.newaxis, ...])[0]
        rand_colors[rand_colors < 0] = 0
        rand_colors[rand_colors > 1] = 1
    rcmap = cm.colors.ListedColormap(np.concatenate((np.zeros((1, 3)),
                                                     rand_colors)))
    return axis.imshow(im, cmap=rcmap)


def show_multiple_images(*images, axes=None, image_type='raw'):
    """Returns a figure with subplots containing multiple images.

    Parameters
    ----------
    images : np.ndarray of int, shape (M, N)
        The input images to be displayed.
    axes: matplotlib.AxesImage, optional
        Whether to pass in multiple axes. Must be equal to the number of
        input images.
    image_type : string, optional
        Displays the images with different colormaps. Set to display
        'raw' by default. Other options that are accepted
        are 'grey' and 'magma', or 'rand'.

    Returns
    -------
    fig : plt.Figure
        The image shown.
    """
    number_of_im = len(images)
    figure = plt.figure()
    for i in range(number_of_im):
        ax = (figure.add_subplot(1, number_of_im, i+1) if axes is None
              else axes[i])
        if image_type == 'grey' or image_type == 'gray':
            imshow_grey(images[i], axis=ax)
        elif image_type == 'magma':
            imshow_magma(images[i], axis=ax)
        elif image_type == 'rand':
            imshow_rand(images[i], axis=ax)
        elif image_type == 'raw':
            ax.imshow(images[i])
        else:
            print("not a valid image type.")
            return None
        ax.set_title(f'Image number {i+1} with a {image_type} colormap.')
    return ax


def draw_seg(seg, im):
    """Return a segmentation map matching the original image color.

    Parameters
    ----------
    seg : np.ndarray of int, shape (M, N, ...)
        The segmentation to be displayed
    im : np.ndarray, shape (M, N, ..., C)
        The image corresponding to the segmentation.

    Returns
    -------
    out : np.ndarray, same shape and type as `im`.
        An image where each segment has uniform color.

    Examples
    --------
    >>> a = np.array([[1, 1, 2, 2],
    ...               [1, 2, 2, 3],
    ...               [2, 2, 3, 3]])
    >>> g = np.array([[0.5, 0.2, 1.0, 0.9],
    ...               [0.2, 0.8, 0.9, 0.6],
    ...               [0.9, 0.9, 0.4, 0.5]])
    >>> draw_seg(a, g)
    array([[ 0.3,  0.3,  0.9,  0.9],
           [ 0.3,  0.9,  0.9,  0.5],
           [ 0.9,  0.9,  0.5,  0.5]])
    """
    out = np.zeros_like(im)
    labels = np.unique(seg)
    if (seg==0).any():
        labels = labels[1:]
    for u in labels:
        mask = (seg == u).nonzero()
        color = im[mask].mean(axis=0)
        out[mask] = color
    return out


def display_3d_segmentations(segs, image=None, probability_map=None, axis=0,
                             z=None, fignum=None):
    """Show slices of multiple 3D segmentations.

    Parameters
    ----------
    segs : list or tuple of np.ndarray of int, shape (M, N, P)
        The segmentations to be examined.
    image : np.ndarray, shape (M, N, P[, 3]), optional
        The image corresponding to the segmentations.
    probability_map : np.ndarray, shape (M, N, P), optional
        The segment boundary probability map.
    axis : int in {0, 1, 2}, optional
        The axis along which to show a slice of the segmentation.
    z : int in [0, `(M, N, P)[axis]`), optional
        The slice to display. Defaults to the middle slice.
    fignum : int, optional
        Which figure number to use. Uses the default (new figure) if none is
        provided.

    Returns
    -------
    fig : plt.Figure
        The figure handle.
    """
    numplots = len(segs)
    if image is not None:
        numplots += 1
    if probability_map is not None:
        numplots += 1
    candidate_plot_arrangements = list(it.combinations_with_replacement(
                                       range(1, 5), 2))
    # get the smallest plot arrangement that can display the number of
    # segmentations we want
    plot_arrangement = [(i, j) for i, j in candidate_plot_arrangements
                        if i * j >= numplots][0]
    fig = plt.figure(fignum)
    current_subplot = 1
    if image is not None:
        plt.subplot(*plot_arrangement + (current_subplot,))
        imshow_grey(np.rollaxis(image, axis)[z])
        current_subplot += 1
    if probability_map is not None:
        plt.subplot(*plot_arrangement + (current_subplot,))
        imshow_magma(np.rollaxis(probability_map, axis)[z])
        current_subplot += 1
    for i, j in enumerate(range(current_subplot, numplots + 1)):
        plt.subplot(*pl
ot_arrangement + (j,))
        imshow_rand(np.rollaxis(segs[i], axis)[z])
    return fig


def plot_vi(g, history, gt, fig=None):
    """Plot the VI from segmentations based on Rag and sequence of merges.

    Parameters
    ----------
    g : agglo.Rag object
        The region adjacency graph.

    history : list of tuples
        The merge history of the RAG.

    gt : np.ndarray

        The ground truth corresponding to the RAG.

    fig : plt.Figure, optional
        Use this figure for plotting. If not provided, a new figure is created.

    Returns
    -------
    None
    """
    v = []
    n = []
    seg = g.get_segmentation()
    for i in history:
        seg[seg==i[1]] = i[0]
        v.append(evaluate.vi(seg, gt))
        n.append(len(np.unique(seg)-1))
    if fig is None:
        fig = plt.figure()
    plt.plot(n, v, figure = fig)
    plt.xlabel('Number of segments', figure = fig)
    plt.ylabel('vi', figure = fig)


def plot_vi_breakdown_panel(px, h, title, xlab, ylab, hlines, scatter_size,
                            **kwargs):
    """Plot a single panel (over or undersegmentation) of VI breakdown plot.

    Parameters
    ----------
    px : np.ndarray of float, shape (N,)
        The probability (size) of each segment.
    h : np.ndarray of float, shape (N,)
        The conditional entropy of that segment.
    title, xlab, ylab : string
        Parameters for `matplotlib.plt.plot`.
    hlines : iterable of float
        Plot hyperbolic lines of same VI contribution. For each value `v` in
        `hlines`, draw the line `h = v/px`.
    scatter_size : int, optional
    **kwargs : dict
        Additional keyword arguments for `matplotlib.pyplot.plot`.

    Returns
    -------
    None
    """
    x = np.arange(max(min(px),1e-10), max(px), (max(px)-min(px))/1