import numpy as np

from dipy.testing.decorators import warning_for_keywords
from dipy.utils.optpkg import optional_package

matplotlib, has_mpl, setup_module = optional_package("matplotlib")
plt, _, _ = optional_package("matplotlib.pyplot")


def _tile_plot(imgs, titles, **kwargs):
    """
    Helper function
    """
    # Create a new figure and plot the three images
    fig, ax = plt.subplots(1, len(imgs))
    for ii, a in enumerate(ax):
        a.set_axis_off()
        a.imshow(imgs[ii], **kwargs)
        a.set_title(titles[ii])

    return fig


def simple_plot(file_name, title, x, y, xlabel, ylabel):
    """Saves the simple plot with given x and y values

    Parameters
    ----------
    file_name : string
        file name for saving the plot
    title : string
        title of the plot
    x : integer list
        x-axis values to be plotted
    y : integer list
        y-axis values to be plotted
    xlabel : string
        label for x-axis
    ylable : string
        label for y-axis

    """

    plt.plot(x, y)
    axes = plt.gca()
    axes.set_ylim([0, 4])
    plt.xlabel(xlabel)
    plt.ylabel(ylabel)
    plt.title(title)

    plt.savefig(file_name)
    plt.clf()


@warning_for_keywords()
def overlay_images(
    img0, img1, *, title0="", title_mid="", title1="", fname=None, **fig_kwargs
):
    r"""Plot two images one on top of the other using red and green channels.

    Creates a figure containing three images: the first image to the left
    plotted on the red channel of a color image, the second to the right
    plotted on the green channel of a color image and the two given images on
    top of each other using the red channel for the first image and the green
    channel for the second one. It is assumed that both images have the same
    shape. The intended use of this function is to visually assess the quality
    of a registration result.

    Parameters
    ----------
    img0 : array, shape(R, C)
        the image to be plotted on the red channel, to the left of the figure
    img1 : array, shape(R, C)
        the image to be plotted on the green channel, to the right of the
        figure
    title0 : string, optional
        the title to be written on top of the image to the left. By default, no
        title is displayed.
    title_mid : string, optional
        the title to be written on top of the middle image. By default, no
        title is displayed.
    title1 : string, optional
        the title to be written on top of the image to the right. By default,
        no title is displayed.
    fname : string, optional
        the file name to write the resulting figure. If None (default), the
        image is not saved.
    fig_kwargs : dict
        Extra parameters for saving figure, e.g. `dpi=300`.
    """
    # Normalize the input images to [0,255]
    img0 = 255 * ((img0 - img0.min()) / (img0.max() - img0.min()))
    img1 = 255 * ((img1 - img1.min()) / (img1.max() - img1.min()))

    # Create the color images
    img0_red = np.zeros(shape=img0.shape + (3,), dtype=np.uint8)
    img1_green = np.zeros(shape=img0.shape + (3,), dtype=np.uint8)
    overlay = np.zeros(shape=img0.shape + (3,), dtype=np.uint8)

    # Copy the normalized intensities into the appropriate channels of the
    # color images
    img0_red[..., 0] = img0
    img1_green[..., 1] = img1
    overlay[..., 0] = img0
    overlay[..., 1] = img1

    fig = _tile_plot([img0_red, overlay, img1_green], [title0, title_mid, title1])

    # If a file name was given, save the figure
    if fname is not None:
        fig.savefig(fname, bbox_inches="tight", **fig_kwargs)

    return fig


def draw_lattice_2d(nrows, ncols, delta):
    r"""Create a regular lattice of nrows x ncols squares.

    Creates an image (2D array) of a regular lattice of nrows x ncols squares.
    The size of each square is delta x delta pixels (not counting the
    separation lines). The lines are one pixel width.

    Parameters
    ----------
    nrows : int
        the number of squares to be drawn vertically
    ncols : int
        the number of squares to be drawn horizontally
    delta : int
        the size of each square of the grid. Each square is delta x delta
        pixels

    Returns
    -------
    lattice : array, shape (R, C)
        the image (2D array) of the segular lattice. The shape (R, C) of the
        array is given by
        R = 1 + (delta + 1) * nrows
        C = 1 + (delta + 1) * ncols
    """
    lattice = np.ndarray(
        (1 + (delta + 1) * nrows, 1 + (delta + 1) * ncols), dtype=np.float64
    )

    # Fill the lattice with "white"
    lattice[...] = 127

    # Draw the horizontal lines in "black"
    for i in range(nrows + 1):
        lattice[i * (delta + 1), :] = 0

    # Draw the vertical lines in "black"
    for j in range(ncols + 1):
        lattice[:, j * (delta + 1)] = 0

    return lattice


@warning_for_keywords()
def plot_2d_diffeomorphic_map(
    mapping,
    *,
    delta=10,
    fname=None,
    direct_grid_shape=None,
    direct_grid2world=-1,
    inverse_grid_shape=None,
    inverse_grid2world=-1,
    show_figure=True,
    **fig_kwargs,
):
    r"""Draw the effect of warping a regular lattice by a diffeomorphic map.

    Draws a diffeomorphic map by showing the effect of the deformation on a
    regular grid. The resulting figure contains two images: the direct
    transformation is plotted to the left, and the inverse transformation is
    plotted to the right.

    Parameters
    ----------
    mapping : DiffeomorphicMap object
        the diffeomorphic map to be drawn
    delta : int, optional
        the size (in pixels) of the squares of the regular lattice to be used
        to plot the warping effects. Each square will be delta x delta pixels.
        By default, the size will be 10 pixels.
    fname : string, optional
        the name of the file the figure will be written to. If None (default),
        the figure will not be saved to disk.
    direct_grid_shape : tuple, shape (2,), optional
        the shape of the grid image after being deformed by the direct
        transformation. By default, the shape of the deformed grid is the
        same as the grid of the displacement field, which is by default
        equal to the shape of the fixed image. In other words, the resulting
        deformed grid (deformed by the direct transformation) will normally
        have the same shape as the fixed image.
    direct_grid2world : array, shape (3, 3), optional
        the affine transformation mapping the direct grid's coordinates to
        physical space. By default, this transformation will correspond to
        the image-to-world transformation corresponding to the default
        direct_grid_shape (in general, if users specify a direct_grid_shape,
        they should also specify direct_grid2world).
    inverse_grid_shape : tuple, shape (2,), optional
        the shape of the grid image after being deformed by the inverse
        transformation. By default, the shape of the deformed grid under the
        inverse transform is the same as the image used as "moving" when
        the diffeomorphic map was generated by a registration algorithm
        (so it corresponds to the effect of warping the static image towards
        the moving).
    inverse_grid2world : array, shape (3, 3), optional
        the affine transformation mapping inverse grid's coordinates to
        physical space. By default, this transformation will correspond to
        the image-to-world transformation corresponding to the default
        inverse_grid_shape (in general, if users specify an inverse_grid_shape,
        they should also specify inverse_grid2world).
    show_figure : bool, optional
        if True (default), the deformed grids will be plotted using matplotlib,
        else the grids are just returned
    fig_kwargs : dict
        Extra parameters for saving figure, e.g. `dpi=300`.

    Returns
    -------
    warped_forward : array
        Image with the grid showing the effect of transforming the moving image to
        the static image.  The shape will be `direct_grid_shape` if specified,
        otherwise the shape of the static image.
    warped_backward : array
        Image with the grid showing the effect of transforming the static image to
        the moving image.  Shape will be `inverse_grid_shape` if specified,
        otherwise the shape of the moving image.

    Notes
    -----
    The default value for the affine transformation is "-1" to handle the case
    in which the user provides "None" as input meaning "identity". If we used
    None as default, we wouldn't know if the user specifically wants to use
    the identity (specifically passing None) or if it was left unspecified,
    meaning to use the appropriate default matrix.
    """
    if mapping.is_inverse:
        # By default, direct_grid_shape is the codomain grid
        if direct_grid_shape is None:
            direct_grid_shape = mapping.codomain_shape
        if direct_grid2world == -1:
            direct_grid2world = mapping.codomain_grid2world

        # By default, the inverse grid is the domain grid
        if inverse_grid_shape is None:
            inverse_grid_shape = mapping.domain_shape
        if inverse_grid2world == -1:
            inverse_grid2world = mapping.domain_grid2world
    else:
        # Now by default, direct_grid_shape is the mapping's input grid
        if direct_grid_shape is None:
            direct_grid_shape = mapping.domain_shape
        if direct_grid2world == -1:
            direct_grid2world = mapping.domain_grid2world

        # By default, the output grid is the mapping's domain grid
        if inverse_grid_shape is None:
            inverse_grid_shape = mapping.codomain_shape
        if inverse_grid2world == -1:
            inverse_grid2world = mapping.codomain_grid2world

    # The world-to-image (image = drawn lattice on the output grid)
    # transformation is the inverse of the output affine
    world_to_image = None
    if inverse_grid2world is not None:
        world_to_image = np.linalg.inv(inverse_grid2world)

    # Draw the squares on the output grid
    lattice_out = draw_lattice_2d(
        (inverse_grid_shape[0] + delta) // (delta + 1),
        (inverse_grid_shape[1] + delta) // (delta + 1),
        delta,
    )
    lattice_out = lattice_out[0 : inverse_grid_shape[0], 0 : inverse_grid_shape[1]]

    # Warp in the forward direction (sampling it on the input grid)
    warped_forward = mapping.transform(
        lattice_out,
        interpolation="linear",
        image_world2grid=world_to_image,
        out_shape=direct_grid_shape,
        out_grid2world=direct_grid2world,
    )

    # Now, the world-to-image (image = drawn lattice on the input grid)
    # transformation is the inverse of the input affine
    world_to_image = None
    if direct_grid2world is not None:
        world_to_image = np.linalg.inv(direct_grid2world)

    # Draw the squares on the input grid
    lattice_in = draw_lattice_2d(
        (direct_grid_shape[0] + delta) // (delta + 1),
        (direct_grid_shape[1] + delta) // (delta + 1),
        delta,
    )
    lattice_in = lattice_in[0 : direct_grid_shape[0], 0 : direct_grid_shape[1]]

    # Warp in the backward direction (sampling it on the output grid)
    warped_backward = mapping.transform_inverse(
        lattice_in,
        interpolation="linear",
        image_world2grid=world_to_image,
        out_shape=inverse_grid_shape,
        out_grid2world=inverse_grid2world,
    )

    # Now plot the grids
    if show_figure:
        plt.figure()

        plt.subplot(1, 2, 1).set_axis_off()
        plt.imshow(warped_forward, cmap=plt.cm.gray)
        plt.title("Direct transform")

        plt.subplot(1, 2, 2).set_axis_off()
        plt.imshow(warped_backward, cmap=plt.cm.gray)
        plt.title("Inverse transform")

        # Finally, save the figure to disk
        if fname is not None:
            plt.savefig(fname, bbox_inches="tight", **fig_kwargs)

    # Return the deformed grids
    return warped_forward, warped_backward


@warning_for_keywords()
def plot_slices(V, *, slice_indices=None, fname=None, **fig_kwargs):
    r"""Plot 3 slices from the given volume: 1 sagittal, 1 coronal and 1 axial

    Creates a figure showing the axial, coronal and sagittal slices at the
    requested positions of the given volume. The requested slices are specified
    by slice_indices.

    Parameters
    ----------
    V : array, shape (S, R, C)
        the 3D volume to extract the slices from
    slice_indices : array, shape (3,), optional
        the indices of the sagittal (slice_indices[0]), coronal
        (slice_indices[1])
        and axial (slice_indices[2]) slices to be displayed. If None, the
        middle slices along each direction are displayed.
    fname : string, optional
        the name of the file to save the figure to. If None (default), the
        figure is not saved to disk.
    fig_kwargs : dict
        Extra parameters for saving figure, e.g. `dpi=300`.
    """
    if slice_indices is None:
        slice_indices = np.array(V.shape) // 2

    # Normalize the intensities to [0, 255]
    V = np.asarray(V, dtype=np.float64)
    V = 255 * (V - V.min()) / (V.max() - V.min())

    # Extract the middle slices
    axial = np.asarray(V[:, :, slice_indices[2]]).astype(np.uint8).T
    coronal = np.asarray(V[:, slice_indices[1], :]).astype(np.uint8).T
    sagittal = np.asarray(V[slice_indices[0], :, :]).astype(np.uint8).T

    fig = _tile_plot(
        [axial, coronal, sagittal],
        ["Axial", "Coronal", "Sagittal"],
        cmap=plt.cm.gray,
        origin="lower",
    )

    # Save the figure if requested
    if fname is not None:
        fig.savefig(fname, bbox_inches="tight", **fig_kwargs)

    return fig


@warning_for_keywords()
def overlay_slices(
    L,
    R,
    *,
    slice_index=None,
    slice_type=1,
    ltitle="Left",
    rtitle="Right",
    fname=None,
    **fig_kwargs,
):
    r"""Plot three overlaid slices from the given volumes.

    Creates a figure containing three images: the gray scale k-th slice of
    the first volume (L) to the left, where k=slice_index, the k-th slice of
    the second volume (R) to the right and the k-th slices of the two given
    images on top of each other using the red channel for the first volume and
    the green channel for the second one. It is assumed that both volumes have
    the same shape. The intended use of this function is to visually assess the
    quality of a registration result.

    Parameters
    ----------
    L : array, shape (S, R, C)
        the first volume to extract the slice from plotted to the left
    R : array, shape (S, R, C)
        the second volume to extract the slice from, plotted to the right
    slice_index : int, optional
        the index of the slices (along the axis given by slice_type) to be
        overlaid. If None, the slice along the specified axis is used
    slice_type : int, optional
        the type of slice to be extracted:
        0=sagittal, 1=coronal (default), 2=axial.
    ltitle : string, optional
        the string to be written as the title of the left image. By default,
        no title is displayed.
    rtitle : string, optional
        the string to be written as the title of the right image. By default,
        no title is displayed.
    fname : string, optional
        the name of the file to write the image to. If None (default), the
        figure is not saved to disk.
    fig_kwargs: extra parameters for saving figure, e.g. `dpi=300`.
    """

    # Normalize the intensities to [0,255]
    sh = L.shape
    L = np.asarray(L, dtype=np.float64)
    R = np.asarray(R, dtype=np.float64)
    L = 255 * (L - L.min()) / (L.max() - L.min())
    R = 255 * (R - R.min()) / (R.max() - R.min())

    # Create the color image to draw the overlapped slices into, and extract
    # the slices (note the transpositions)
    if slice_type == 0:
        if slice_index is None:
            slice_index = sh[0] // 2
        colorImage = np.zeros(shape=(sh[2], sh[1], 3), dtype=np.uint8)
        ll = np.asarray(L[slice_index, :, :]).astype(np.uint8).T
        rr = np.asarray(R[slice_index, :, :]).astype(np.uint8).T
    elif slice_type == 1:
        if slice_index is None:
            slice_index = sh[1] // 2
        colorImage = np.zeros(shape=(sh[2], sh[0], 3), dtype=np.uint8)
        ll = np.asarray(L[:, slice_index, :]).astype(np.uint8).T
        rr = np.asarray(R[:, slice_index, :]).astype(np.uint8).T
    elif slice_type == 2:
        if slice_index is None:
            slice_index = sh[2] // 2
        colorImage = np.zeros(shape=(sh[1], sh[0], 3), dtype=np.uint8)
        ll = np.asarray(L[:, :, slice_index]).astype(np.uint8).T
        rr = np.asarray(R[:, :, slice_index]).astype(np.uint8).T
    else:
        print("Slice type must be 0, 1 or 2.")
        return

    # Draw the intensity images to the appropriate channels of the color image
    # The "(ll > ll[0, 0])" condition is just an attempt to eliminate the
    # background when its intensity is not exactly zero (the [0,0] corner is
    # usually background)
    colorImage[..., 0] = ll * (ll > ll[0, 0])
    colorImage[..., 1] = rr * (rr > rr[0, 0])

    fig = _tile_plot(
        [ll, colorImage, rr],
        [ltitle, "Overlay", rtitle],
        cmap=plt.cm.gray,
        origin="lower",
    )

    # Save the figure to disk, if requested
    if fname is not None:
        fig.savefig(fname, bbox_inches="tight", **fig_kwargs)

    return fig
