#!/usr/bin/env python
"""
Creates 3 animated .gifs of the input NIFTI in axial, transverse, and coronal
orientation.

All backgound images are printed in grayscale. In the overlays, positive values
will always be red, and negative values will always be blue.

Outputs will be written as <prefix>_axial.gif, etc.

Usage:
    qc-montage [options] <image> <prefix>

Arguments:
    <image>      background image file name
    <prefix>     output

Options:
    --label=<label>  an optional overlay statmap
    --alpha=<alpha>  opacity of overlay (1 = full, 0 = none) [default: 0.75]

Dependencies:

    imagemagick, matplotlib version 1.3+

"""

import os, sys
import logging
import numpy as np
import nibabel as nib
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import matplotlib.animation as animation
from scipy import ndimage as nd
from epitome.docopt import docopt

logging.basicConfig(level=logging.WARN, format="[%(name)s] %(levelname)s: %(message)s")
logger = logging.getLogger(os.path.basename(__file__))

def bounding_box(filename):
    """
    Finds a box that only includes all nonzero voxels in a 3D image. Output box
    is represented as 3 x 2 numpy array with rows denoting x, y, z, and columns
    denoting stand and end slices.

    Usage:
        box = bounding_box(filename)
    """

    # find 3D bounding box
    box = np.zeros((3,2))  # init bounding box
    flag = 0  # switch to ascending

    for i, dim in enumerate(filename.shape): # loop through (x, y, z)

        # ascending search
        while flag == 0:
            for dim_test in np.arange(dim):

                # get sum of all values in each slice
                if i == 0:
                    test = np.sum(filename[dim_test, :, :])
                elif i == 1:
                    test = np.sum(filename[:, dim_test, :])
                elif i == 2:
                    test = np.sum(filename[:, :, dim_test])

                if test >= 1:  # if slice is nonzero, set starting bound
                    box[i, 0] = dim_test
                    flag = 1  # switch to descending
                    break  # get us out of this nested nonsense

        # descending search
        while flag == 1:
            for dim_test in np.arange(dim):

                dim_test = dim-dim_test - 1  # we have to reverse things

                # get sum of all values in each slice
                if i == 0:
                    test = np.sum(filename[dim_test, :, :])
                elif i == 1:
                    test = np.sum(filename[:, dim_test, :])
                elif i == 2:
                    test = np.sum(filename[:, :, dim_test])

                if test >= 1:  # if slice is nonzero, set ending bound
                    box[i, 1] = dim_test
                    flag = 0  # switch to ascending
                    break  # get us out of this nested nonsense

    return box

def animate(i, im, image):
    im.set_array(image[:, i, :])
    return [im]

def main():
    arguments = docopt(__doc__)
    image     = arguments['<image>']
    prefix    = arguments['<prefix>']
    label     = arguments['--label']
    alpha     = arguments['--alpha']

    logging.info('Starting')

    # load in the data
    image = nib.load(image).get_data()

    if label:
        label = nib.load(label).get_data()

    # reorient the data to radiological (does this generalize?)
    image = np.transpose(image, (2,0,1))
    image = np.rot90(image, 2)

    if label:
        label = np.transpose(label, (2,0,1))
        label = np.rot90(label, 2)

        # get size ratio between over + underlay, and set to match
        dsfactor = [a/float(r) for a,r in zip(image.shape, label.shape)]
        label = nd.interpolation.zoom(label, zoom=dsfactor)

        # crop data to statmap bounds
        box = bounding_box(label) # get the label bounds
        image = image[box[0,0]:box[0,1], box[1,0]:box[1,1], box[2,0]:box[2,1]]
        label = label[box[0,0]:box[0,1], box[1,0]:box[1,1], box[2,0]:box[2,1]]
        label = np.ma.masked_where(label < 1, label) # small values transparent

    stepsx=image.shape[0]
    stepsy=image.shape[1]
    stepsz=image.shape[2]

    cmap = plt.cm.RdBu_r  # the overlay colormap
    cmap.set_bad('g', 0)  # value for transparent pixels in the overlay

    fig = plt.figure()
    #ax = plt.axes(xlim=(0, 2), ylim=(-2, 2))
    im = plt.imshow(image[:, 0, :], cmap=plt.cm.gray, interpolation='nearest', vmin=image.min(), vmax=image.max())
    im.axes.get_xaxis().set_visible(False)
    im.axes.get_yaxis().set_visible(False)

    #im = plt.imshow(label[0, :, :], cmap=cmap, alpha=alpha, interpolation='nearest', animated=True)
    #im.set_frame_on(False)  # clean up unnecessary detail

    ani = animation.FuncAnimation(fig, animate, frames=stepsx, fargs = (im, image))
    ani.save(prefix + '.gif', writer='imagemagick', fps=15)

if __name__ == "__main__":
    main()

