Plot an orthographic view of a 3D image

    TODO:
        - add colorbar option
    
    ANTsR function: N/A

    Arguments
    ---------
    image : ANTsImage
        image to plot

    overlay : ANTsImage
        image to overlay on base image

    slices : list or tuple of 3 integers
        slice indices along each axis to plot
        This can be absolute array indices (e.g. (80,100,120)), or
        this can be relative array indices (e.g. (0.4,0.5,0.6)).
        The default is to take the middle slice along each axis.
    
    xyz : list or tuple of 3 integers
        if given, solid lines will be drawn to converge at this coordinate.
        This is useful for pinpointing a specific location in the image.

    flat : boolean
        if true, the ortho image will be plot in one row
        if false, the ortho image will be a 2x2 grid with the bottom
            left corner blank

    cmap : string
        colormap to use for base image. See matplotlib.

    overlay_cmap : string
        colormap to use for overlay images, if applicable. See matplotlib.

    overlay_alpha : float
        level of transparency for any overlays. Smaller value means 
        the overlay is more transparent. See matplotlib.

    axis : integer
        which axis to plot along if image is 3D

    black_bg : boolean
        if True, the background of the image(s) will be black.
        if False, the background of the image(s) will be determined by the
            values `bg_thresh_quant` and `bg_val_quant`.

    bg_thresh_quant : float 
        if white_bg=True, the background will be determined by thresholding
        the image at the `bg_thresh` quantile value and setting the background
        intensity to the `bg_val` quantile value. 
        This value should be in [0, 1] - somewhere around 0.01 is recommended.
            - equal to 1 will threshold the entire image
            - equal to 0 will threshold none of the image

    bg_val_quant : float
        if white_bg=True, the background will be determined by thresholding
        the image at the `bg_thresh` quantile value and setting the background
        intensity to the `bg_val` quantile value.
        This value should be in [0, 1] 
            - equal to 1 is pure white
            - equal to 0 is pure black
            - somewhere in between is gray

    domain_image_map : ANTsImage
        this input ANTsImage or list of ANTsImage types contains a reference image
        `domain_image` and optional reference mapping named `domainMap`.
        If supplied, the image(s) to be plotted will be mapped to the domain
        image space before plotting - useful for non-standard image orientations.

    crop : boolean
        if true, the image(s) will be cropped to their bounding boxes, resulting
        in a potentially smaller image size.
        if false, the image(s) will not be cropped

    scale : boolean or 2-tuple
        if true, nothing will happen to intensities of image(s) and overlay(s)
        if false, dynamic range will be maximized when visualizing overlays
        if 2-tuple, the image will be dynamically scaled between these quantiles

    title : string 
        add a title to the plot

    filename : string
        if given, the resulting image will be saved to this file

    dpi : integer
        determines resolution of image if saved to file. Higher values
        result in higher resolution images, but at a cost of having a 
        larger file size

    Example
    -------
    >>> import ants
    >>> mni = ants.image_read(ants.get_data('mni'))
    >>> ants.plot_ortho(mni, xyz=(100,100,100))
    >>> mni2 = mni.threshold_image(7000, mni.max())
    >>> ants.plot_ortho(mni, overlay=mni2)
    >>> ants.plot_ortho(mni, overlay=mni2, flat=True)
    >>> ants.plot_ortho(mni, overlay=mni2, xyz=(110,110,110), xyz_lines=False,
                        text='Lines Turned Off', textfontsize=22)
    >>> ants.plot_ortho(mni, mni2, xyz=(120,100,100),
                        text=' Example \nOrtho Text', textfontsize=26, 
                        title='Example Ortho Title', titlefontsize=26)

def plot_ortho(image, overlay=None, reorient=True, blend=False,
    # xyz arguments
    xyz=None, xyz_lines=True, xyz_color='red', xyz_alpha=0.6, xyz_linewidth=2, xyz_pad=5,
    orient_labels=True,
    # base image arguments
    alpha=1, cmap='Greys_r', 
    # overlay arguments
    overlay_cmap='jet', overlay_alpha=0.9,   
    # background arguments
    black_bg=True, bg_thresh_quant=0.01, bg_val_quant=0.99, 
    # scale/crop/domain arguments
    crop=False, scale=False, domain_image_map=None, 
    # title arguments
    title=None, titlefontsize=24, title_dx=0, title_dy=0, 
    # 4th panel text arguemnts
    text=None, textfontsize=24, textfontcolor='white', text_dx=0, text_dy=0,
    # save & size arguments
    filename=None, dpi=500, figsize=1., flat=False, transparent=True):
    """
    Plot an orthographic view of a 3D image

    TODO:
        - add colorbar option
    
    ANTsR function: N/A

    Arguments
    ---------
    image : ANTsImage
        image to plot

    overlay : ANTsImage
        image to overlay on base image

    slices : list or tuple of 3 integers
        slice indices along each axis to plot
        This can be absolute array indices (e.g. (80,100,120)), or
        this can be relative array indices (e.g. (0.4,0.5,0.6)).
        The default is to take the middle slice along each axis.
    
    xyz : list or tuple of 3 integers
        if given, solid lines will be drawn to converge at this coordinate.
        This is useful for pinpointing a specific location in the image.

    flat : boolean
        if true, the ortho image will be plot in one row
        if false, the ortho image will be a 2x2 grid with the bottom
            left corner blank

    cmap : string
        colormap to use for base image. See matplotlib.

    overlay_cmap : string
        colormap to use for overlay images, if applicable. See matplotlib.

    overlay_alpha : float
        level of transparency for any overlays. Smaller value means 
        the overlay is more transparent. See matplotlib.

    axis : integer
        which axis to plot along if image is 3D

    black_bg : boolean
        if True, the background of the image(s) will be black.
        if False, the background of the image(s) will be determined by the
            values `bg_thresh_quant` and `bg_val_quant`.

    bg_thresh_quant : float 
        if white_bg=True, the background will be determined by thresholding
        the image at the `bg_thresh` quantile value and setting the background
        intensity to the `bg_val` quantile value. 
        This value should be in [0, 1] - somewhere around 0.01 is recommended.
            - equal to 1 will threshold the entire image
            - equal to 0 will threshold none of the image

    bg_val_quant : float
        if white_bg=True, the background will be determined by thresholding
        the image at the `bg_thresh` quantile value and setting the background
        intensity to the `bg_val` quantile value.
        This value should be in [0, 1] 
            - equal to 1 is pure white
            - equal to 0 is pure black
            - somewhere in between is gray

    domain_image_map : ANTsImage
        this input ANTsImage or list of ANTsImage types contains a reference image
        `domain_image` and optional reference mapping named `domainMap`.
        If supplied, the image(s) to be plotted will be mapped to the domain
        image space before plotting - useful for non-standard image orientations.

    crop : boolean
        if true, the image(s) will be cropped to their bounding boxes, resulting
        in a potentially smaller image size.
        if false, the image(s) will not be cropped

    scale : boolean or 2-tuple
        if true, nothing will happen to intensities of image(s) and overlay(s)
        if false, dynamic range will be maximized when visualizing overlays
        if 2-tuple, the image will be dynamically scaled between these quantiles

    title : string 
        add a title to the plot

    filename : string
        if given, the resulting image will be saved to this file

    dpi : integer
        determines resolution of image if saved to file. Higher values
        result in higher resolution images, but at a cost of having a 
        larger file size

    Example
    -------
    >>> import ants
    >>> mni = ants.image_read(ants.get_data('mni'))
    >>> ants.plot_ortho(mni, xyz=(100,100,100))
    >>> mni2 = mni.threshold_image(7000, mni.max())
    >>> ants.plot_ortho(mni, overlay=mni2)
    >>> ants.plot_ortho(mni, overlay=mni2, flat=True)
    >>> ants.plot_ortho(mni, overlay=mni2, xyz=(110,110,110), xyz_lines=False,
                        text='Lines Turned Off', textfontsize=22)
    >>> ants.plot_ortho(mni, mni2, xyz=(120,100,100),
                        text=' Example \nOrtho Text', textfontsize=26, 
                        title='Example Ortho Title', titlefontsize=26)
    """
    def mirror_matrix(x):
        return x[::-1,:]
    def rotate270_matrix(x):
        return mirror_matrix(x.T)
    def reorient_slice(x, axis):
        return rotate270_matrix(x)

    # need this hack because of a weird NaN warning from matplotlib with overlays
    warnings.simplefilter('ignore')

    # handle `image` argument
    if isinstance(image, str):
        image = iio2.image_read(image)
    if not isinstance(image, iio.ANTsImage):
        raise ValueError('image argument must be an ANTsImage')
    if image.dimension != 3:
        raise ValueError('Input image must have 3 dimensions!')

    # handle `overlay` argument 
    if overlay is not None:
        if isinstance(overlay, str):
            overlay = iio2.image_read(overlay)
        if not isinstance(overlay, iio.ANTsImage):
            raise ValueError('overlay argument must be an ANTsImage')
        if overlay.dimension != 3:
            raise ValueError('Overlay image must have 3 dimensions!')

        if not iio.image_physical_space_consistency(image, overlay):
            overlay = reg.resample_image_to_target(overlay, image, interp_type='linear')

    if blend:
        if alpha == 1:
            alpha = 0.5
        image = image*alpha + overlay*(1-alpha)
        overlay = None
        alpha = 1.

    if image.pixeltype not in {'float', 'double'}:
        scale = False # turn off scaling if image is discrete

    # reorient images
    if reorient != False:
        if reorient == True:
            reorient = 'RPI'
        image = image.reorient_image2('RPI')
        if overlay is not None:
            overlay = overlay.reorient_image2('RPI')

    # handle `slices` argument
    if xyz is None:
        xyz = [int(s/2) for s in image.shape]
    for i in range(3):
        if xyz[i] is None:
            xyz[i] = int(image.shape[i]/2)

    # resample image if spacing is very unbalanced
    spacing = [s for i,s in enumerate(image.spacing)]
    if (max(spacing) / min(spacing)) > 3.:
        new_spacing = (1,1,1)
        image = image.resample_image(tuple(new_spacing))
        if overlay is not None:
            overlay = overlay.resample_image(tuple(new_spacing))
        xyz = [int(sl*(sold/snew)) for sl,sold,snew in zip(xyz,spacing,new_spacing)]

    # potentially crop image
    if crop:
        plotmask = image.get_mask(cleanup=0)
        if plotmask.max() == 0:
            plotmask += 1
        image = image.crop_image(plotmask)
        if overlay is not None:
            overlay = overlay.crop_image(plotmask)

    # pad images
    image, lowpad, uppad = image.pad_image(return_padvals=True)
    xyz = [v+l for v,l in zip(xyz,lowpad)]
    if overlay is not None:
        overlay = overlay.pad_image()

    # handle `domain_image_map` argument
    if domain_image_map is not None:
        if isinstance(domain_image_map, iio.ANTsImage):
            tx = tio2.new_ants_transform(precision='float', transform_type='AffineTransform',
                                         dimension=image.dimension)
            image = tio.apply_ants_transform_to_image(tx, image, domain_image_map)
            if overlay is not None:
                overlay = tio.apply_ants_transform_to_image(tx, overlay, 
                                                            domain_image_map, 
                                                            interpolation='linear')
        elif isinstance(domain_image_map, (list, tuple)):
            # expect an image and transformation
            if len(domain_image_map) != 2:
                raise ValueError('domain_image_map list or tuple must have length == 2')
            
            dimg = domain_image_map[0]
            if not isinstance(dimg, iio.ANTsImage):
                raise ValueError('domain_image_map first entry should be ANTsImage')

            tx = domain_image_map[1]
            image = reg.apply_transforms(dimg, image, transform_list=tx)
            if overlay is not None:
                overlay = reg.apply_transforms(dimg, overlay, transform_list=tx,
                                               interpolator='linear')

    ## single-channel images ##
    if image.components == 1:

        # potentially find dynamic range
        if scale == True:
            vmin, vmax = image.quantile((0.05,0.95))
        elif isinstance(scale, (list,tuple)):
            if len(scale) != 2:
                raise ValueError('scale argument must be boolean or list/tuple with two values')
            vmin, vmax = image.quantile(scale)
        else:
            vmin = None
            vmax = None

        if not flat:
            nrow = 2
            ncol = 2
        else:
            nrow = 1
            ncol = 3

        fig = plt.figure(figsize=(9*figsize,9*figsize))
        if title is not None:
            basey = 0.88 if not flat else 0.66
            basex = 0.5
            fig.suptitle(title, fontsize=titlefontsize, x=basex+title_dx, y=basey+title_dy)

        gs = gridspec.GridSpec(nrow, ncol,
                 wspace=0.0, hspace=0.0, 
                 top=1.-0.5/(nrow+1), bottom=0.5/(nrow+1), 
                 left=0.5/(ncol+1), right=1-0.5/(ncol+1))

        # pad image to have isotropic array dimensions
        image = image.numpy()
        if overlay is not None:
            overlay = overlay.numpy()
            overlay[np.abs(overlay) == 0] = np.nan

        yz_slice = reorient_slice(image[xyz[0],:,:],0)
        ax = plt.subplot(gs[0,0])
        ax.imshow(yz_slice, cmap=cmap, vmin=vmin, vmax=vmax)
        if overlay is not None:
            yz_overlay = reorient_slice(overlay[xyz[0],:,:],0)
            ax.imshow(yz_overlay, alpha=overlay_alpha, cmap=overlay_cmap)
        if xyz_lines:
            # add lines
            l = mlines.Line2D([yz_slice.shape[0]-xyz[1],yz_slice.shape[0]-xyz[1]], 
                              [xyz_pad,yz_slice.shape[0]-xyz_pad],
                              color=xyz_color, alpha=xyz_alpha, linewidth=xyz_linewidth)
            ax.add_line(l)
            l = mlines.Line2D([xyz_pad,yz_slice.shape[1]-xyz_pad], 
                              [yz_slice.shape[1]-xyz[2],yz_slice.shape[1]-xyz[2]],
                color=xyz_color, alpha=xyz_alpha, linewidth=xyz_linewidth)
            ax.add_line(l)
        if orient_labels:
            ax.text(0.5,0.98, 'S',
                    horizontalalignment='center',
                    verticalalignment='top',
                    fontsize=20*figsize, color=textfontcolor,
                    transform=ax.transAxes)
            ax.text(0.5,0.02, 'I',
                    horizontalalignment='center',
                    verticalalignment='bottom',
                    fontsize=20*figsize, color=textfontcolor,
                    transform=ax.transAxes)
            ax.text(0.98,0.5, 'A',
                    horizontalalignment='right',
                    verticalalignment='center',
                    fontsize=20*figsize, color=textfontcolor,
                    transform=ax.transAxes)
            ax.text(0.02,0.5, 'P',
                    horizontalalignment='left',
                    verticalalignment='center',
                    fontsize=20*figsize, color=textfontcolor,
                    transform=ax.transAxes)
        ax.axis('off')

        xz_slice = reorient_slice(image[:,xyz[1],:],1)
        ax = plt.subplot(gs[0,1])
        ax.imshow(xz_slice, cmap=cmap, vmin=vmin, vmax=vmax)
        if overlay is not None:
            xz_overlay = reorient_slice(overlay[:,xyz[1],:],1)
            ax.imshow(xz_overlay, alpha=overlay_alpha, cmap=overlay_cmap)
        if xyz_lines:
            # add lines
            l = mlines.Line2D([xz_slice.shape[0]-xyz[0],xz_slice.shape[0]-xyz[0]], 
                              [xyz_pad,xz_slice.shape[0]-xyz_pad],
                               color=xyz_color, alpha=xyz_alpha, linewidth=xyz_linewidth)
            ax.add_line(l)
            l = mlines.Line2D([xyz_pad,xz_slice.shape[1]-xyz_pad], 
                              [xz_slice.shape[1]-xyz[2],xz_slice.shape[1]-xyz[2]],
                              color=xyz_color, alpha=xyz_alpha, linewidth=xyz_linewidth)
            ax.add_line(l)
        if orient_labels:
            ax.text(0.5,0.98, 'S',
                    horizontalalignment='center',
                    verticalalignment='top',
                    fontsize=20*figsize, color=textfontcolor,
                    transform=ax.transAxes)
            ax.text(0.5,0.02, 'I',
                    horizontalalignment='center',
                    verticalalignment='bottom',
                    fontsize=20*figsize, color=textfontcolor,
                    transform=ax.transAxes)
            ax.text(0.98,0.5, 'R',
                    horizontalalignment='right',
                    verticalalignment='center',
                    fontsize=20*figsize, color=textfontcolor,
                    transform=ax.transAxes)
            ax.text(0.02,0.5, 'L',
                    horizontalalignment='left',
                    verticalalignment='center',
                    fontsize=20*figsize, color=textfontcolor,
                    transform=ax.transAxes)
        ax.axis('off')

        xy_slice = reorient_slice(image[:,:,xyz[2]],2)
        if not flat:
            ax = plt.subplot(gs[1,1])
        else:
            ax = plt.subplot(gs[0,2])
        ax.imshow(xy_slice, cmap=cmap, vmin=vmin, vmax=vmax)
        if overlay is not None:
            xy_overlay = reorient_slice(overlay[:,:,xyz[2]],2)
            ax.imshow(xy_overlay, alpha=overlay_alpha, cmap=overlay_cmap)
        if xyz_lines:
            # add lines
            l = mlines.Line2D([xy_slice.shape[0]-xyz[0],xy_slice.shape[0]-xyz[0]], 
                              [xyz_pad,xy_slice.shape[0]-xyz_pad],
                              color=xyz_color, alpha=xyz_alpha, linewidth=xyz_linewidth)
            ax.add_line(l)
            l = mlines.Line2D([xyz_pad,xy_slice.shape[1]-xyz_pad], 
                              [xy_slice.shape[1]-xyz[1],xy_slice.shape[1]-xyz[1]],
                              color=xyz_color, alpha=xyz_alpha, linewidth=xyz_linewidth)
            ax.add_line(l)
        if orient_labels:
            ax.text(0.5,0.98, 'A',
                    horizontalalignment='center',
                    verticalalignment='top',
                    fontsize=20*figsize, color=textfontcolor,
                    transform=ax.transAxes)
            ax.text(0.5,0.02, 'P',
                    horizontalalignment='center',
                    verticalalignment='bottom',
                    fontsize=20*figsize, color=textfontcolor,
                    transform=ax.transAxes)
            ax.text(0.98,0.5, 'R',
                    horizontalalignment='right',
                    verticalalignment='center',
                    fontsize=20*figsize, color=textfontcolor,
                    transform=ax.transAxes)
            ax.text(0.02,0.5, 'L',
                    horizontalalignment='left',
                    verticalalignment='center',
                    fontsize=20*figsize, color=textfontcolor,
                    transform=ax.transAxes)
        ax.axis('off')

        if not flat:
            # empty corner
            ax = plt.subplot(gs[1,0])
            if text is not None:
                # add text
                left, width = .25, .5
                bottom, height = .25, .5
                right = left + width
                top = bottom + height
                ax.text(0.5*(left+right)+text_dx, 0.5*(bottom+top)+text_dy, text,
                        horizontalalignment='center',
                        verticalalignment='center',
                        fontsize=textfontsize, color=textfontcolor,
                        transform=ax.transAxes)
            #ax.text(0.5, 0.5)
            ax.imshow(np.zeros(image.shape[:-1]), cmap='Greys_r')
            ax.axis('off')
            
    ## multi-channel images ##
    elif image.components > 1:
        raise ValueError('Multi-channel images not currently supported!')

    if filename is not None:
        plt.savefig(filename, dpi=dpi, transparent=transparent)
        plt.close(fig)
    else:
        plt.show()

    # turn warnings back to default
    warnings.simplefilter('default')