#!/usr/bin/env python3
"""
This module contains functions for the input and output of images, data, and
measurement results for BNMs
"""

import warnings

import numpy as np
from skimage.util import img_as_ubyte
from scipy.misc import imread, imsave
import matplotlib.pyplot as plt
import matplotlib.animation as animation
from scipy.stats import lognorm, norm
import skimage.measure
import scipy.ndimage.morphology as morphology

from aquami import display


__author__ = "Joshua Stuckner"


    

def uint8(img):
    with warnings.catch_warnings():
        warnings.simplefilter("ignore")
        img = img_as_ubyte(img)
    return img

def load(path):
    img = imread(path)
    img = uint8(img)
    return img

    
def _uint16_to_uint8(image, display_min=0, display_max=0):
    """Helper function for uint16_to_uint8"""
    image = np.array(image, copy=True)
    if display_min == 0 and display_max == 0:
        sigma = 3
        display_min = image.mean() - sigma * image.std()
        display_max = image.mean() + sigma * image.std()       
    image.clip(display_min, display_max, out=image)
    image = image - display_min
    np.true_divide(image, (display_max - display_min + 1) / 256., out=image, casting="unsafe")
    return image.astype(np.uint8)

def uint16_to_uint8(image, display_min=0, display_max=0):
    """
    Uses a lookup table to quickly convert an image from uin16 to uint 8.

    Parameters
    ----------
    image : ndarray of uint16
        Image data to convert.
    display_min : int, optional
        Pixels in image with a value less than this will be clipped to 0.
    display_max : int, optional
        Pixels in image with a value higher than this will be clipped to 255.
        If both display_max and display_min are set to 0, the image intensity
        range will be automatically selected.

    Returns
    -------
    ndarray of uint8
        Image converted to uint8
    """
    
    lut = np.arange(2**16, dtype='uint16')
    lut = _uint16_to_uint8(lut, display_min, display_max)
    return np.take(lut, image)

def save_movie(frames, save_name, fps=2, bit_rate=-1):
    """
    Saves a movie of the images in frames.

    Parameters
    ----------
    frames : list of images
        Contains images to create a movie from.
    fps : double, optional
        Frames per second of movie.  Defaults is 20.
    save_name : string
        Path to save the movie.
    bit_rate : int
        bits per second of movie.  See matplotlib.animation.writers
    """

    fig = plt.figure()
    plt.axis('off')
    plt.gca().set_position([0, 0, 1, 1])
    ims = []

    for i in frames:
        ims.append([plt.imshow(i, animated=True)])

    try: #FFmpeg
        ani = animation.ArtistAnimation(fig, ims, fps / 1000, True, fps / 1000)
        writer = animation.writers['ffmpeg'](fps=fps, bitrate=bit_rate,
                                             extra_args=['-vcodec', 'libx264',
                                                         '-pix_fmt', 'yuv420p'])
        ani.save(save_name, writer=writer)

    except RuntimeError:
        print("To save the movie you must install FFmpeg.", 
                  " Go to http://adaptivesamples.com/how-to-install-ffmpeg-on",
                  "-windows/ and follow the directions to install the FFmpeg ",
                  "movie writer.", sep='')
    
##def save_movie(stack, fps=20, directory=None,
##               save_name=None, initialdir='/'):
##    """
##    Saves a movie of the images to the specified directory.
##
##    Parameters
##    ----------
##    stack : ndarray (3D)
##        Contains images to create a movie from.
##    fps : double, optional
##        Frames per second of movie.  Defaults is 20.
##    directory : string, optional
##        Directory to save the movie.  If not provided the user will be asked.
##    image_names : string, optional
##        What to name the movie.  If not provided the user will be asked.
##    initialdir : string, optional
##        If the save directory is not given, this parameter is used as the
##        starting directory for the save as dialog.
##    """
##          
##    if directory is None:
##        # User selects the directory to save the movie.
##        root = tkinter.Tk()
##        directory = tkinter.filedialog.askdirectory(
##                parent=root,
##                initialdir=initialdir,
##                title='Select a folder to save your movie.'
##                )
##        root.destroy()
##        if len(directory) == 0:
##            print("No folder selected.  Movie will not be saved.")
##            return
##
##    if save_name is None:
##        save_name = input("What do you want to name your movie? ")
##
##    save_path = directory + '/' + save_name + '.mp4'
##
##    # Get the size of the images
##    try:
##        rows, cols = stack[0].shape
##        color = False
##    except ValueError:
##        rows, cols, _ = stack[0].shape
##        color = True
##
##    try: #FFmpeg
##        fig = plt.figure()
##        plt.axis('off')
##        ims = []
##        for im in stack:
##            ims.append([plt.imshow(im, cmap=plt.get_cmap('gray'), animated=True)])
##            ims.append([plt.imshow(im, animated=True)])
##        ani = animation.ArtistAnimation(fig, ims, fps/1000, True, fps/1000)
##        writer = animation.writers['ffmpeg'](fps=fps)
##        ani.save(save_path, writer=writer)
##    except RuntimeError:
##        try: #openCV
##            import cv2
##            video = cv2.VideoWriter(save_path, 0,
##                                    fps, (int(cols), int(rows)), isColor=True)
##
##            for i in range(len(stack)):
##                video.write(stack[i])
##            cv2.destroyAllWindows()
##            video.release()
##        except ImportError:
##            print("To save the movie you must install either openCV or FFmpeg.", 
##                  " Go to http://adaptivesamples.com/how-to-install-ffmpeg-on",
##                  "-windows/ and follow the directions to install the FFmpeg ",
##                  "movie writer.", sep='')

def saveBinary(mask, path):
    """
    Saves the binary mask to the path.

    Parameters
    ----------
    mask : 2darray
        Contains mask to save.
    path : str
        Path to binary file.  ex: "c:/images/binary.tiff"
    """
    imsave(path, mask)

def pdfSaveImage(pdf, image, title='', cmap=None):
    plt.imshow(image, cmap=cmap)
    plt.title(title)
    plt.axis('off')
    pdf.savefig()
    plt.close()

def pdfSaveOverlay(pdf, image, mask, title=''):
    overlay = display.overlayMask(image, mask, return_overlay=True)
    plt.imshow(overlay)
    plt.title(title)
    plt.axis('off')
    pdf.savefig()
    plt.close()

def pdfSaveHist(pdf, data, title=None, xlabel=None,
             numBins=None, gauss=False, log=False):
    
    # Determine optimum number of bins.
    if numBins is None:
        u = len(np.unique(data))
        numBins = int(2*u**(1/2))
        if numBins < 4:
            numBins = len(np.unique(data))

    # Create the histogram.
    n, bins, patches = plt.hist(data, bins=numBins, normed=1, edgecolor='black')
    if log:
        logfit = lognorm.fit(data.flatten(), floc=0)
        pdf_plot = lognorm.pdf(bins, logfit[0], loc=logfit[1], scale=logfit[2])
        plt.plot(bins, pdf_plot, 'r--', linewidth=3, label='lognorm')
    if gauss:
        gfit = norm.fit(data.flatten())
        gauss_plot = norm.pdf(bins, gfit[0], gfit[1])
        plt.plot(bins, gauss_plot, 'g--', linewidth=3, label='gaussian')

    # Save the histogram
    if title is not None:
        plt.title(title)
    if xlabel is not None:
        plt.xlabel(xlabel)
    ax = plt.gca()
    ax.yaxis.set_visible(False)
    if gauss or log:
        plt.legend()
    pdf.savefig()
    plt.close()

def pdfSaveSkel(pdf, skeleton, mask, dialate=False, title=None):
    """
    Displays skelatal data on top of an outline of a binary mask. For example,
    displays a medial axis transform over an outline of segmented ligaments.

    Parameters
    ----------
    skeleton : 2D array
        Data to be displayed.
    mask : binary 2D array
        Mask of segmentation data, the outline of which is displayed along with
        the skel data.
    dialate : boolean, optional
        If dialate is true, the skelatal data will be made thicker in the
        display.
    title : str, optional
        Text to be displayed above the image.
    """

    skel = np.copy(skeleton)

    # Find the outlines of the mask and make an outline mask called outlines.
    contours = skimage.measure.find_contours(mask, 0.5)
    outlines = np.zeros((mask.shape), dtype='uint8')
    for n, contour in enumerate(contours):
        for i in range(len(contour)):
            outlines[int(contour[i,0]), int(contour[i,1])] = 255

    # Make the skel data thicker if dialate is true.
    if dialate:
        skel = morphology.grey_dilation(skel, size=(3,3))


    # Scale the skel data to uint8 and add the outline mask
    with warnings.catch_warnings():
        warnings.simplefilter("ignore")
        skel = skel.astype(np.float32) # convert to float
        skel -= skel.min() # ensure the minimal value is 0.0
        if skel.max() != 0:
            skel /= skel.max() # maximum value in image is now 1.0
    skel = np.uint8(plt.cm.spectral(skel)*255) # apply colormap to skel data.
    for i in range(3):
        skel[:,:,i] += outlines

    # Display the results.
    plt.imshow(skel, cmap = plt.cm.spectral, interpolation='none')
    plt.axis('off')
    figManager = plt.get_current_fig_manager()
    figManager.window.showMaximized()
    if title is None:
        plt.gca().set_position([0, 0, 1, 1])
    else:
        plt.gca().set_position([0, 0, 1, 0.95])
        plt.title(title)
    pdf.savefig()
    plt.close()

def parsePath(path):
    folder = '/'.join(path.split('/')[:-1])
    fname = path.split('/')[-1].split('.')[:-1][0]
    ftype = path.split('.')[-1]
    fnametype = fname + '.' + ftype
    return folder, fname, ftype, fnametype
