import os
from os import path
from glob import glob
import numpy as np
import pandas as pd
import h5py
import datashader as ds
import datashader.transfer_functions as tf
# from plots import plot_occupancy_projections

HDF_FORMAT = 'fixed'


def strip_sessions(pos):
	"""Remove the beginning and end of each session's tracking data, using its position XYZ info."""
	# Get Only the data after entry into the arena for each session.
	def get_first_index(pos):
	    indices = np.where(
	        (pos.Y < .2)   # Rat is low.
	        & (pos.Z.abs() < .2)  # Rat is inside the Arena.
	    )
	    return indices[0][0] if len(indices[0]) > 0 else 0

	def get_last_index(pos):
	    indices = np.where(
	        (.20 < np.abs(pos.X))  # Rat is off board
	        & (np.abs(pos.X) < .45)  # Rat is still close to board, though.
	        & (pos.Y > .06)  # Rat is above the floor, not sitting on it.
	        & (pos.Y < .20)  # Rat is below the board.
	        & (np.sign(pos.X) == np.sign(pos.X.diff()))  # Rat is moving away from the board.
	        & (np.abs(pos.X.diff() / pos.Time.diff()) > .08)  # Rat is moving away from the board at speed.
	        & ((pos.Y.diff() / pos.Time.diff()) < -.20)
	    )
	    
	    return indices[0][0] if len(indices[0]) > 0 else 0

	get_window = lambda df: df.iloc[(get_first_index(df) + 300):get_last_index(df) + 28 ]
	pos = pos.groupby('Session_id').apply(get_window)
	pos.reset_index('Session_id', inplace=True, drop=True)
	return pos


def filter_sessions(pos):
	"""Remove Bad Sessions"""
	pos = pos.groupby('Session_id').filter(lambda df: (np.abs(df.X) > .38).any() == False )
	pos = pos.groupby('Session_id').filter(lambda df: (np.abs(df.Y) > .49).any() == False )
	return pos.copy()


def add_relative_time(pos):
	# Get Relative Time in Session
	pos = pos.copy()
	normalize = lambda t: (t - t.min()) / (t.max() - t.min())
	pos['TimeNormed'] = pos.groupby('Session_id').Time.apply(normalize)	
	return pos


def smooth_trajectories(df, **rolling_kwargs):
	"""Transform Data (Velocity and Acceleration)"""
	pos_s = df.copy()
	g = pos_s.groupby('Session_id')
	for dim in 'XYZ':
	    pos_s[dim] = g[dim].rolling(**rolling_kwargs).mean().values
	    pos_s['Vel' + dim] = g.apply(lambda df: df[dim].diff() / df.Time.diff()).values
	    pos_s['Acc' + dim] = g.apply(lambda df: df['Vel' + dim].diff() / df.Time.diff()).values
	return pos_s



def get_jump_trajectory(df, search_size=100, win_size=40):
	def get_end_of_fall(df):
	    end = np.where(df.VelY == np.min(df.VelY.iloc[-search_size:]))[0][0]
	    return df.iloc[(end - win_size):end]
	falls = df.copy().groupby('Session_id').apply(get_end_of_fall)
	falls.reset_index('Session_id', drop=True, inplace=True)
	return falls


def get_jump_point(df, session_filter='VelY < -.3'):
	"""Get last frame of falling: The minimum VelY"""
	jumps = df.groupby('Session_id').apply(lambda df: df[df.VelY == df.VelY.min()])
	jumps.reset_index('Session_id', inplace=True, drop=True)
	jumps = jumps.query(session_filter)
	return jumps.copy()


if __name__ == '__main__':

	# Load Data
	fname = '../data/VRCliff_AllData.h5'
	if not path.exists(fname):
	    os.system("python merge_sessions.py")
	pos = pd.read_hdf(fname, '/Position').dropna()
	pos.reset_index(inplace=True, drop=True)  # Because original file put frame as index.

	sessions = pd.read_hdf(fname, '/Session')
	sessions.set_index('id', inplace=True)

	# Correct for different Board Height
	sessions['BoardCorrection'] = sessions.BoardHeight.apply(lambda x: -.01 if x > .15 else 0)
	boardcorr = pd.merge(pos[['Session_id']], sessions[['BoardCorrection']], left_on='Session_id', right_index=True,)
	pos.Y.update(pos.Y + boardcorr.BoardCorrection)


    # Process, Clean, Filter Data
	pos = strip_sessions(pos)
	pos = filter_sessions(pos)
	pos = add_relative_time(pos)
	pos_s = smooth_trajectories(pos, window=15, center=True, min_periods=9, win_type='hamming')
	falls = get_jump_trajectory(pos_s, search_size=100, win_size=40)
	jumps = get_jump_point(falls, session_filter='VelY < -.3')


    
    # # Save Cleaned Data to a new file.
    # savename = '../data/VRCliff_AllData_cleaned.h5'
    # pos.to_hdf(savename, '/Position', format='table')
    # sessions.to_hdf(savename, '/Session', format='table')

# ###########




# Correct_for Left Cliff Condition (onlyh in RealCliff=Left condition)
# dfsub = sessions.query('CliffType == "Real" & CliffSide == "L"')
# swap_lr = lambda x: {'L': 'R', 'R': 'L', np.nan: np.nan}[x]
# sessions.loc[dfsub.index, 'JumpSide'] = dfsub.JumpSide.apply(swap_lr)
# todo: Move Left Cliff correction to earlier in the process.

# Evaluate Whether jump was on correct side or not
# sessions['JumpRight'] = ((sessions.JumpSide == 'R') - 0.5) * 2
# sessions['Correct'] = sessions.JumpSide != sessions.CliffSide

# #########





# falls.to_hdf(fname, '/Falls', format=HDF_FORMAT)



# jumps.to_hdf(fname, '/Jumps', format=HDF_FORMAT)