"""
Functions for plotting ants images
"""


__all__ = [
    "plot_directory"
]

import fnmatch
import math
import os
import warnings

from matplotlib import gridspec
import matplotlib.pyplot as plt
import matplotlib.patheffects as path_effects
import matplotlib.lines as mlines
import matplotlib.patches as patches
import matplotlib.mlab as mlab
import matplotlib.animation as animation
from mpl_toolkits.axes_grid1.inset_locator import inset_axes


import numpy as np

import ants


def plot_directory(
    directory,
    recursive=False,
    regex="*",
    save_prefix="",
    save_suffix="",
    axis=None,
    **kwargs
):
    """
    Create and save an ANTsPy plot for every image matching a given regular
    expression in a directory, optionally recursively. This is a good function
    for quick visualize exploration of all of images in a directory

    ANTsR function: N/A

    Arguments
    ---------
    directory : string
        directory in which to search for images and plot them

    recursive : boolean
        If true, this function will search through all directories under
        the given directory recursively to make plots.
        If false, this function will only create plots for images in the
        given directory

    regex : string
        regular expression used to filter out certain filenames or suffixes

    save_prefix : string
        sub-string that will be appended to the beginning of all saved plot filenames.
        Default is to add nothing.

    save_suffix : string
        sub-string that will be appended to the end of all saved plot filenames.
        Default is add nothing.

    kwargs : keyword arguments
        any additional arguments to pass onto the `ants.plot` function.
        e.g. overlay, alpha, cmap, etc. See `ants.plot` for more options.

    Example
    -------
    >>> import ants
    >>> ants.plot_directory(directory='~/desktop/testdir',
                            recursive=False, regex='*')
    """

    def has_acceptable_suffix(fname):
        suffixes = {".nii.gz"}
        return sum([fname.endswith(sx) for sx in suffixes]) > 0

    if directory.startswith("~"):
        directory = os.path.expanduser(directory)

    if not os.path.isdir(directory):
        raise ValueError("directory %s does not exist!" % directory)

    for root, dirnames, fnames in os.walk(directory):
        for fname in fnames:
            if fnmatch.fnmatch(fname, regex) and has_acceptable_suffix(fname):
                load_fname = os.path.join(root, fname)
                fname = fname.replace(".".join(fname.split(".")[1:]), "png")
                fname = fname.replace(".png", "%s.png" % save_suffix)
                fname = "%s%s" % (save_prefix, fname)
                save_fname = os.path.join(root, fname)
                img = ants.image_read(load_fname)

                if axis is None:
                    axis_range = [i for i in range(img.dimension)]
                else:
                    axis_range = axis if isinstance(axis, (list, tuple)) else [axis]

                if img.dimension > 2:
                    for axis_idx in axis_range:
                        filename = save_fname.replace(".png", "_axis%i.png" % axis_idx)
                        ncol = int(math.sqrt(img.shape[axis_idx]))
                        ants.plot(
                            img,
                            axis=axis_idx,
                            nslices=img.shape[axis_idx],
                            ncol=ncol,
                            filename=filename,
                            **kwargs
                        )
                else:
                    filename = save_fname
                    ants.plot(img, filename=filename, **kwargs)