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)

