import pandas as pd 
import numpy as np 
import matplotlib.pyplot as plt

hdf_path = 'Data.h5'
fixation_data = pd.read_hdf(hdf_path, 'fixation_data')
fixslide_fixation_data = pd.read_hdf(hdf_path, 'fixslide_fixation_data')
saccade_data = pd.read_hdf(hdf_path, 'saccade_data')
fixslide_saccade_data = pd.read_hdf(hdf_path, 'fixslide_saccade_data')

monkeys = ['Ned','Apu','Marge','Lenny','Walter']
monkey_color_map=dict(zip(monkeys, 'gbcmk'))

fixations = {
    'front': fixation_data.reset_index().query('front and not scrambled'),
    'side': fixation_data.reset_index().query('not front and not scrambled'),
    'scrambled': fixation_data.reset_index().query('scrambled'),
    'fix_slide': fixslide_fixation_data
}
saccades = {
    'front': saccade_data.reset_index().query('front and not scrambled'),
    'side': saccade_data.reset_index().query('not front and not scrambled'),
    'scrambled': saccade_data.reset_index().query('scrambled'),
    'fix_slide': fixslide_saccade_data
}

fig, axes = plt.subplots(4, 1, figsize=(3,7), sharex=True, sharey=True)
conditions = ['fix_slide', 'scrambled', 'front', 'side']
for condition, ax in zip(conditions, axes):
    for monkey, data in fixations[condition].groupby('monkey').duration:
        weights = np.ones_like(data) / data.size
        counts, bin_edges = np.histogram(data, bins=15, range=(0,600), weights=weights)
        x = np.mean([bin_edges[:-1], bin_edges[1:]], axis=0)
        ax.plot(x, counts, c=monkey_color_map[monkey], marker='o', markersize=2.5)
    ax.axvline(fixations[condition].groupby('monkey').duration.median().mean(), color='red')
    ax.set_title(condition)
    ax.set_ylabel('Proportion')
    ax.spines['right'].set_visible(False)
    ax.spines['top'].set_visible(False)
fig.savefig(f'FixationDuration.png')

fig, axes = plt.subplots(4, 1, figsize=(3,7), sharex=True, sharey=True)
conditions = ['fix_slide', 'scrambled', 'front', 'side']
for condition, ax in zip(conditions, axes):
    for monkey, data in saccades[condition].groupby('monkey').amplitude:
        weights = np.ones_like(data) / data.size
        counts, bin_edges = np.histogram(data, bins=8, range=(0,16), weights=weights)
        x = np.mean([bin_edges[:-1], bin_edges[1:]], axis=0)
        ax.plot(x, counts, c=monkey_color_map[monkey], marker='o', markersize=2.5)
    ax.axvline(saccades[condition].groupby('monkey').amplitude.median().mean(), color='red')
    ax.set_title(condition)
    ax.set_ylabel('Proportion')
    ax.spines['right'].set_visible(False)
    ax.spines['top'].set_visible(False)
fig.savefig(f'SaccadeAmplitude.png')