import itertools

import numpy as np
import seaborn as sns

from dosma import defaults

import matplotlib.pyplot as plt
from matplotlib.lines import Line2D

__all__ = ["downsample_slice", "write_regions"]


def downsample_slice(img_array, ds_factor, is_mask=False):
    """
    Takes in a 3D array and then downsamples in the z-direction by a
    user-specified downsampling factor.

    Args:
        img_array (np.ndarray): 3D numpy array for now (xres x yres x zres)
        ds_factor (int): Downsampling factor
        is_mask (:obj:`bool`, optional): If ``True``, ``img_array`` is a mask and will be binarized
            after downsampling. Defaults to `False`.

    Returns:
        np.ndarray: 3D numpy array of dimensions (xres x yres x zres//ds_factor)

    Examples:
        >>> input_image  = numpy.random.rand(4,4,4)
        >>> input_mask   = (a > 0.5) * 1.0
        >>> output_image = downsample_slice(input_mask, ds_factor = 2, is_mask = False)
        >>> output_mask  = downsample_slice(input_mask, ds_factor = 2, is_mask = True)
    """

    img_array = np.transpose(img_array, (2, 0, 1))
    L = list(img_array)

    def grouper(iterable, n):
        args = [iter(iterable)] * n
        return itertools.zip_longest(fillvalue=0, *args)

    final = np.array([sum(x) for x in grouper(L, ds_factor)])
    final = np.transpose(final, (1, 2, 0))

    # Binarize if it is a mask.
    if is_mask is True:
        final = (final >= 1) * 1

    return final


def write_regions(file_path, arr, plt_dict=None):
    """Write 2D array to region image where colors correspond to the region.

    All finite values should be >= 1.
    nan/inf value are ignored - written as white.

    Args:
        file_path (str): File path to save image.
        arr (np.ndarray): The 2D numpy array to convert to region image.
            Unique non-zero values correspond to different regions.
            Values that are `0` or `np.nan` will be written as white pixels.
        plt_dict (:obj:`dict`, optional): Dictionary of values to use when plotting with
            ``matplotlib.pyplot``. Keys are strings like `xlabel`, `ylabel`, etc.
            Use Key `labels` to specify a mapping from unique non-zero values in the array
            to names for the legend.
    """

    if len(arr.shape) != 2:
        raise ValueError("`arr` must be a 2D numpy array")

    unique_vals = np.unique(arr.flatten())
    if 0 in unique_vals:
        raise ValueError("All finite values in `arr` must be >=1")

    unique_vals = unique_vals[np.isfinite(unique_vals)]
    num_unique_vals = len(unique_vals)

    plt_dict_int = {"xlabel": "", "ylabel": "", "title": "", "labels": None}
    if plt_dict:
        plt_dict_int.update(plt_dict)
    plt_dict = plt_dict_int

    labels = plt_dict["labels"]
    if labels is None:
        labels = list(unique_vals)

    if len(labels) != num_unique_vals:
        raise ValueError(
            "len(labels) != num_unique_vals - %d != %d" % (len(labels), num_unique_vals)
        )

    cpal = sns.color_palette("pastel", num_unique_vals)

    arr_c = np.array(arr)
    arr_c = np.nan_to_num(arr_c)
    arr_c[arr_c > np.max(unique_vals)] = 0
    arr_rgb = np.ones([arr_c.shape[0], arr_c.shape[1], 3])

    plt.figure()
    plt.clf()
    custom_lines = []
    for i in range(num_unique_vals):
        unique_val = unique_vals[i]
        i0, i1 = np.where(arr_c == unique_val)
        arr_rgb[i0, i1, ...] = np.asarray(cpal[i])

        custom_lines.append(
            Line2D([], [], color=cpal[i], marker="o", linestyle="None", markersize=5)
        )

    plt.xlabel(plt_dict["xlabel"])
    plt.ylabel(plt_dict["ylabel"])
    plt.title(plt_dict["title"])

    lgd = plt.legend(
        custom_lines,
        labels,
        loc="upper center",
        bbox_to_anchor=(0.5, -defaults.DEFAULT_TEXT_SPACING),
        fancybox=True,
        shadow=True,
        ncol=3,
    )
    plt.imshow(arr_rgb)

    plt.savefig(file_path, bbox_extra_artists=(lgd,), bbox_inches="tight")