import pkg_resources

import numpy as np
import pandas as pd

from ..core.pycompat import basestring

def _load_default_cmap(fname='default_colormap.csv'):
    """
    Returns viridis color map
    """
    from matplotlib.colors import LinearSegmentedColormap

    # Not sure what the first arg here should be
    f = pkg_resources.resource_stream(__name__, fname)
    cm_data = pd.read_csv(f, header=None).values

    return LinearSegmentedColormap.from_list('viridis', cm_data)

def _determine_extend(calc_data, vmin, vmax):
    extend_min = calc_data.min() < vmin
    extend_max = calc_data.max() > vmax
    if extend_min and extend_max:
        extend = 'both'
    elif extend_min:
        extend = 'min'
    elif extend_max:
        extend = 'max'
    else:
        extend = 'neither'
    return extend

def _build_discrete_cmap(cmap, levels, extend, filled):
    """
    Build a discrete colormap and normalization of the data.
    """
    import matplotlib as mpl

    if not filled:
        # non-filled contour plots
        extend = 'max'

    if extend == 'both':
        ext_n = 2
    elif extend in ['min', 'max']:
        ext_n = 1
    else:
        ext_n = 0

    n_colors = len(levels) + ext_n - 1
    pal = _color_palette(cmap, n_colors)

    new_cmap, cnorm = mpl.colors.from_levels_and_colors(
        levels, pal, extend=extend)
    # copy the old cmap name, for easier testing
    new_cmap.name = getattr(cmap, 'name', cmap)

    return new_cmap, cnorm

def _color_palette(cmap, n_colors):
    import matplotlib.pyplot as plt
    from matplotlib.colors import ListedColormap
    colors_i = np.linspace(0, 1., n_colors)
    if isinstance(cmap, (list, tuple)):
        # we have a list of colors
        try:
            # first try to turn it into a palette with seaborn
            from seaborn.apionly import color_palette
            pal = color_palette(cmap, n_colors=n_colors)
        except ImportError:
            # if that fails, use matplotlib
            # in this case, is there any difference between mpl and seaborn?
            cmap = ListedColormap(cmap, N=n_colors)
            pal = cmap(colors_i)
    elif isinstance(cmap, basestring):
        # we have some sort of named palette
        try:
            # first try to turn it into a palette with seaborn
            from seaborn.apionly import color_palette
            pal = color_palette(cmap, n_colors=n_colors)
        except (ImportError, ValueError):
            # ValueError is raised when seaborn doesn't like a colormap
            # (e.g. jet). If that fails, use matplotlib
            try:
                # is this a matplotlib cmap?
                cmap = plt.get_cmap(cmap)
            except ValueError:
                # or maybe we just got a single color as a string
                cmap = ListedColormap([cmap], N=n_colors)
            pal = cmap(colors_i)
    else:
        # cmap better be a LinearSegmentedColormap (e.g. viridis)
        pal = cmap(colors_i)

    return pal

def _determine_cmap_params(plot_data, vmin=None, vmax=None, cmap=None,
                           center=None, robust=False, extend=None,
                           levels=None, filled=True, cnorm=None):
    """
    Use some heuristics to set good defaults for colorbar and range.

    Adapted from Seaborn:
    https://github.com/mwaskom/seaborn/blob/v0.6/seaborn/matrix.py#L158

    Parameters
    ==========
    plot_data: Numpy array
        Doesn't handle xarray objects

    Returns
    =======
    cmap_params : dict
        Use depends on the type of the plotting function
    """
    ROBUST_PERCENTILE = 2.0
    import matplotlib as mpl

    calc_data = np.ravel(plot_data[~pd.isnull(plot_data)])

    # Setting center=False prevents a divergent cmap
    possibly_divergent = center is not False

    # Set center to 0 so math below makes sense but remember its state
    center_is_none = False
    if center is None:
        center = 0
        center_is_none = True

    # Setting both vmin and vmax prevents a divergent cmap
    if (vmin is not None) and (vmax is not None):
        possibly_divergent = False

    # vlim might be computed below
    vlim = None

    if vmin is None:
        if robust:
            vmin = np.percentile(calc_data, ROBUST_PERCENTILE)
        else:
            vmin = calc_data.min()
    elif possibly_divergent:
        vlim = abs(vmin - center)

    if vmax is None:
        if robust:
            vmax = np.percentile(calc_data, 100 - ROBUST_PERCENTILE)
        else:
            vmax = calc_data.max()
    elif possibly_divergent:
        vlim = abs(vmax - center)

    if possibly_divergent:
        # kwargs not specific about divergent or not: infer defaults from data
        divergent = ((vmin < 0) and (vmax > 0)) or not center_is_none
    else:
        divergent = False

    # A divergent map should be symmetric around the center value
    if divergent:
        if vlim is None:
            vlim = max(abs(vmin - center), abs(vmax - center))
        vmin, vmax = -vlim, vlim

    # Now add in the centering value and set the limits
    vmin += center
    vmax += center

    # Choose default colormaps if not provided
    if cmap is None:
        if divergent:
            cmap = "RdBu_r"
        else:
            cmap = "viridis"

    # Allow viridis before matplotlib 1.5
    if cmap == "viridis":
        cmap = _load_default_cmap()

    # Handle discrete levels
    if levels is not None:
        if isinstance(levels, int):
            ticker = mpl.ticker.MaxNLocator(levels)
            levels = ticker.tick_values(vmin, vmax)
        vmin, vmax = levels[0], levels[-1]

    if extend is None:
        extend = _determine_extend(calc_data, vmin, vmax)

    if levels is not None:
        cmap, cnorm = _build_discrete_cmap(cmap, levels, extend, filled)

    return dict(vmin=vmin, vmax=vmax, cmap=cmap, extend=extend,
                levels=levels, norm=cnorm)

def _infer_xy_labels(darray, x, y):
    """
    Determine x and y labels. For use in _plot2d

    darray must be a 2 dimensional data array.
    """

    if x is None and y is None:
        if darray.ndim != 2:
            raise ValueError('DataArray must be 2d')
        y, x = darray.dims
    elif x is None or y is None:
        raise ValueError('cannot supply only one of x and y')
    elif any(k not in darray.coords for k in (x, y)):
        raise ValueError('x and y must be coordinate variables')
    return x, y


"""Functions to visualize matrices of data."""
import itertools

import colorsys
import matplotlib as mpl
import matplotlib.pyplot as plt
from matplotlib import gridspec
import numpy as np
import pandas as pd
from scipy.spatial import distance
from scipy.cluster import hierarchy

from .axisgrid import Grid
from .palettes import cubehelix_palette
from .utils import despine, axis_ticklabels_overlap
from .external.six.moves import range

def _index_to_label(index):
    """Convert a pandas index or multiindex to an axis label."""
    if isinstance(index, pd.MultiIndex):
        return "-".join(map(str, index.names))
    else:
        return index.name

def _index_to_ticklabels(index):
    """Convert a pandas index or multiindex into ticklabels."""
    if isinstance(index, pd.MultiIndex):
        return ["-".join(map(str, i)) for i in index.values]
    else:
        return index.values

def _convert_colors(colors):
    """Convert either a list of colors or nested lists of colors to RGB."""
    to_rgb = mpl.colors.colorConverter.to_rgb
    try:
        to_rgb(colors[0])
        # If this works, there is only one level of colors
        return list(map(to_rgb, colors))
    except ValueError:
        # If we get here, we have nested lists
        return [list(map(to_rgb, l)) for l in colors]

def _matrix_mask(data, mask):
    """Ensure that data and mask are compatabile and add missing values.

    Values will be plotted for cells where ``mask`` is ``False``.

    ``data`` is expected to be a DataFrame; ``mask`` can be an array or
    a DataFrame.

    """
    if mask is None:
        mask = np.zeros(data.shape, np.bool)

    if isinstance(mask, np.ndarray):
        # For array masks, ensure that shape matches data then convert
        if mask.shape != data.shape:
            raise ValueError("Mask must have the same shape as data.")

        mask = pd.DataFrame(mask,
                            index=data.index,
                            columns=data.columns,
                            dtype=np.bool)

    elif isinstance(mask, pd.DataFrame):
        # For DataFrame masks, ensure that semantic labels match data
        if not mask.index.equals(data.index) \
           and mask.columns.equals(data.columns):
            err = "Mask must have the same index and columns as data."
            raise ValueError(err)

    # Add any cells with missing data to the mask
    # This works around an issue where `plt.pcolormesh` doesn't represent
    # missing data properly
    mask = mask | pd.isnull(data)

    return mask

class _HeatMapper(object):
    """Draw a heatmap plot of a matrix with nice labels and colormaps."""

    def __init__(self, data, vmin, vmax, cmap, center, robust, annot, fmt,
                 annot_kws, cbar, cbar_kws,
                 xticklabels=True, yticklabels=True, mask=None):
        """Initialize the plotting object."""
        # We always want to have a DataFrame with semantic information
        # and an ndarray to pass to matplotlib
        if isinstance(data, pd.DataFrame):
            plot_data = data.values
        else:
            plot_data = np.asarray(data)
            data = pd.DataFrame(plot_data)

        # Validate the mask and convet to DataFrame
        mask = _matrix_mask(data, mask)

        # Reverse the rows so the plot looks like the matrix
        plot_data = plot_data[::-1]
        data = data.ix[::-1]
        mask = mask.ix[::-1]

        plot_data = np.ma.masked_where(np.asarray(mask), plot_data)

        # Get good names for the rows and columns
        xtickevery = 1
        if isinstance(xticklabels, int) and xticklabels > 1:
            xtickevery = xticklabels
            xticklabels = _index_to_ticklabels(data.columns)
        elif isinstance(xticklabels, bool) and xticklabels:
            xticklabels = _index_to_ticklabels(data.columns)
        elif isinstance(xticklabels, bool) and not xticklabels:
            xticklabels = ['' for _ in range(data.shape[1])]

        ytickevery = 1
        if isinstance(yticklabels, int) and yticklabels > 1:
            ytickevery = yticklabels
            yticklabels = _index_to_ticklabels(data.index)
        elif isinstance(yticklabels, bool) and yticklabels:
            yticklabels = _index_to_ticklabels(data.index)
        elif isinstance(yticklabels, bool) and not yticklabels:
            yticklabels = ['' for _ in range(data.shape[0])]
        else:
            yticklabels = yticklabels[::-1]

        # Get the positions and used label for the ticks
        nx, ny = data.T.shape
        xstart, xend, xstep = 0, nx, xtickevery
        self.xticks = np.arange(xstart, xend, xstep) + .5
        self.xticklabels = xticklabels[xstart:xend:xstep]
        ystart, yend, ystep = (ny - 1) % ytickevery, ny, ytickevery
        self.yticks = np.arange(ystart, yend, ystep) + .5
        self.yticklabels = yticklabels[ystart:yend:ystep]

        # Get good names for the axis labels
        xlabel = _index_to_label(data.columns)
        ylabel = _index_to_label(data.index)
        self.xlabel = xlabel if xlabel is not None else ""
        self.ylabel = ylabel if ylabel is not None else ""

        # Determine good default values for the colormapping
        self._determine_cmap_params(plot_data, vmin, vmax,
                                    cmap, center, robust)

        # Save other attributes to the object
        self.data = data
        self.plot_data = plot_data
        self.annot = annot
        self.fmt = fmt
        self.annot_kws = {} if annot_kws is None else annot_kws
        self.cbar = cbar
        self.cbar_kws = {} if cbar_kws is None else cbar_kws

    def _determine_cmap_params(self, plot_data, vmin, vmax,
                               cmap, center, robust):
        """Use some heuristics to set good defaults for colorbar and range."""
        calc_data = plot_data.data[~np.isnan(plot_data.data)]
        if vmin is None:
            vmin = np.percentile(calc_data, 2) if robust else calc_data.min()
        if vmax is None:
            vmax = np.percentile(calc_data, 98) if robust else calc_data.max()

        # Simple heuristics for whether these data should  have a divergent map
        divergent = ((vmin < 0) and (vmax > 0)) or center is not None

        # Now set center to 0 so math below makes sense
        if center is None:
            center = 0

        # A divergent map should be symmetric around the center value
        if divergent:
            vlim = max(abs(vmin - center), abs(vmax - center))
            vmin, vmax = -vlim, vlim
        self.divergent = divergent

        # Now add in the centering value and set the limits
        vmin += center
        vmax += center
        self.vmin = vmin
        self.vmax = vmax

        # Choose default colormaps if not provided
        if cmap is None:
            if divergent:
                self.cmap = "RdBu_r"
            else:
                self.cmap = cubehelix_palette(light=.95, as_cmap=True)
        else:
            self.cmap = cmap

    def _annotate_heatmap(self, ax, mesh):
        """Add textual labels with the value in each cell."""
        xpos, ypos = np.meshgrid(ax.get_xticks(), ax.get_yticks())
        for x, y, val, color in zip(xpos.flat, ypos.flat,
                                    mesh.get_array(), mesh.get_facecolors()):
            if val is not np.ma.masked:
                _, l, _ = colorsys.rgb_to_hls(*color[:3])
                text_color = ".15" if l > .5 else "w"
                val = ("{:" + self.fmt + "}").format(val)
                ax.text(x, y, val, color=text_color,
                        ha="center", va="center", **self.annot_kws)

    def plot(self, ax, cax, kws):
        """Draw the heatmap on the provided Axes."""
        # Remove all the Axes spines
        despine(ax=ax, left=True, bottom=True)

        # Draw the heatmap
        mesh = ax.pcolormesh(self.plot_data, vmin=self.vmin, vmax=self.vmax,
                             cmap=self.cmap, **kws)

        # Set the axis limits
        ax.set(xlim=(0, self.data.shape[1]), ylim=(0, self.data.shape[0]))

        # Add row and column labels
        ax.set(xticks=self.xticks, yticks=self.yticks)
        xtl = ax.set_xticklabels(self.xticklabels)
        ytl = ax.set_yticklabels(self.yticklabels, rotation="vertical")

        # Possibly rotate them if they overlap
        plt.draw()
        if axis_ticklabels_overlap(xtl):
            plt.setp(xtl, rotation="vertical")
        if axis_ticklabels_overlap(ytl):
            plt.setp(ytl, rotation="horizontal")

        # Add the axis labels
        ax.set(xlabel=self.xlabel, ylabel=self.ylabel)

        # Annotate the cells with the formatted values
        if self.annot:
            self._annotate_heatmap(ax, mesh)

        # Possibly add a colorbar
        if self.cbar:
            ticker = mpl.ticker.MaxNLocator(6)
            cb = ax.figure.colorbar(mesh, cax, ax,
                                    ticks=ticker, **self.cbar_kws)
            cb.outline.set_linewidth(0)

def heatmap(data, vmin=None, vmax=None, cmap=None, center=None, robust=False,
            annot=False, fmt=".2g", annot_kws=None,
            linewidths=0, linecolor="white",
            cbar=True, cbar_kws=None, cbar_ax=None,
            square=False, ax=None, xticklabels=True, yticklabels=True,
            mask=None,
            **kwargs):
    """Plot rectangular data as a color-encoded matrix.

    This function tries to infer a good colormap to use from the data, but
    this is not guaranteed to work, so take care to make sure the kind of
    colormap (sequential or diverging) and its limits are appropriate.

    This is an Axes-level function and will draw the heatmap into the
    currently-active Axes if none is provided to the ``ax`` argument.  Part of
    this Axes space will be taken and used to plot a colormap, unless ``cbar``
    is False or a separate Axes is provided to ``cbar_ax``.

    Parameters
    ----------
    data : rectangular dataset
        2D dataset that can be coerced into an ndarray. If a Pandas DataFrame
        is provided, the index/column information will be used to label the
        columns and rows.
    vmin, vmax : floats, optional
        Values to anchor the colormap, otherwise they are inferred from the
        data and other keyword arguments. When a diverging dataset is inferred,
        one of these values may be ignored.
    cmap : matplotlib colormap name or object, optional
        The mapping from data values to color space. If not provided, this
        will be either a cubehelix map (if the function infers a sequential
        dataset) or ``RdBu_r`` (if the function infers a diverging dataset).
    center : float, optional
        The value at which to center the colormap. Passing this value implies
        use of a diverging colormap.
    robust : bool, optional
        If True and ``vmin`` or ``vmax`` are absent, the colormap range is
        computed with robust quantiles instead of the extreme values.
    annot : bool, optional
        If True, write the data value in each cell.
    fmt : string, optional
        String formatting code to use when ``annot`` is True.
    annot_kws : dict of key, value mappings, optional
        Keyword arguments for ``ax.text`` when ``annot`` is True.
    linewidths : float, optional
        Width of the lines that will divide each cell.
    linecolor : color, optional
        Color of the lines that will divide each cell.
    cbar : boolean, optional
        Whether to draw a colorbar.
    cbar_kws : dict of key, value mappings, optional
        Keyword arguments for `fig.colorbar`.
    cbar_ax : matplotlib Axes, optional
        Axes in which to draw the colorbar, otherwise take space from the
        main Axes.
    square : boolean, optional
        If True, set the Axes aspect to "equal" so each cell will be
        square-shaped.
    ax : matplotlib Axes, optional
        Axes in which to draw the plot, otherwise use the currently-active
        Axes.
    xticklabels : list-like, int, or bool, optional
        If True, plot the column names of the dataframe. If False, don't plot
        the column names. If list-like, plot these alternate labels as the
        xticklabels. If an integer, use the column names but plot only every
        n label.
    yticklabels : list-like, int, or bool, optional
        If True, plot the row names of the dataframe. If False, don't plot
        the row names. If list-like, plot these alternate labels as the
        yticklabels. If an integer, use the index names but plot only every
        n label.
    mask : boolean array or DataFrame, optional
        If passed, data will not be shown in cells where ``mask`` is True.
        Cells with missing values are automatically masked.
    kwargs : other keyword arguments
        All other keyword arguments are passed to ``ax.pcolormesh``.

    Returns
    -------
    ax : matplotlib Axes
        Axes object with the heatmap.

    Examples
    --------

    Plot a heatmap for a numpy array:

    .. plot::
        :context: close-figs

        >>> import numpy as np; np.random.seed(0)
        >>> import seaborn as sns; sns.set()
        >>> uniform_data = np.random.rand(10, 12)
        >>> ax = sns.heatmap(uniform_data)

    Change the limits of the colormap:

    .. plot::
        :context: close-figs

        >>> ax = sns.heatmap(uniform_data, vmin=0, vmax=1)

    Plot a heatmap for data centered on 0:

    .. plot::
        :context: close-figs

        >>> normal_data = np.random.randn(10, 12)
        >>> ax = sns.heatmap(normal_data)

    Plot a dataframe with meaningful row and column labels:

    .. plot::
        :context: close-figs

        >>> flights = sns.load_dataset("flights")
        >>> flights = flights.pivot("month", "year", "passengers")
        >>> ax = sns.heatmap(flights)

    Annotate each cell with the numeric value using integer formatting:

    .. plot::
        :context: close-figs

        >>> ax = sns.heatmap(flights, annot=True, fmt="d")

    Add lines between each cell:

    .. plot::
        :context: close-figs

        >>> ax = sns.heatmap(flights, linewidths=.5)

    Use a different colormap:

    .. plot::
        :context: close-figs

        >>> ax = sns.heatmap(flights, cmap="YlGnBu")

    Center the colormap at a specific value:

    .. plot::
        :context: close-figs

        >>> ax = sns.heatmap(flights, center=flights.loc["January", 1955])

    Plot every other column label and don't plot row labels:

    .. plot::
        :context: close-figs

        >>> data = np.random.randn(50, 20)
        >>> ax = sns.heatmap(data, xticklabels=2, yticklabels=False)

    Don't draw a colorbar:

    .. plot::
        :context: close-figs

        >>> ax = sns.heatmap(flights, cbar=False)

    Use different axes for the colorbar:

    .. plot::
        :context: close-figs

        >>> grid_kws = {"height_ratios": (.9, .05), "hspace": .3}
        >>> f, (ax, cbar_ax) = plt.subplots(2, gridspec_kw=grid_kws)
        >>> ax = sns.heatmap(flights, ax=ax,
        ...                  cbar_ax=cbar_ax,
        ...                  cbar_kws={"orientation": "horizontal"})

    Use a mask to plot only part of a matrix

    .. plot::
        :context: close-figs

        >>> corr = np.corrcoef(np.random.randn(10, 200))
        >>> mask = np.zeros_like(corr)
        >>> mask[np.triu_indices_from(mask)] = True
        >>> with sns.axes_style("white"):
        ...     ax = sns.heatmap(corr, mask=mask, vmax=.3, square=True)

    """
    # Initialize the plotter object
    plotter = _HeatMapper(data, vmin, vmax, cmap, center, robust, annot, fmt,
                          annot_kws, cbar, cbar_kws, xticklabels, yticklabels,
                          mask)

    # Add the pcolormesh kwargs here
    kwargs["linewidths"] = linewidths
    kwargs["edgecolor"] = linecolor

    # Draw the plot and return the Axes
    if ax is None:
        ax = plt.gca()
    if square:
        ax.set_aspect("equal")
    plotter.plot(ax, cbar_ax, kwargs)
    return ax

class _DendrogramPlotter(object):
    """Object for drawing tree of similarities between data rows/columns"""

    def __init__(self, data, linkage, metric, method, axis, label, rotate):
        """Plot a dendrogram of the relationships between the columns of data

        Parameters
        ----------
        data : pandas.DataFrame
            Rectangular data
        """
        self.axis = axis
        if self.axis == 1:
            data = data.T

        if isinstance(data, pd.DataFrame):
            array = data.values
        else:
            array = np.asarray(data)
            data = pd.DataFrame(array)

        self.array = array
        self.data = data

        self.shape = self.data.shape
        self.metric = metric
        self.method = method
        self.axis = axis
        self.label = label
        self.rotate = rotate

        if linkage is None:
            self.linkage = self.calculated_linkage
        else:
            self.linkage = linkage
        self.dendrogram = self.calculate_dendrogram()

        # Dendrogram ends are always at multiples of 5, who knows why
        ticks = 10 * np.arange(self.data.shape[0]) + 5

        if self.label:
            ticklabels = _index_to_ticklabels(self.data.index)
            ticklabels = [ticklabels[i] for i in self.reordered_ind]
            if self.rotate:
                self.xticks = []
                self.yticks = ticks
                self.xticklabels = []

                self.yticklabels = ticklabels
                self.ylabel = _index_to_label(self.data.index)
                self.xlabel = ''
            else:
                self.xticks = ticks
                self.yticks = []
                self.xticklabels = ticklabels
                self.yticklabels = []
                self.ylabel = ''
                self.xlabel = _index_to_label(self.data.index)
        else:
            self.xticks, self.yticks = [], []
            self.yticklabels, self.xticklabels = [], []
            self.xlabel, self.ylabel = '', ''

        if self.rotate:
            self.X = self.dendrogram['dcoord']
            self.Y = self.dendrogram['icoord']
        else:
            self.X = self.dendrogram['icoord']
            self.Y = self.dendrogram['dcoord']

    def _calculate_linkage_scipy(self):
        if np.product(self.shape) >= 10000:
            UserWarning('This will be slow... (gentle suggestion: '
                        '"pip install fastcluster")')

        pairwise_dists = distance.pdist(self.array, metric=self.metric)
        linkage = hierarchy.linkage(pairwise_dists, method=self.method)
        del pairwise_dists
        return linkage

    def _calculate_linkage_fastcluster(self):
        import fastcluster
        # Fastcluster has a memory-saving vectorized version, but only
        # with certain linkage methods, and mostly with euclidean metric
        vector_methods = ('single', 'centroid', 'median', 'ward')
        euclidean_methods = ('centroid', 'median', 'ward')
        euclidean = self.metric == 'euclidean' and self.method in \
            euclidean_methods
        if euclidean or self.method == 'single':
            return fastcluster.linkage_vector(self.array,
                                              method=self.method,
                                              metric=self.metric)
        else:
            pairwise_dists = distance.pdist(self.array, metric=self.metric)
            linkage = fastcluster.linkage(pairwise_dists, method=self.method)
            del pairwise_dists
            return linkage

    @property
    def calculated_linkage(self):
        try:
            return self._calculate_linkage_fastcluster()
        except ImportError:
            return self._calculate_linkage_scipy()

    def calculate_dendrogram(self):
        """Calculates a dendrogram based on the linkage matrix

        Made a separate function, not a property because don't want to
        recalculate the dendrogram every time it is accessed.

        Returns
        -------
        dendrogram : dict
            Dendrogram dictionary as returned by scipy.cluster.hierarchy
            .dendrogram. The important key-value pairing is
            "reordered_ind" which indicates the re-ordering of the matrix
        """
        return hierarchy.dendrogram(self.linkage, no_plot=True,
                                    color_list=['k'], color_threshold=-np.inf)

    @property
    def reordered_ind(self):
        """Indices of the matrix, reordered by the dendrogram"""
        return self.dendrogram['leaves']

    def plot(self, ax):
        """Plots a dendrogram of the similarities between data on the axes

        Parameters
        ----------
        ax : matplotlib.axes.Axes
            Axes object upon which the dendrogram is plotted

        """
        for x, y in zip(self.X, self.Y):
            ax.plot(x, y, color='k', linewidth=.5)

        if self.rotate and self.axis == 0:
            ax.invert_xaxis()
            ax.yaxis.set_ticks_position('right')

            ymax = min(map(min, self.Y)) + max(map(max, self.Y))
            ax.set_ylim(0, ymax)
            ax.invert_yaxis()
        else:
            xmax = min(map(min, self.X)) + max(map(max, self.X))
            ax.set_xlim(0, xmax)

        despine(ax=ax, bottom=True, left=True)

        ax.set(xticks=self.xticks, yticks=self.yticks,
               xlabel=self.xlabel, ylabel=self.ylabel)
        xtl = ax.set_xticklabels(self.xticklabels)
        ytl = ax.set_yticklabels(self.yticklabels, rotation='vertical')

        # Force a draw of the plot to avoid matplotlib window error
        plt.draw()
        if len(ytl) > 0 and axis_ticklabels_overlap(ytl):
            plt.setp(ytl, rotation="horizontal")
        if len(xtl) > 0 and axis_ticklabels_overlap(xtl):
            plt.setp(xtl, rotation="vertical")
        return self

def dendrogram(data, linkage=None, axis=1, label=True, metric='euclidean',
               method='average', rotate=False, ax=None):
    """Draw a tree diagram of relationships within a matrix

    Parameters
    ----------
    data : pandas.DataFrame
        Rectangular data
    linkage : numpy.array, optional
        Linkage matrix
    axis : int, optional
        Which axis to use to calculate linkage. 0 is rows, 1 is columns.
    label : bool, optional
        If True, label the dendrogram at leaves with column or row names
    metric : str, optional
        Distance metric. Anything valid for scipy.spatial.distance.pdist
    method : str, optional
        Linkage method to use. Anything valid for
        scipy.cluster.hierarchy.linkage
    rotate : bool, optional
        When plotting the matrix, whether to rotate it 90 degrees
        counter-clockwise, so the leaves face right
    ax : matplotlib axis, optional
        Axis to plot on, otherwise uses current axis

    Returns
    -------
    dendrogramplotter : _DendrogramPlotter
        A Dendrogram plotter object.

    Notes
    -----
    Access the reordered dendrogram indices with
    dendrogramplotter.reordered_ind

    """
    plotter = _DendrogramPlotter(data, linkage=linkage, axis=axis,
                                 metric=metric, method=method,
                                 label=label, rotate=rotate)
    if ax is None:
        ax = plt.gca()
    return plotter.plot(ax=ax)

class ClusterGrid(Grid):
    def __init__(self, data, pivot_kws=None, z_score=None, standard_scale=None,
                 figsize=None, row_colors=None, col_colors=None, mask=None):
        """Grid object for organizing clustered heatmap input on to axes"""

        if isinstance(data, pd.DataFrame):
            self.data = data
        else:
            self.data = pd.DataFrame(data)

        self.data2d = self.format_data(self.data, pivot_kws, z_score,
                                       standard_scale)

        self.mask = _matrix_mask(self.data2d, mask)

        if figsize is None:
            width, height = 10, 10
            figsize = (width, height)
        self.fig = plt.figure(figsize=figsize)

        if row_colors is not None:
            row_colors = _convert_colors(row_colors)
        self.row_colors = row_colors
        if col_colors is not None:
            col_colors = _convert_colors(col_colors)
        self.col_colors = col_colors

        width_ratios = self.dim_ratios(self.row_colors,
                                       figsize=figsize,
                                       axis=1)

        height_ratios = self.dim_ratios(self.col_colors,
                                        figsize=figsize,
                                        axis=0)
        nrows = 3 if self.col_colors is None else 4
        ncols = 3 if self.row_colors is None else 4

        self.gs = gridspec.GridSpec(nrows, ncols, wspace=0.01, hspace=0.01,
                                    width_ratios=width_ratios,
                                    height_ratios=height_ratios)

        self.ax_row_dendrogram = self.fig.add_subplot(self.gs[nrows - 1, 0:2],
                                                      axisbg="white")
        self.ax_col_dendrogram = self.fig.add_subplot(self.gs[0:2, ncols - 1],
                                                      axisbg="white")

        self.ax_row_colors = None
        self.ax_col_colors = None

        if self.row_colors is not None:
            self.ax_row_colors = self.fig.add_subplot(
                self.gs[nrows - 1, ncols - 2])
        if self.col_colors is not None:
            self.ax_col_colors = self.fig.add_subplot(
                self.gs[nrows - 2, ncols - 1])

        self.ax_heatmap = self.fig.add_subplot(self.gs[nrows - 1, ncols - 1])

        # colorbar for scale to left corner
        self.cax = self.fig.add_subplot(self.gs[0, 0])

        self.dendrogram_row = None
        self.dendrogram_col = None

    def format_data(self, data, pivot_kws, z_score=None,
                    standard_scale=None):
        """Extract variables from data or use directly."""

        # Either the data is already in 2d matrix format, or need to do a pivot
        if pivot_kws is not None:
            data2d = data.pivot(**pivot_kws)
        else:
            data2d = data

        if z_score is not None and standard_scale is not None:
            raise ValueError(
                'Cannot perform both z-scoring and standard-scaling on data')

        if z_score is not None:
            data2d = self.z_score(data2d, z_score)
        if standard_scale is not None:
            data2d = self.standard_scale(data2d, standard_scale)
        return data2d

    @staticmethod
    def z_score(data2d, axis=1):
        """Standarize the mean and variance of the data axis

        Parameters
        ----------
        data2d : pandas.DataFrame
            Data to normalize
        axis : int
            Which axis to normalize across. If 0, normalize across rows, if 1,
            normalize across columns.

        Returns
        -------
        normalized : pandas.DataFrame
            Noramlized data with a mean of 0 and variance of 1 across the
            specified axis.
        """
        if axis == 1:
            z_scored = data2d
        else:
            z_scored = data2d.T

        z_scored = (z_scored - z_scored.mean()) / z_scored.std()

        if axis == 1:
            return z_scored
        else:
            return z_scored.T

    @staticmethod
    def standard_scale(data2d, axis=1):
        """Divide the data by the difference between the max and min

        Parameters
        ----------
        data2d : pandas.DataFrame
            Data to normalize
        axis : int
            Which axis to normalize across. If 0, normalize across rows, if 1,
            normalize across columns.
        vmin : int
            If 0, then subtract the minimum of the data before dividing by
            the range.

        Returns
        -------
        standardized : pandas.DataFrame
            Noramlized data with a mean of 0 and variance of 1 across the
            specified axis.

        >>> import numpy as np
        >>> d = np.arange(5, 8, 0.5)
        >>> ClusterGrid.standard_scale(d)
        array([ 0. ,  0.2,  0.4,  0.6,  0.8,  1. ])
        """
        # Normalize these values to range from 0 to 1
        if axis == 1:
            standardized = data2d
        else:
            standardized = data2d.T

        subtract = standardized.min()
        standardized = (standardized - subtract) / (
            standardized.max() - standardized.min())

        if axis == 1:
            return standardized
        else:
            return standardized.T

    def dim_ratios(self, side_colors, axis, figsize, side_colors_ratio=0.05):
        """Get the proportions of the figure taken up by each axes
        """
        figdim = figsize[axis]
        # Get resizing proportion of this figure for the dendrogram and
        # colorbar, so only the heatmap gets bigger but the dendrogram stays
        # the same size.
        dendrogram = min(2. / figdim, .2)

        # add the colorbar
        colorbar_width = .8 * dendrogram
        colorbar_height = .2 * dendrogram
        if axis == 0:
            ratios = [colorbar_width, colorbar_height]
        else:
            ratios = [colorbar_height, colorbar_width]

        if side_colors is not None:
            # Add room for the colors
            ratios += [side_colors_ratio]

        # Add the ratio for the heatmap itself
        ratios += [.8]

        return ratios

    @staticmethod
    def color_list_to_matrix_and_cmap(colors, ind, axis=0):
        """Turns a list of colors into a numpy matrix and matplotlib colormap

        These arguments can now be plotted using heatmap(matrix, cmap)
        and the provided colors will be plotted.

        Parameters
        ----------
        colors : list of matplotlib colors
            Colors to label the rows or columns of a dataframe.
        ind : list of ints
            Ordering of the rows or columns, to reorder the original colors
            by the clustered dendrogram order
        axis : int
            Which axis this is labeling

        Returns
        -------
        matrix : numpy.array
            A numpy array of integer values, where each corresponds to a color
            from the originally provided list of colors
        cmap : matplotlib.colors.ListedColormap

        """
        # check for nested lists/color palettes.
        # Will fail if matplotlib color is list not tuple
        if any(issubclass(type(x), list) for x in colors):
            all_colors = set(itertools.chain(*colors))
            n = len(colors)
            m = len(colors[0])
        else:
            all_colors = set(colors)
            n = 1
            m = len(colors)
            colors = [colors]
        color_to_value = dict((col, i) for i, col in enumerate(all_colors))

        matrix = np.array([color_to_value[c]
                           for color in colors for c in color])

        shape = (n, m)
        matrix = matrix.reshape(shape)
        matrix = matrix[:, ind]
        if axis == 0:
            # row-side:
            matrix = matrix.T

        cmap = mpl.colors.ListedColormap(all_colors)
        return matrix, cmap

    def savefig(self, *args, **kwargs):
        if 'bbox_inches' not in kwargs:
            kwargs['bbox_inches'] = 'tight'
        self.fig.savefig(*args, **kwargs)

    def plot_dendrograms(self, row_cluster, col_cluster, metric, method,
                         row_linkage, col_linkage):
        # Plot the row dendrogram
        if row_cluster:
            self.dendrogram_row = dendrogram(
                self.data2d, metric=metric, method=method, label=False, axis=0,
                ax=self.ax_row_dendrogram, rotate=True, linkage=row_linkage)
        else:
            self.ax_row_dendrogram.set_xticks([])
            self.ax_row_dendrogram.set_yticks([])
        # PLot the column dendrogram
        if col_cluster:
            self.dendrogram_col = dendrogram(
                self.data2d, metric=metric, method=method, label=False,
                axis=1, ax=self.ax_col_dendrogram, linkage=col_linkage)
        else:
            self.ax_col_dendrogram.set_xticks([])
            self.ax_col_dendrogram.set_yticks([])
        despine(ax=self.ax_row_dendrogram, bottom=True, left=True)
        despine(ax=self.ax_col_dendrogram, bottom=True, left=True)

    def plot_colors(self, xind, yind, **kws):
        """Plots color labels between the dendrogram and the heatmap

        Parameters
        ----------
        heatmap_kws : dict
            Keyword arguments heatmap
        """
        # Remove any custom colormap and centering
        kws = kws.copy()
        kws.pop('cmap', None)
        kws.pop('center', None)
        kws.pop('vmin', None)
        kws.pop('vmax', None)
        kws.pop('xticklabels', None)
        kws.pop('yticklabels', None)
        if self.row_colors is not None:
            matrix, cmap = self.color_list_to_matrix_and_cmap(
                self.row_colors, yind, axis=0)
            heatmap(matrix, cmap=cmap, cbar=False, ax=self.ax_row_colors,
                    xticklabels=False, yticklabels=False,
                    **kws)
        else:
            despine(self.ax_row_colors, left=True, bottom=True)

        if self.col_colors is not None:
            matrix, cmap = self.color_list_to_matrix_and_cmap(
                self.col_colors, xind, axis=1)
            heatmap(matrix, cmap=cmap, cbar=False, ax=self.ax_col_colors,
                    xticklabels=False, yticklabels=False,
                    **kws)
        else:
            despine(self.ax_col_colors, left=True, bottom=True)

    def plot_matrix(self, colorbar_kws, xind, yind, **kws):
        self.data2d = self.data2d.iloc[yind, xind]
        self.mask = self.mask.iloc[yind, xind]

        # Try to reorganize specified tick labels, if provided
        xtl = kws.pop("xticklabels", True)
        try:
            xtl = np.asarray(xtl)[xind]
        except (TypeError, IndexError):
            pass
        ytl = kws.pop("yticklabels", True)
        try:
            ytl = np.asarray(ytl)[yind]
        except (TypeError, IndexError):
            pass

        heatmap(self.data2d, ax=self.ax_heatmap, cbar_ax=self.cax,
                cbar_kws=colorbar_kws, mask=self.mask,
                xticklabels=xtl, yticklabels=ytl, **kws)
        self.ax_heatmap.yaxis.set_ticks_position('right')
        self.ax_heatmap.yaxis.set_label_position('right')

    def plot(self, metric, method, colorbar_kws, row_cluster, col_cluster,
             row_linkage, col_linkage, **kws):
        colorbar_kws = {} if colorbar_kws is None else colorbar_kws
        self.plot_dendrograms(row_cluster, col_cluster, metric, method,
                              row_linkage=row_linkage, col_linkage=col_linkage)
        try:
            xind = self.dendrogram_col.reordered_ind
        except AttributeError:
            xind = np.arange(self.data2d.shape[1])
        try:
            yind = self.dendrogram_row.reordered_ind
        except AttributeError:
            yind = np.arange(self.data2d.shape[0])

        self.plot_colors(xind, yind, **kws)
        self.plot_matrix(colorbar_kws, xind, yind, **kws)
        return self

def clustermap(data, pivot_kws=None, method='average', metric='euclidean',
               z_score=None, standard_scale=None, figsize=None, cbar_kws=None,
               row_cluster=True, col_cluster=True,
               row_linkage=None, col_linkage=None,
               row_colors=None, col_colors=None, mask=None, **kwargs):
    """Plot a hierarchically clustered heatmap of a pandas DataFrame

    Parameters
    ----------
    data: pandas.DataFrame
        Rectangular data for clustering. Cannot contain NAs.
    pivot_kws : dict, optional
        If `data` is a tidy dataframe, can provide keyword arguments for
        pivot to create a rectangular dataframe.
    method : str, optional
        Linkage method to use for calculating clusters.
        See scipy.cluster.hierarchy.linkage documentation for more information:
        http://docs.scipy.org/doc/scipy/reference/generated/scipy.cluster.hierarchy.linkage.html
    metric : str, optional
        Distance metric to use for the data. See
        scipy.spatial.distance.pdist documentation for more options
        http://docs.scipy.org/doc/scipy/reference/generated/scipy.spatial.distance.pdist.html
    z_score : int or None, optional
        Either 0 (rows) or 1 (columns). Whether or not to calculate z-scores
        for the rows or the columns. Z scores are: z = (x - mean)/std, so
        values in each row (column) will get the mean of the row (column)
        subtracted, then divided by the standard deviation of the row (column).
        This ensures that each row (column) has mean of 0 and variance of 1.
    standard_scale : int or None, optional
        Either 0 (rows) or 1 (columns). Whether or not to standardize that
        dimension, meaning for each row or column, subtract the minimum and
        divide each by its maximum.
    figsize: tuple of two ints, optional
        Size of the figure to create.
    cbar_kws : dict, optional
        Keyword arguments to pass to ``cbar_kws`` in ``heatmap``, e.g. to
        add a label to the colorbar.
    {row,col}_cluster : bool, optional
        If True, cluster the {rows, columns}.
    {row,col}_linkage : numpy.array, optional
        Precomputed linkage matrix for the rows or columns. See
        scipy.cluster.hierarchy.linkage for specific formats.
    {row,col}_colors : list-like, optional
        List of colors to label for either the rows or columns. Useful to
        evaluate whether samples within a group are clustered together. Can
        use nested lists for multiple color levels of labeling.
    mask : boolean array or DataFrame, optional
        If passed, data will not be shown in cells where ``mask`` is True.
        Cells with missing values are automatically masked. Only used for
        visualizing, not for calculating.
    kwargs : other keyword arguments
        All other keyword arguments are passed to ``sns.heatmap``

    Returns
    -------
    clustergrid : ClusterGrid
        A ClusterGrid instance.

    Notes
    -----
    The returned object has a ``savefig`` method that should be used if you
    want to save the figure object without clipping the dendrograms.

    To access the reordered row indices, use:
    ``clustergrid.dendrogram_row.reordered_ind``

    Column indices, use:
    ``clustergrid.dendrogram_col.reordered_ind``

    Examples
    --------

    Plot a clustered heatmap:

    .. plot::
        :context: close-figs

        >>> import seaborn as sns; sns.set()
        >>> flights = sns.load_dataset("flights")
        >>> flights = flights.pivot("month", "year", "passengers")
        >>> g = sns.clustermap(flights)

    Don't cluster one of the axes:

    .. plot::
        :context: close-figs

        >>> g = sns.clustermap(flights, col_cluster=False)

    Use a different colormap and add lines to separate the cells:

    .. plot::
        :context: close-figs

        >>> cmap = sns.cubehelix_palette(as_cmap=True, rot=-.3, light=1)
        >>> g = sns.clustermap(flights, cmap=cmap, linewidths=.5)

    Use a different figure size:

    .. plot::
        :context: close-figs

        >>> g = sns.clustermap(flights, cmap=cmap, figsize=(7, 5))

    Standardize the data across the columns:

    .. plot::
        :context: close-figs

        >>> g = sns.clustermap(flights, standard_scale=1)

    Normalize the data across the rows:

    .. plot::
        :context: close-figs

        >>> g = sns.clustermap(flights, z_score=0)

    Use a different clustering method:

    .. plot::
        :context: close-figs

        >>> g = sns.clustermap(flights, method="single", metric="cosine")

    Add colored labels on one of the axes:

    .. plot::
        :context: close-figs

        >>> season_colors = (sns.color_palette("BuPu", 3) +
        ...                  sns.color_palette("RdPu", 3) +
        ...                  sns.color_palette("YlGn", 3) +
        ...                  sns.color_palette("OrRd", 3))
        >>> g = sns.clustermap(flights, row_colors=season_colors)

    """
    plotter = ClusterGrid(data, pivot_kws=pivot_kws, figsize=figsize,
                          row_colors=row_colors, col_colors=col_colors,
                          z_score=z_score, standard_scale=standard_scale,
                          mask=mask)

    return plotter.plot(metric=metric, method=method,
                        colorbar_kws=cbar_kws,
                        row_cluster=row_cluster, col_cluster=col_cluster,
                        row_linkage=row_linkage, col_linkage=col_linkage,
                        **kwargs)


import os, sys, time
from glob import glob

import cv2

from pylab import *
from mpl_toolkits.mplot3d import Axes3D
from matplotlib.backends.backend_pdf import PdfPages
matplotlib.rcParams['figure.facecolor'] = 'w'

from scipy.signal import argrelextrema
import scipy.stats as stats
import scipy.io as sio
from scipy import signal

from xlwt import Workbook

# specify these in mm to match your behavior chamber.
CHMAMBER_LENGTH=235
WATER_HIGHT=40

# quick plot should also show xy_within and location_one_third etc
# summary PDF: handle exception when a pickle file missing some fish in other pickle file

## these three taken from http://stackoverflow.com/a/18420730/566035
def strided_sliding_std_dev(data, radius=5):
    windowed = rolling_window(data, (2*radius, 2*radius))
    shape = windowed.shape
    windowed = windowed.reshape(shape[0], shape[1], -1)
    return windowed.std(axis=-1)

def rolling_window(a, window):
    """Takes a numpy array *a* and a sequence of (or single) *window* lengths
    and returns a view of *a* that represents a moving window."""
    if not hasattr(window, '__iter__'):
        return rolling_window_lastaxis(a, window)
    for i, win in enumerate(window):
        if win > 1:
            a = a.swapaxes(i, -1)
            a = rolling_window_lastaxis(a, win)
            a = a.swapaxes(-2, i)
    return a

def rolling_window_lastaxis(a, window):
    """Directly taken from Erik Rigtorp's post to numpy-discussion.
    <http://www.mail-archive.com/numpy-discussion@scipy.org/msg29450.html>"""
    if window < 1:
       raise ValueError, "`window` must be at least 1."
    if window > a.shape[-1]:
       raise ValueError, "`window` is too long."
    shape = a.shape[:-1] + (a.shape[-1] - window + 1, window)
    strides = a.strides + (a.strides[-1],)
    return np.lib.stride_tricks.as_strided(a, shape=shape, strides=strides)
## stealing ends here... //

def filterheadxy(headx,heady,thrs_denom=10):

    b, a = signal.butter(8, 0.125)
    dhy = np.abs(np.hstack((0, np.diff(heady,1))))
    thrs = np.nanstd(dhy)/thrs_denom
    ind2remove = dhy>thrs
    headx[ind2remove] = np.nan
    heady[ind2remove] = np.nan

    headx = interp_nan(headx)
    heady = interp_nan(heady)
    headx = signal.filtfilt(b, a, headx, padlen=150)
    heady = signal.filtfilt(b, a, heady, padlen=150)

    return headx,heady

def smoothRad(theta, thrs=np.pi/4*3):
    jumps = (np.diff(theta) > thrs).nonzero()[0]
    print 'jumps.size', jumps.size
    while jumps.size:
        # print '%d/%d' % (jumps[0], theta.size)
        theta[jumps+1] -= np.pi
        jumps = (np.diff(theta) > thrs).nonzero()[0]

    return theta

def datadct2array(data, key1, key2):

    # put these in a MATLAB CELL
    trialN = len(data[key1][key2])
    matchedUSnameP = np.zeros((trialN,), dtype=np.object)
    fnameP = np.zeros((trialN,), dtype=np.object)
    # others to append to a list
    eventsP = []
    speed3DP = []
    movingSTDP = []
    d2inflowP = []
    xP, yP, zP = [], [], []
    XP, YP, ZP = [], [], []
    ringpixelsP = []
    peaks_withinP = []
    swimdir_withinP = []
    xy_withinP = []
    location_one_thirdP = []
    dtheta_shapeP = []
    dtheta_velP = []
    turns_shapeP = []
    turns_velP = []

    for n, dct in enumerate(data[key1][key2]):
        # MATLAB CELL
        matchedUSnameP[n] = dct['matchedUSname']
        fnameP[n] = dct['fname']

        # 2D array
        eventsP.append([ele if type(ele) is not list else ele[0] for ele in dct['events']])
        speed3DP.append(dct['speed3D'])
        movingSTDP.append(dct['movingSTD'])
        d2inflowP.append(dct['d2inflow'])
        xP.append(dct['x'])
        yP.append(dct['y'])
        zP.append(dct['z'])
        XP.append(dct['X'])
        YP.append(dct['Y'])
        ZP.append(dct['Z'])
        ringpixelsP.append(dct['ringpixels'])
        peaks_withinP.append(dct['peaks_within'])
        swimdir_withinP.append(dct['swimdir_within'])
        xy_withinP.append(dct['xy_within'])
        location_one_thirdP.append(dct['location_one_third'])
        dtheta_shapeP.append(dct['dtheta_shape'])
        dtheta_velP.append(dct['dtheta_vel'])
        turns_shapeP.append(dct['turns_shape'])
        turns_velP.append(dct['turns_vel'])

    TVroi = np.array(dct['TVroi'])
    SVroi = np.array(dct['SVroi'])

    return matchedUSnameP, fnameP, np.array(eventsP), np.array(speed3DP), np.array(d2inflowP), \
            np.array(xP), np.array(yP), np.array(zP), np.array(XP), np.array(YP), np.array(ZP), \
            np.array(ringpixelsP), np.array(peaks_withinP), np.array(swimdir_withinP), \
            np.array(xy_withinP), np.array(dtheta_shapeP), np.array(dtheta_velP), \
            np.array(turns_shapeP), np.array(turns_velP), TVroi, SVroi

def pickle2mat(fp, data=None):
    # fp : full path to pickle file
    # data : option to provide data to skip np.load(fp)
    if not data:
        data = np.load(fp)

    for key1 in data.keys():
        for key2 in data[key1].keys():

            matchedUSname, fname, events, speed3D, d2inflow, x, y, z, X, Y, Z, \
            ringpixels, peaks_within, swimdir_within, xy_within, dtheta_shape, dtheta_vel, \
            turns_shape, turns_vel, TVroi, SVroi = datadct2array(data, key1, key2)

            datadict = {
                'matchedUSname' : matchedUSname,
                'fname' : fname,
                'events' : events,
                'speed3D' : speed3D,
                'd2inflow' : d2inflow,
                'x' : x,
                'y' : y,
                'z' : z,
                'X' : X,
                'Y' : Y,
                'Z' : Z,
                'ringpixels' : ringpixels,
                'peaks_within' : peaks_within,
                'swimdir_within' : swimdir_within,
                'xy_within' : xy_within,

               'dtheta_shape' : dtheta_shape,
               'dtheta_vel' : dtheta_vel,
               'turns_shape' : turns_shape,
               'turns_vel' : turns_vel,

                'TVroi' : TVroi,
                'SVroi' : SVroi,
            }
            outfp = '%s_%s_%s.mat' % (fp[:-7],key1,key2)
            sio.savemat(outfp, datadict, oned_as='row', do_compression=True)

def interp_nan(x):
    '''
    Replace nan by interporation
    http://stackoverflow.com/questions/6518811/interpolate-nan-values-in-a-numpy-array
    '''
    ok = -np.isnan(x)
    if (ok == False).all():
        return x
    else:
        xp = ok.ravel().nonzero()[0]
        fp = x[ok]
        _x = np.isnan(x).ravel().nonzero()[0]
        x[-ok] = np.interp(_x, xp, fp)
        return x

def polytest(x,y,rx,ry,rw,rh,rang):
    points=cv2.ellipse2Poly(
            (rx,ry),
            axes=(rw/2,rh/2),
            angle=rang,
            arcStart=0,
            arcEnd=360,
            delta=3
        )
    return cv2.pointPolygonTest(np.array(points), (x,y), measureDist=1)

def depthCorrection(z,x,TVx1,TVx2,SVy1,SVy2,SVy3):
    z0 = z - SVy1
    x0 = x - TVx1
    mid = (SVy2-SVy1)/2
    adj = (z0 - mid) / (SVy2-SVy1) * (SVy2-SVy3) * (1-(x0)/float(TVx2-TVx1))
    return z0 + adj + SVy1 # back to abs coord

def putNp2xls(array, ws):
    for r, row in enumerate(array):
        for c, val in enumerate(row):
            ws.write(r, c, val)

def drawLines(mi, ma, events, fps=30.0):

    CS, USs, preRange = events

    plot([CS-preRange, CS-preRange], [mi,ma], '--c')                    # 2 min prior odor
    plot([CS         , CS         ], [mi,ma], '--g', linewidth=2)       # CS onset
    if USs:
        if len(USs) > 3:
            colors = 'r' * len(USs)
        else:
            colors = [_ for _ in ['r','b','c'][:len(USs)]]
        for c,us in zip(colors, USs):
            plot([us, us],[mi,ma], linestyle='--', color=c, linewidth=2)            # US onset
        plot([USs[0]+preRange/2,USs[0]+preRange/2], [mi,ma], linestyle='--', color=c, linewidth=2)    # end of US window
        xtck = np.arange(0, max(CS+preRange, max(USs)), 0.5*60*fps) # every 0.5 min tick
    else:
        xtck = np.arange(0, CS+preRange, 0.5*60*fps) # every 0.5 min tick

    xticks(xtck, xtck/fps/60)
    gca().xaxis.set_minor_locator(MultipleLocator(5*fps)) # 5 s minor ticks

def approachevents(x,y,z, ringpolyTVArray, ringpolySVArray, fishlength=134, thrs=None):
    '''
    fishlength: some old scrits may call this with fishlength
    thrs: multitrack GUI provides this by ringAppearochLevel spin control.
          can be an numpy array (to track water level change etc)
    '''

    smoothedz = np.convolve(np.hanning(10)/np.hanning(10).sum(), z, 'same')
    peaks = argrelextrema(smoothedz, np.less)[0] # less because 0 is top in image.

    # now filter peaks by height.
    ringLevel = ringpolySVArray[:,1]
    if thrs is None:
        thrs = ringLevel+fishlength/2

    if type(thrs) == int:  # can be numpy array or int
        thrs = ringLevel.mean() + thrs
        peaks = peaks[ z[peaks] < thrs ]
    else: # numpy array should be ready to use
        peaks = peaks[ z[peaks] < thrs[peaks] ]

    # now filter out by TVringCenter
    peaks_within = get_withinring(ringpolyTVArray, peaks, x, y)

    return smoothedz, peaks_within

def get_withinring(ringpolyTVArray, timepoints, x, y):

    rx = ringpolyTVArray[:,0].astype(np.int)
    ry = ringpolyTVArray[:,1].astype(np.int)
    rw = ringpolyTVArray[:,2].astype(np.int)
    rh = ringpolyTVArray[:,3].astype(np.int)
    rang = ringpolyTVArray[:,4].astype(np.int)
    # poly test
    peaks_within = []
    for p in timepoints:
        points=cv2.ellipse2Poly(
                (rx[p],ry[p]),
                axes=(rw[p]/2,rh[p]/2),
                angle=rang[p],
                arcStart=0,
                arcEnd=360,
                delta=3
            )
        inout = cv2.pointPolygonTest(np.array(points), (x[p],y[p]), measureDist=1)
        if inout > 0:
            peaks_within.append(p)

    return peaks_within

def location_ring(x,y,ringpolyTVArray):

    rx = ringpolyTVArray[:,0].astype(np.int)
    ry = ringpolyTVArray[:,1].astype(np.int)
    rw = ringpolyTVArray[:,2].astype(np.int)
    rh = ringpolyTVArray[:,3].astype(np.int)
    d2ringcenter = np.sqrt((x-rx)**2 + (y-ry)**2)

    # filter by radius 20% buffer in case the ring moves around
    indices = (d2ringcenter < 1.2*max(rw.max(), rh.max())).nonzero()[0]
    xy_within = get_withinring(ringpolyTVArray, indices, x, y)

    return xy_within

def swimdir_analysis(x,y,z,ringpolyTVArray,ringpolySVArray,TVx1,TVy1,TVx2,TVy2,fps=30.0):
    # smoothing
    # z = np.convolve(np.hanning(16)/np.hanning(16).sum(), z, 'same')

    # two cameras have different zoom settings. So, distance per pixel is different. But, for
    # swim direction, it does not matter how much x,y are compressed relative to z.

    # ring z level from SV
    rz = ringpolySVArray[:,1].astype(np.int)
    # ring all other params from TV
    rx = ringpolyTVArray[:,0].astype(np.int)
    ry = ringpolyTVArray[:,1].astype(np.int)
    rw = ringpolyTVArray[:,2].astype(np.int)
    rh = ringpolyTVArray[:,3].astype(np.int)
    rang = ringpolyTVArray[:,4].astype(np.int)

    speed3D = np.sqrt( np.diff(x)**2 + np.diff(y)**2 + np.diff(z)**2 )
    speed3D = np.hstack(([0], speed3D))

    # line in 3D http://tutorial.math.lamar.edu/Classes/CalcIII/EqnsOfLines.aspx
    # x-x0   y-y0   z-z0
    # ---- = ---- = ----
    #   a      b      c
    # solve them for z = rz. x0,y0,z0 are tvx, tvy, svy
    # x = (a * (rz-z)) / c + x0

    dt = 3 # define slope as diff between current and dt frame before
    a = np.hstack( (np.ones(dt), x[dt:]-x[:-dt]) )
    b = np.hstack( (np.ones(dt), y[dt:]-y[:-dt]) )
    c = np.hstack( (np.ones(dt), z[dt:]-z[:-dt]) )
    c[c==0] = np.nan # avoid zero division

    water_x = (a * (rz-z) / c) + x
    water_y = (b * (rz-z) / c) + y

    upwards = c<-2/30.0*fps # not accurate when c is small or negative
    xok = (TVx1 < water_x) & (water_x < TVx2)
    yok = (TVy1 < water_y) & (water_y < TVy2)

    filtered = upwards & xok & yok# & -np.isinf(water_x) & -np.isinf(water_y)

    water_x[-filtered] = np.nan
    water_y[-filtered] = np.nan

    # figure()
    # ax = subplot(111)
    # ax.imshow(npData['TVbg'], cmap=cm.gray) # clip out from TVx1,TVy1
    # ax.plot(x-TVx1, y-TVy1, 'c')
    # ax.plot(water_x-TVx1, water_y-TVy1, 'r.')
    # xlim([0, TVx2-TVx1]); ylim([TVy2-TVy1, 0])
    # draw(); show()

    SwimDir = []
    for n in filtered.nonzero()[0]:
        inout = polytest(water_x[n],water_y[n],rx[n],ry[n],rw[n],rh[n],rang[n])
        SwimDir.append((n, inout, speed3D[n])) # inout>0 are inside

    return SwimDir, water_x, water_y

def plot_eachTr(events, x, y, z, inflowpos, ringpixels, peaks_within, swimdir_within=None,
    pp=None, _title=None, fps=30.0, inmm=False):

    CS, USs, preRange = events
    # preRange = 3600 2 min prior and 1 min after CS. +900 for 0.5 min
    if USs:
        xmin, xmax = CS-preRange-10*fps, USs[0]+preRange/2+10*fps
    else:
        xmin, xmax = CS-preRange-10*fps, CS+preRange/2+(23+10)*fps

    fig = figure(figsize=(12,8), facecolor='w')

    subplot(511)        # Swimming speed
    speed3D = np.sqrt( np.diff(x)**2 + np.diff(y)**2 + np.diff(z)**2 )
    drawLines(np.nanmin(speed3D), np.nanmax(speed3D), events, fps) # go behind
    plot(speed3D)

    movingSTD = np.append( np.zeros(fps*10), strided_sliding_std_dev(speed3D, fps*10) )
    plot(movingSTD, linewidth=2)

    plot(np.ones_like(speed3D) * speed3D.std()*6, '-.', color='gray')
    ylim([-5, speed3D[xmin:xmax].max()])
    xlim([xmin,xmax]); title(_title)
    if inmm:
        ylabel('Speed 3D (mm),\n6SD thr');
    else:
        ylabel('Speed 3D, 6SD thr');

    ax = subplot(512)   # z level
    drawLines(z.min(), z.max(), events)
    plot(z, 'b')

    pkx = peaks_within.nonzero()[0]
    if inmm:
        plot(pkx, peaks_within[pkx]*z[xmin:xmax].max()*0.97, 'mo')
        if swimdir_within is not None:
            ___x = swimdir_within.nonzero()[0]
            plot(___x, swimdir_within[___x]*z[xmin:xmax].max()*0.96, 'g+')
        ylim([z[xmin:xmax].min()*0.95, z[xmin:xmax].max()])
        xlim([xmin,xmax]); ylabel('Z (mm)')
    else:
        plot(pkx, peaks_within[pkx]*z[xmin:xmax].min()*0.97, 'mo')
        if swimdir_within is not None:
            ___x = swimdir_within.nonzero()[0]
            plot(___x, swimdir_within[___x]*z[xmin:xmax].min()*0.96, 'g+')
        ylim([z[xmin:xmax].min()*0.95, z[xmin:xmax].max()])
        ax.invert_yaxis(); xlim([xmin,xmax]); ylabel('z')

    subplot(513)        # x
    drawLines(x.min(), x.max(), events)
    plot(x, 'b')
    plot(y, 'g')
    xlim([xmin,xmax]); ylabel('x,y')

    subplot(514)        # Distance to the inflow tube
    xin, yin, zin = inflowpos
    d2inflow = np.sqrt((x-xin) ** 2 + (y-yin) ** 2 + (z-zin) ** 2 )
    drawLines(d2inflow.min(), d2inflow.max(), events)
    plot(d2inflow)
    ylim([d2inflow[xmin:xmax].min(), d2inflow[xmin:xmax].max()])
    xlim([xmin,xmax]); ylabel('distance to\ninflow tube')

    subplot(515)        # ringpixels: it seems i never considered TV x,y for this
    rpmax, rpmin = np.nanmax(ringpixels[xmin:xmax]), np.nanmin(ringpixels[xmin:xmax])
    drawLines(rpmin, rpmax, events)
    plot(ringpixels)
    plot(pkx, peaks_within[pkx]*rpmax*1.06, 'mo')
    if swimdir_within is not None:
        plot(___x, swimdir_within[___x]*rpmax*1.15, 'g+')
    ylim([-100, rpmax*1.2])
    xlim([xmin,xmax]); ylabel('ringpixels')

    tight_layout()

    if pp:
        fig.savefig(pp, format='pdf')

    rng = np.arange(CS-preRange, CS+preRange, dtype=np.int)

    return speed3D[rng], movingSTD[rng], d2inflow[rng], ringpixels[rng]

def plot_turnrates(events, dthetasum_shape,dthetasum_vel,turns_shape,turns_vel,
            pp=None, _title=None, thrs=np.pi/4*(133.33333333333334/120), fps=30.0):

    CS, USs, preRange = events
    # preRange = 3600 2 min prior and 1 min after CS. +900 for 0.5 min
    if USs:
        xmin, xmax = CS-preRange-10*fps, USs[0]+preRange/2+10*fps
    else:
        xmin, xmax = CS-preRange-10*fps, CS+preRange/2+(23+10)*fps

    fig = figure(figsize=(12,8), facecolor='w')

    subplot(211)
    drawLines(dthetasum_shape.min(), dthetasum_shape.max(), events)
    plot(np.ones_like(dthetasum_shape)*thrs,'gray',linestyle='--')
    plot(-np.ones_like(dthetasum_shape)*thrs,'gray',linestyle='--')
    plot(dthetasum_shape)
    dmax = dthetasum_shape[xmin:xmax].max()
    plot(turns_shape, (0.5+dmax)*np.ones_like(turns_shape), 'o')
    temp = np.zeros_like(dthetasum_shape)
    temp[turns_shape] = 1
    shape_cumsum = np.cumsum(temp)
    shape_cumsum -= shape_cumsum[xmin]
    plot( shape_cumsum / shape_cumsum[xmax] * (dmax-dthetasum_shape.min()) + dthetasum_shape.min())
    xlim([xmin,xmax]); ylabel('Shape based'); title('Orientation change per 4 frames: ' + _title)
    ylim([dthetasum_shape[xmin:xmax].min()-1, dmax+1])

    subplot(212)
    drawLines(dthetasum_vel.min(), dthetasum_vel.max(), events)
    plot(np.ones_like(dthetasum_vel)*thrs,'gray',linestyle='--')
    plot(-np.ones_like(dthetasum_vel)*thrs,'gray',linestyle='--')
    plot(dthetasum_vel)
    dmax = dthetasum_vel[xmin:xmax].max()
    plot(turns_vel, (0.5+dmax)*np.ones_like(turns_vel), 'o')
    temp = np.zeros_like(dthetasum_vel)
    temp[turns_vel] = 1
    vel_cumsum = np.cumsum(temp)
    vel_cumsum -= vel_cumsum[xmin]
    plot( vel_cumsum / vel_cumsum[xmax] * (dmax-dthetasum_shape.min()) + dthetasum_shape.min())

    ylim([dthetasum_vel[xmin:xmax].min()-1, dmax+1])
    xlim([xmin,xmax]); ylabel('Velocity based')

    tight_layout()

    if pp:
        fig.savefig(pp, format='pdf')

def trajectory(x, y, z, rng, ax, _xlim=[0,640], _ylim=[480,480+300], _zlim=[150,340],
                    color='b', fps=30.0, ringpolygon=None):
    ax.plot(x[rng],y[rng],z[rng], color=color)
    ax.view_init(azim=-75, elev=-180+15)
    if ringpolygon:
        rx, ry, rz = ringpolygon
        ax.plot(rx, ry, rz, color='gray')
    ax.set_xlim(_xlim[0],_xlim[1])
    ax.set_ylim(_ylim[0],_ylim[1])
    ax.set_zlim(_zlim[0],_zlim[1])
    title(("(%2.1f min to %2.1f min)" % (rng[0]/fps/60.0,(rng[-1]+1)/60.0/fps)))
    draw()

def plotTrajectory(x, y, z, events, _xlim=None, _ylim=None, _zlim=None, fps=30.0, pp=None, ringpolygon=None):
    CS, USs, preRange = events
    rng1 = np.arange(CS-preRange, CS-preRange/2, dtype=int)
    rng2 = np.arange(CS-preRange/2, CS, dtype=int)
    if USs:
        rng3 = np.arange(CS, min(USs), dtype=int)
        rng4 = np.arange(min(USs), min(USs)+preRange/2, dtype=int)
        combined = np.hstack((rng1,rng2,rng3,rng4))
    else:
        combined = np.hstack((rng1,rng2))

    if _xlim is None:
        _xlim = map( int, ( x[combined].min(), x[combined].max() ) )
    if _ylim is None:
        _ylim = map( int, ( y[combined].min(), y[combined].max() ) )
    if _zlim is None:
        _zlim = map( int, ( z[combined].min(), z[combined].max() ) )
        if ringpolygon:
            _zlim[0] = min( _zlim[0], int(ringpolygon[2][0]) )

    fig3D = plt.figure(figsize=(12,8), facecolor='w')
    ax = fig3D.add_subplot(221, projection='3d'); trajectory(x,y,z,rng1,ax,_xlim,_ylim,_zlim,'c',fps,ringpolygon)
    ax = fig3D.add_subplot(222, projection='3d'); trajectory(x,y,z,rng2,ax,_xlim,_ylim,_zlim,'c',fps,ringpolygon)
    if USs:
        ax = fig3D.add_subplot(223, projection='3d'); trajectory(x,y,z,rng3,ax,_xlim,_ylim,_zlim,'g',fps,ringpolygon)
        ax = fig3D.add_subplot(224, projection='3d'); trajectory(x,y,z,rng4,ax,_xlim,_ylim,_zlim,'r',fps,ringpolygon)
    tight_layout()

    if pp:
        fig3D.savefig(pp, format='pdf')

def add2DataAndPlot(fp, fish, data, createPDF):

    if createPDF:
        pp = PdfPages(fp[:-7]+'_'+fish+'.pdf')
    else:
        pp = None

    params = np.load(fp)
    fname = os.path.basename(fp).split('.')[0] + '.avi'
    dirname = os.path.dirname(fp)
    preRange = params[(fname, 'mog')]['preRange']
    fps = params[(fname, 'mog')]['fps']
    TVx1 = params[(fname, fish)]['TVx1']
    TVy1 = params[(fname, fish)]['TVy1']
    TVx2 = params[(fname, fish)]['TVx2']
    TVy2 = params[(fname, fish)]['TVy2']

    SVx1 = params[(fname, fish)]['SVx1']
    SVx2 = params[(fname, fish)]['SVx2']
    SVx3 = params[(fname, fish)]['SVx3']
    SVy1 = params[(fname, fish)]['SVy1']
    SVy2 = params[(fname, fish)]['SVy2']
    SVy3 = params[(fname, fish)]['SVy3']
    ringAppearochLevel = params[(fname, fish)]['ringAppearochLevel']

    _npz = os.path.join(dirname, os.path.join('%s_%s.npz' % (fname[:-4], fish)))
    # if os.path.exists(_npz):
    npData = np.load(_npz)
    tvx = npData['TVtracking'][:,0] # x with nan
    tvy = npData['TVtracking'][:,1] # y
    headx = npData['TVtracking'][:,3] # headx
    heady = npData['TVtracking'][:,4] # heady
    svy = npData['SVtracking'][:,1] # z
    InflowTubeTVArray = npData['InflowTubeTVArray']
    InflowTubeSVArray = npData['InflowTubeSVArray']
    inflowpos = InflowTubeTVArray[:,0], InflowTubeTVArray[:,1], InflowTubeSVArray[:,1]
    ringpixels = npData['ringpixel']
    ringpolyTVArray = npData['ringpolyTVArray']
    ringpolySVArray = npData['ringpolySVArray']
    TVbg = npData['TVbg']
    print os.path.basename(_npz), 'loaded.'

    x,y,z = map(interp_nan, [tvx,tvy,svy])

    # z level correction by depth (x)
    z = depthCorrection(z,x,TVx1,TVx2,SVy1,SVy2,SVy3)

    smoothedz, peaks_within = approachevents(x, y, z,
        ringpolyTVArray, ringpolySVArray, thrs=ringAppearochLevel)
    # convert to numpy array from list
    temp = np.zeros_like(x)
    temp[peaks_within] = 1
    peaks_within = temp

    # normalize to mm
    longaxis = float(max((TVx2-TVx1), (TVy2-TVy1))) # before rotation H is applied they are orthogonal
    waterlevel = float(SVy2-SVy1)
    X = (x-TVx1) / longaxis * CHMAMBER_LENGTH
    Y = (TVy2-y) / longaxis * CHMAMBER_LENGTH
    Z = (SVy2-z) / waterlevel * WATER_HIGHT  # bottom of chamber = 0, higher more positive
    inflowpos_mm = ((inflowpos[0]-TVx1) / longaxis * CHMAMBER_LENGTH,
                    (TVy2-inflowpos[1]) / longaxis * CHMAMBER_LENGTH,
                    (SVy2-inflowpos[2]) / waterlevel * WATER_HIGHT )

    # do the swim direction analysis here
    swimdir, water_x, water_y = swimdir_analysis(x,y,z,
                ringpolyTVArray,ringpolySVArray,TVx1,TVy1,TVx2,TVy2,fps)
    # all of swimdir are within ROI (frame#, inout, speed) but not necessary within ring
    sdir = np.array(swimdir)
    withinRing = sdir[:,1]>0 # inout>0 are inside ring
    temp = np.zeros_like(x)
    temp[ sdir[withinRing,0].astype(int) ] = 1
    swimdir_within = temp

    # location_ring
    xy_within = location_ring(x,y, ringpolyTVArray)
    temp = np.zeros_like(x)
    temp[xy_within] = 1
    xy_within = temp

    # location_one_third
    if (TVx2-TVx1) > (TVy2-TVy1):
        if np.abs(np.arange(TVx1, longaxis+TVx1, longaxis/3) + longaxis/6 - inflowpos[0].mean()).argmin() == 2:
            location_one_third = x-TVx1 > longaxis/3*2
        else:
            location_one_third = x < longaxis/3
    else:
        if np.abs(np.arange(TVy1, longaxis+TVy1, longaxis/3) + longaxis/6 - inflowpos[1].mean()).argmin() == 2:
            location_one_third = y-TVy1 > longaxis/3*2
        else:
            location_one_third = y < longaxis/3

    # turn rate analysis (shape based)
    heady, headx = map(interp_nan, [heady, headx])
    headx, heady = filterheadxy(headx, heady)
    dy = heady - y
    dx = headx - x
    theta_shape = np.arctan2(dy, dx)

    # velocity based
    cx, cy = filterheadxy(x.copy(), y.copy())  # centroid x,y

    vx = np.append(0, np.diff(cx))
    vy = np.append(0, np.diff(cy))
    theta_vel = np.arctan2(vy, vx)

    # prepare ringpolygon for trajectory plot
    rx, ry, rw, rh, rang = ringpolyTVArray.mean(axis=0).astype(int) # use mm ver above
    rz = ringpolySVArray.mean(axis=0)[1].astype(int)

    RX = (rx-TVx1) / longaxis * CHMAMBER_LENGTH
    RY = (TVy2-ry) / longaxis * CHMAMBER_LENGTH
    RW = rw / longaxis * CHMAMBER_LENGTH / 2
    RH = rh / longaxis * CHMAMBER_LENGTH / 2
    RZ = (SVy2-rz) / waterlevel * WATER_HIGHT
    points = cv2.ellipse2Poly(
            (RX.astype(int),RY.astype(int)),
            axes=(RW.astype(int),RH.astype(int)),
            angle=rang,
            arcStart=0,
            arcEnd=360,
            delta=3
        )
    ringpolygon = [points[:,0], points[:,1], np.ones(points.shape[0]) * RZ]

    eventTypeKeys = params[(fname, fish)]['EventData'].keys()
    CSs = [_ for _ in eventTypeKeys if _.startswith('CS')]
    USs = [_ for _ in eventTypeKeys if _.startswith('US')]
    # print CSs, USs

    # events
    for CS in CSs:
        CS_Timings = params[(fname, fish)]['EventData'][CS]
        CS_Timings.sort()
        # initialize when needed
        if CS not in data[fish].keys():
            data[fish][CS] = []

        # now look around for US after it within preRange
        for t in CS_Timings:
            tr = len(data[fish][CS])+1
            rng = np.arange(t-preRange, t+preRange, dtype=np.int)

            matchedUSname = None
            for us in USs:
                us_Timings = params[(fname, fish)]['EventData'][us]
                matched = [_ for _ in us_Timings if t-preRange < _ < t+preRange]
                if matched:
                    events = [t, matched, preRange] # ex. CS+
                    matchedUSname = us
                    break
                else:
                    continue

            _title = '(%s, %s) trial#%02d %s (%s)' % (CS, matchedUSname[0], tr, fname, fish)
            print _title, events

            _speed3D, _movingSTD, _d2inflow, _ringpixels = plot_eachTr(events, X, Y, Z, inflowpos_mm,
                                    ringpixels, peaks_within, swimdir_within, pp, _title, fps, inmm=True)

            # 3d trajectory
            _xlim = (0, CHMAMBER_LENGTH)
            _zlim = (RZ.max(),0)
            plotTrajectory(X, Y, Z, events, _xlim=_xlim, _zlim=_zlim, fps=fps, pp=pp, ringpolygon=ringpolygon)

            # turn rate analysis
            # shape based
            theta_shape[rng] = smoothRad(theta_shape[rng].copy(), thrs=np.pi/2)
            dtheta_shape = np.append(0, np.diff(theta_shape)) # full length

            kernel = np.ones(4)
            dthetasum_shape = np.convolve(dtheta_shape, kernel, 'same')

            # 4 frames = 1000/30.0*4 = 133.3 ms
            thrs = (np.pi / 2) * (133.33333333333334/120) # Braubach et al 2009 90 degree in 120 ms
            peaks_shape = argrelextrema(abs(dthetasum_shape), np.greater)[0]
            turns_shape = peaks_shape[ (abs(dthetasum_shape[peaks_shape]) > thrs).nonzero()[0] ]

            # velocity based
            theta_vel[rng] = smoothRad(theta_vel[rng].copy(), thrs=np.pi/2)
            dtheta_vel = np.append(0, np.diff(theta_vel))

            dthetasum_vel = np.convolve(dtheta_vel, kernel, 'same')

            peaks_vel = argrelextrema(abs(dthetasum_vel), np.greater)[0]
            turns_vel = peaks_vel[ (abs(dthetasum_vel[peaks_vel]) > thrs).nonzero()[0] ]

            plot_turnrates(events, dthetasum_shape, dthetasum_vel, turns_shape, turns_vel, pp, _title, fps=fps)

            _temp = np.zeros_like(dtheta_shape)
            _temp[turns_shape] = 1
            turns_shape_array = _temp

            _temp = np.zeros_like(dtheta_vel)
            _temp[turns_vel] = 1
            turns_vel_array = _temp

            # plot swim direction analysis
            fig = figure(figsize=(12,8), facecolor='w')
            ax1 = subplot(211)
            ax1.imshow(TVbg, cmap=cm.gray) # TVbg is clip out of ROI
            ax1.plot(x[rng]-TVx1, y[rng]-TVy1, 'gray')
            ax1.plot(water_x[t-preRange:t]-TVx1, water_y[t-preRange:t]-TVy1, 'c.')
            if matched:
                ax1.plot(   water_x[t:matched[0]]-TVx1,
                            water_y[t:matched[0]]-TVy1, 'g.')
                ax1.plot(   water_x[matched[0]:matched[0]+preRange/4]-TVx1,
                            water_y[matched[0]:matched[0]+preRange/4]-TVy1, 'r.')
            xlim([0, TVx2-TVx1]); ylim([TVy2-TVy1, 0])
            title(_title)

            ax2 = subplot(212)
            ax2.plot( swimdir_within )
            ax2.plot( peaks_within*1.15-0.1, 'mo' )
            if matched:
                xmin, xmax = t-preRange-10*fps, matched[0]+preRange/4
            else:
                xmin, xmax = t-preRange-10*fps, t+preRange/2+10*fps
            gzcs = np.cumsum(swimdir_within)
            gzcs -= gzcs[xmin]
            ax2.plot( gzcs/gzcs[xmax] )
            drawLines(0,1.2, events)
            ylim([0,1.2])
            xlim([xmin, xmax])
            ylabel('|: SwimDirection\no: approach events')

            data[fish][CS].append( {
                    'fname' : fname,
                    'x': x[rng], 'y': y[rng], 'z': z[rng],
                    'X': X[rng], 'Y': Y[rng], 'Z': Z[rng], # calibrate space (mm)
                    'speed3D': _speed3D,                   # calibrate space (mm)
                    'movingSTD' : _movingSTD,              # calibrate space (mm)
                    'd2inflow': _d2inflow,                 # calibrate space (mm)
                    'ringpixels': _ringpixels,

                    'peaks_within': peaks_within[rng],
                    'xy_within': xy_within[rng],
                    'location_one_third' : location_one_third[rng],
                    'swimdir_within' : swimdir_within[rng],

                    'dtheta_shape': dtheta_shape[rng],
                    'dtheta_vel': dtheta_vel[rng],
                    'turns_shape': turns_shape_array[rng], # already +/- preRange
                    'turns_vel': turns_vel_array[rng],

                    'events' : events,
                    'matchedUSname' : matchedUSname,
                    'TVroi' : (TVx1,TVy1,TVx2,TVy2),
                    'SVroi' : (SVx1,SVy1,SVx2,SVy2),
                    } )
            if pp:
                fig.savefig(pp, format='pdf')
            close('all') # release memory ASAP!

    if pp:
        pp.close()

def getPDFs(pickle_files, fishnames=None, createPDF=True):

    # type checking args
    if type(pickle_files) is str:
        pickle_files = [pickle_files]

    # convert to a list or set of fish names
    if type(fishnames) is str:
        fishnames = [fishnames]
    elif not fishnames:
        fishnames = set()

    # re-organize trials into a dict "data"
    data = {}
    # figure out trial number (sometime many trials in one files) for each fish
    # go through all pickle_files and use timestamps of file to sort events.
    timestamps = []

    for fp in pickle_files:
        # collect ctime of pickled files
        fname = os.path.basename(fp).split('.')[0] + '.avi'
        timestamps.append( time.strptime(fname, "%b-%d-%Y_%H_%M_%S.avi") )
        # look into the pickle and collect fish analyzed
        params = np.load(fp) # loading pickled file!
        if type(fishnames) is set:
            for fish in [fs for fl,fs in params.keys() if fl == fname and fs != 'mog']:
                fishnames.add(fish)
    timestamps = sorted(range(len(timestamps)), key=timestamps.__getitem__)

    # For each fish, go thru all pickled files
    for fish in fishnames:
        data[fish] = {}

        # now go thru the sorted
        for ind in timestamps:

            fp = pickle_files[ind]
            print 'processing #%d\n%s' % (ind, fp)
            add2DataAndPlot(fp, fish, data, createPDF)

    return data

def plotTrials(data, fish, CSname, key, step, offset=0, pp=None):
    fig = figure(figsize=(12,8), facecolor='w')

    ax1 = fig.add_subplot(121) # raw trace
    ax2 = fig.add_subplot(222) # learning curve
    ax3 = fig.add_subplot(224) # bar plot
    preP, postP, postP2 = [], [], []
    longestUS = 0
    for n, measurement in enumerate(data[fish][CSname]):
        tr = n+1
        CS, USs, preRange = measurement['events']
        subplot(ax1)
        mi = -step*(tr-1)
        ma = mi + step
        drawLines(mi, ma, (preRange, [preRange+(USs[0]-CS)], preRange))
        longestUS = max([us-CS+preRange*3/2 for us in USs]+[longestUS])

        # 'measurement[key]': vector around the CS timing (+/-) preRange. i.e., preRange is the center
        ax1.plot(measurement[key]-step*(tr-1)+offset)
        title(CSname+': '+key)                                                                  # cf. preRange = 3600 frames
        pre = measurement[key][:preRange].mean()+offset                                       # 2 min window
        post = measurement[key][preRange:preRange+(USs[0]-CS)].mean()+offset                  # 23 s window
        post2 = measurement[key][preRange+(USs[0]-CS):preRange*3/2+(USs[0]-CS)].mean()+offset # 1 min window after US
        preP.append(pre)
        postP.append(post)
        postP2.append(post2)

        ax3.plot([1, 2, 3], [pre, post, post2],'o-')

    ax1.set_xlim([0,longestUS])
    ax1.axis('off')

    subplot(ax2)
    x = range(1, tr+1)
    y = np.diff((preP,postP), axis=0).ravel()
    ax2.plot( x, y, 'ko-', linewidth=2 )
    ax2.plot( x, np.zeros_like(x), '-.', linewidth=1, color='gray' )
    # grid()
    slope, intercept, rvalue, pval, stderr = stats.stats.linregress(x,y)
    title('slope = zero? p-value = %f' % pval)
    ax2.set_xlabel("Trial#")
    ax2.set_xlim([0.5,tr+0.5])
    ax2.set_ylabel('CS - pre')

    subplot(ax3)
    ax3.bar([0.6, 1.6, 2.6], [np.nanmean(preP), np.nanmean(postP), np.nanmean(postP2)], facecolor='none')
    t, pval = stats.ttest_rel(postP, preP)
    title('paired t p-value = %f' % pval)
    ax3.set_xticks([1,2,3])
    ax3.set_xticklabels(['pre', CSname, measurement['matchedUSname']])
    ax3.set_xlim([0.5,3.5])
    ax3.set_ylabel('Raw mean values')

    tight_layout(2, h_pad=1, w_pad=1)

    if pp:
        fig.savefig(pp, format='pdf')
    close('all')

    return np.vstack((preP, postP, postP2))

def getSummary(data, dirname=None):

    for fish in data.keys():
        for CSname in data[fish].keys():

            if dirname:
                pp = PdfPages(os.path.join(dirname, '%s_for_%s.pdf' % (CSname,fish)))
                print 'generating %s_for_%s.pdf' % (CSname,fish)

            book = Workbook()
            sheet1 = book.add_sheet('speed3D')
            avgs = plotTrials(data, fish, CSname, 'speed3D', 30, pp=pp)
            putNp2xls(avgs, sheet1)

            sheet2 = book.add_sheet('d2inflow')
            avgs = plotTrials(data, fish, CSname, 'd2inflow', 200, pp=pp)
            putNp2xls(avgs, sheet2)

            # sheet3 = book.add_sheet('smoothedz')
            sheet3 = book.add_sheet('Z')
            # avgs = plotTrials(data, fish, CSname, 'smoothedz', 100, pp=pp)
            avgs = plotTrials(data, fish, CSname, 'Z', 30, pp=pp)
            putNp2xls(avgs, sheet3)

            sheet4 = book.add_sheet('ringpixels')
            avgs = plotTrials(data, fish, CSname, 'ringpixels', 1200, pp=pp)
            putNp2xls(avgs, sheet4)

            sheet5 = book.add_sheet('peaks_within')
            avgs = plotTrials(data, fish, CSname, 'peaks_within', 1.5, pp=pp)
            putNp2xls(avgs, sheet5)

            sheet6 = book.add_sheet('swimdir_within')
            avgs = plotTrials(data, fish, CSname, 'swimdir_within', 1.5, pp=pp)
            putNp2xls(avgs, sheet6)

            sheet7 = book.add_sheet('xy_within')
            avgs = plotTrials(data, fish, CSname, 'xy_within', 1.5, pp=pp)
            putNp2xls(avgs, sheet7)

            sheet8 = book.add_sheet('turns_shape')
            avgs = plotTrials(data, fish, CSname, 'turns_shape', 1.5, pp=pp)
            putNp2xls(avgs, sheet8)

            sheet9 = book.add_sheet('turns_vel')
            avgs = plotTrials(data, fish, CSname, 'turns_vel', 1.5, pp=pp)
            putNp2xls(avgs, sheet9)

            if dirname:
                pp.close()
                book.save(os.path.join(dirname, '%s_for_%s.xls' % (CSname,fish)))
                close('all')
            else:
                show()

def add2Pickles(dirname, pickle_files):
    # dirname : folder to look for pickle files
    # pickle_files : output, a list to be concatenated.
    pattern = os.path.join(dirname, '*.pickle')
    temp = [_ for _ in glob(pattern) if not _.endswith('- Copy.pickle') and
                    not os.path.basename(_).startswith('Summary')]
    pickle_files += temp

if __name__ == '__main__':

    pickle_files = []

    # small test data
    # add2Pickles('R:/Data/itoiori/behav/adult whitlock/conditioning/NeuroD/Aug4/test', pickle_files)
    # outputdir = 'R:/Data/itoiori/behav/adult whitlock/conditioning/NeuroD/Aug4/test'

    # show me what you got
    for pf in pickle_files:
        print pf

    fp = os.path.join(outputdir, 'Summary.pickle')
    createPDF = True # useful when plotting etc code updated
    if 1: # refresh analysis
        data = getPDFs(pickle_files, createPDF=createPDF)
        import cPickle as pickle
        with open(os.path.join(outputdir, 'Summary.pickle'), 'wb') as f:
            pickle.dump(data, f)
    else: # or reuse previous
        data = np.load(fp)

    getSummary(data, outputdir)
    pickle2mat(fp, data)


# Copyright 2016 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Tests for histogram_ops."""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import numpy as np
import tensorflow as tf

from tensorflow.contrib.metrics.python.ops import histogram_ops

class Strict1dCumsumTest(tf.test.TestCase):
  """Test this private function."""

  def test_empty_tensor_returns_empty(self):
    with self.test_session():
      tensor = tf.constant([])
      result = histogram_ops._strict_1d_cumsum(tensor, 0)
      expected = tf.constant([])
      np.testing.assert_array_equal(expected.eval(), result.eval())

  def test_length_1_tensor_works(self):
    with self.test_session():
      tensor = tf.constant([3], dtype=tf.float32)
      result = histogram_ops._strict_1d_cumsum(tensor, 1)
      expected = tf.constant([3], dtype=tf.float32)
      np.testing.assert_array_equal(expected.eval(), result.eval())

  def test_length_3_tensor_works(self):
    with self.test_session():
      tensor = tf.constant([1, 2, 3], dtype=tf.float32)
      result = histogram_ops._strict_1d_cumsum(tensor, 3)
      expected = tf.constant([1, 3, 6], dtype=tf.float32)
      np.testing.assert_array_equal(expected.eval(), result.eval())

class AUCUsingHistogramTest(tf.test.TestCase):

  def setUp(self):
    self.rng = np.random.RandomState(0)

  def test_empty_labels_and_scores_gives_nan_auc(self):
    with self.test_session():
      labels = tf.constant([], shape=[0], dtype=tf.bool)
      scores = tf.constant([], shape=[0], dtype=tf.float32)
      score_range = [0, 1.]
      auc, update_op = tf.contrib.metrics.auc_using_histogram(labels, scores,
                                                              score_range)
      tf.local_variables_initializer().run()
      update_op.run()
      self.assertTrue(np.isnan(auc.eval()))

  def test_perfect_scores_gives_auc_1(self):
    self._check_auc(nbins=100,
                    desired_auc=1.0,
                    score_range=[0, 1.],
                    num_records=50,
                    frac_true=0.5,
                    atol=0.05,
                    num_updates=1)

  def test_terrible_scores_gives_auc_0(self):
    self._check_auc(nbins=100,
                    desired_auc=0.0,
                    score_range=[0, 1.],
                    num_records=50,
                    frac_true=0.5,
                    atol=0.05,
                    num_updates=1)

  def test_many_common_conditions(self):
    for nbins in [50]:
      for desired_auc in [0.3, 0.5, 0.8]:
        for score_range in [[-1, 1], [-10, 0]]:
          for frac_true in [0.3, 0.8]:
            # Tests pass with atol = 0.03.  Moved up to 0.05 to avoid flakes.
            self._check_auc(nbins=nbins,
                            desired_auc=desired_auc,
                            score_range=score_range,
                            num_records=100,
                            frac_true=frac_true,
                            atol=0.05,
                            num_updates=50)

  def test_large_class_imbalance_still_ok(self):
    # With probability frac_true ** num_records, each batch contains only True
    # records.  In this case, ~ 95%.
    # Tests pass with atol = 0.02.  Increased to 0.05 to avoid flakes.
    self._check_auc(nbins=100,
                    desired_auc=0.8,
                    score_range=[-1, 1.],
                    num_records=10,
                    frac_true=0.995,
                    atol=0.05,
                    num_updates=1000)

  def test_super_accuracy_with_many_bins_and_records(self):
    # Test passes with atol = 0.0005.  Increased atol to avoid flakes.
    self._check_auc(nbins=1000,
                    desired_auc=0.75,
                    score_range=[0, 1.],
                    num_records=1000,
                    frac_true=0.5,
                    atol=0.005,
                    num_updates=100)

  def _check_auc(self,
                 nbins=100,
                 desired_auc=0.75,
                 score_range=None,
                 num_records=50,
                 frac_true=0.5,
                 atol=0.05,
                 num_updates=10):
    """Check auc accuracy against synthetic data.

    Args:
      nbins:  nbins arg from contrib.metrics.auc_using_histogram.
      desired_auc:  Number in [0, 1].  The desired auc for synthetic data.
      score_range:  2-tuple, (low, high), giving the range of the resultant
        scores.  Defaults to [0, 1.].
      num_records:  Positive integer.  The number of records to return.
      frac_true:  Number in (0, 1).  Expected fraction of resultant labels that
        will be True.  This is just in expectation...more or less may actually
        be True.
      atol:  Absolute tolerance for final AUC estimate.
      num_updates:  Update internal histograms this many times, each with a new
        batch of synthetic data, before computing final AUC.

    Raises:
      AssertionError: If resultant AUC is not within atol of theoretical AUC
        from synthetic data.
    """
    score_range = [0, 1.] or score_range
    with self.test_session():
      labels = tf.placeholder(tf.bool, shape=[num_records])
      scores = tf.placeholder(tf.float32, shape=[num_records])
      auc, update_op = tf.contrib.metrics.auc_using_histogram(labels,
                                                              scores,
                                                              score_range,
                                                              nbins=nbins)
      tf.local_variables_initializer().run()
      # Updates, then extract auc.
      for _ in range(num_updates):
        labels_a, scores_a = synthetic_data(desired_auc, score_range,
                                            num_records, self.rng, frac_true)
        update_op.run(feed_dict={labels: labels_a, scores: scores_a})
      labels_a, scores_a = synthetic_data(desired_auc, score_range, num_records,
                                          self.rng, frac_true)
      # Fetch current auc, and verify that fetching again doesn't change it.
      auc_eval = auc.eval()
      self.assertAlmostEqual(auc_eval, auc.eval(), places=5)

    msg = ('nbins: %s, desired_auc: %s, score_range: %s, '
           'num_records: %s, frac_true: %s, num_updates: %s') % (nbins,
                                                                 desired_auc,
                                                                 score_range,
                                                                 num_records,
                                                                 frac_true,
                                                                 num_updates)
    np.testing.assert_allclose(desired_auc, auc_eval, atol=atol, err_msg=msg)

def synthetic_data(desired_auc, score_range, num_records, rng, frac_true):
  """Create synthetic boolean_labels and scores with adjustable auc.

  Args:
    desired_auc:  Number in [0, 1], the theoretical AUC of resultant data.
    score_range:  2-tuple, (low, high), giving the range of the resultant scores
    num_records:  Positive integer.  The number of records to return.
    rng:  Initialized np.random.RandomState random number generator
    frac_true:  Number in (0, 1).  Expected fraction of resultant labels that
      will be True.  This is just in expectation...more or less may actually be
      True.

  Returns:
    boolean_labels:  np.array, dtype=bool.
    scores:  np.array, dtype=np.float32
  """
  # We prove here why the method (below) for computing AUC works.  Of course we
  # also checked this against sklearn.metrics.roc_auc_curve.
  #
  # First do this for score_range = [0, 1], then rescale.
  # WLOG assume AUC >= 0.5, otherwise we will solve for AUC >= 0.5 then swap
  # the labels.
  # So for AUC in [0, 1] we create False and True labels
  # and corresponding scores drawn from:
  # F ~ U[0, 1],  T ~ U[x, 1]
  # We have,
  # AUC
  #  = P[T > F]
  #  = P[T > F | F < x] P[F < x] + P[T > F | F > x] P[F > x]
  #  = (1 * x) + (0.5 * (1 - x)).
  # Inverting, we have:
  # x = 2 * AUC - 1, when AUC >= 0.5.
  assert 0 <= desired_auc <= 1
  assert 0 < frac_true < 1

  if desired_auc < 0.5:
    flip_labels = True
    desired_auc = 1 - desired_auc
    frac_true = 1 - frac_true
  else:
    flip_labels = False
  x = 2 * desired_auc - 1

  labels = rng.binomial(1, frac_true, size=num_records).astype(bool)
  num_true = labels.sum()
  num_false = num_records - labels.sum()

  # Draw F ~ U[0, 1], and T ~ U[x, 1]
  false_scores = rng.rand(num_false)
  true_scores = x + rng.rand(num_true) * (1 - x)

  # Reshape [0, 1] to score_range.
  def reshape(scores):
    return score_range[0] + scores * (score_range[1] - score_range[0])

  false_scores = reshape(false_scores)
  true_scores = reshape(true_scores)

  # Place into one array corresponding with the labels.
  scores = np.nan * np.ones(num_records, dtype=np.float32)
  scores[labels] = true_scores
  scores[~labels] = false_scores

  if flip_labels:
    labels = ~labels

  return labels, scores

if __name__ == '__main__':
  tf.test.main()