import pandas as pd
import matplotlib.pyplot as plt
from scipy import signal
import numpy as np
import math
import easygui

scorer=['DLC_resnet50_Cylinder_camera-1Jun15shuffle1_250000', 'DLC_resnet50_Cylinder_camera-2Jun15shuffle1_250000', 'DLC_resnet50_GridWalk_camera-1Jun15shuffle1_250000', 'DLC_resnet50_GridWalk_camera-2Jun15shuffle1_250000']
path=easygui.fileopenbox(multiple=True, default='/home/user/owncloud/3D_videos/3D/GW', filetypes="*.csv")


def likelihood(data, threshold):
    coln=data.columns
    for i in range(1, len(coln), 3):
        data.loc[data[coln[i][0], 'likelihood']<threshold, (coln[i][0], 'x')]=math.nan
        data.loc[data[coln[i][0], 'likelihood']<threshold, (coln[i][0], 'y')]=math.nan
    return data

def prep_dlc(pathway, threshold, height):
    data=pd.read_csv(pathway, index_col=0, delimiter=',', skiprows=0, header=[1,2])
    data=likelihood(data, threshold)
    #data=y_invert(data, height)
    return data


for i in path:
    header = list(range(3))
    data=prep_dlc(i, 0.9, 1080)
    names=[i[0] for i in data.columns[::3]]
    for j in names:
        if np.count_nonzero(data['snout', 'x'].value_counts())<2:
            data[j, 'x']=data[j, 'x'].interpolate(method='polynomial', order=1)
            data[j, 'y']=data[j, 'y'].interpolate(method='polynomial', order=1)
            data[j, 'likelihood']=0.95
        data.loc[pd.isna(data[j, 'x']), (j, 'likelihood')]=0.1
        data.loc[pd.isna(data[j, 'y']), (j, 'likelihood')]=0.1
    if (('CY') and ('camera-1')) in i:    
        data=pd.concat([data], keys=[scorer[0]], axis=1, names=['scorer'])
    elif (('CY') and ('camera-2')) in i:    
        data=pd.concat([data], keys=[scorer[1]], axis=1, names=['scorer'])
    elif (('GW') and ('camera-1')) in i:    
        data=pd.concat([data], keys=[scorer[2]], axis=1, names=['scorer'])
    elif (('GW') and ('camera-2')) in i:    
        data=pd.concat([data], keys=[scorer[3]], axis=1, names=['scorer'])  
    if 'filtered' in i:
        data.to_hdf(i.replace('_filtered.csv', '.h5'), key='df_with_missing', mode='w')
        data.to_hdf(i.replace('_filtered.csv', '_filtered.h5'), key='df_with_missing', mode='w')
        data.to_hdf(i, key='df_with_missing', mode='w')
        data.to_hdf(i.replace('_filtered.csv', '.csv'), key='df_with_missing', mode='w')
    else:
        data.to_hdf(i.replace('.csv', '.h5'), key='df_with_missing', mode='w')