import pandas as pd
import matplotlib.pyplot as plt
from scipy import signal
import numpy as np
import math
from mpl_toolkits import mplot3d
import easygui
from master_funcs import *
import dill
import dataframe_image as dfi
%matplotlib qt


patients=['PB_T2_3_1', 'PB_T2_3_2', 'PB_T2_3_3', 'PB_T2_4_2', 'PB_T2_5_1', 'PB_T2_6_1', 'PB_T2_6_2', 'PB_T3_23_1', 'PB_T3_23_2', 'PB_T3_23_3', 'PB_T3_24_1', 'PB_T3_24_2', 'PB_T3_24_3']

for z in patients:
    Cyl_3D='/home/user/owncloud/3D_videos/CY_3D/'+z+'_P3_CY_DLC_3D.csv'
    data=pd.read_csv(Cyl_3D, delimiter=',', skiprows=0, index_col=0, header=[1,2])
    names=[i[0] for i in data.columns[::3]]
    zero='cylinder_left'
    x='cylinder_top'
    y='cylinder_bottom'
    data=transform(data, zero, x, y)
    vel, acc=vel_acc(data, '3D')
    df_vel=pd.DataFrame(vel, columns=names)
    df_acc=pd.DataFrame(acc, columns=names)
    touch_reset=0
    drag_reset=0
    print(z)
    touch=[]
    drag=[]
    for i in range(0, len(data-3)):
        if (distance_single(data, 'snout', 'hindpaw_right', i)>25) and (distance_single(data, 'snout', 'hindpaw_left', i)>25):
            #touch right
            if (touch_reset==0) and ((circle_dis(data, 'forepaw_right', i)>=0.90) or (circle_dis(data, 'forepaw_left', i)>=0.90)):
                touch.append(i)
                touch_reset=1
                print('touch at ', i)
    
            #drag right
            if (drag_reset==0) and ((df_acc['forepaw_right'][i]>=0.02) or (df_acc['forepaw_left'][i]>=0.02)):
                drag.append(i)
                drag_reset=1
                print('drag at ', i)
    
        if (touch_reset==1) and (distance_single(data, 'snout', 'hindpaw_right', i)<5):
            touch_reset=0
            drag_reset=0
    if z=='PB_T2_3_1':
        touch_3d=pd.DataFrame([touch], index=[z])
        drag_3d=pd.DataFrame([drag], index=[z])
    else:
        touch_3d_temp=pd.DataFrame([touch], index=[z])
        touch_3d=touch_3d.append(touch_3d_temp)
        drag_3d_temp=pd.DataFrame([drag], index=[z])
        drag_3d=drag_3d.append(drag_3d_temp)

touch_3d.columns=['Frame #']*touch_3d.shape[1]
dfi.export(touch_3d, '/home/user/owncloud/thesis_figures/DLC_CY_touch_3d.png')
drag_3d.columns=['Frame #']*drag_3d.shape[1]
dfi.export(drag_3d, '/home/user/owncloud/thesis_figures/DLC_CY_drag_3d.png')

patients=['PB_T2_3_1', 'PB_T2_3_2', 'PB_T2_3_3', 'PB_T2_4_2', 'PB_T2_5_1', 'PB_T2_6_1', 'PB_T2_6_2', 'PB_T3_23_1', 'PB_T3_23_2', 'PB_T3_23_3', 'PB_T3_24_1', 'PB_T3_24_2', 'PB_T3_24_3']
for z in patients:
    print(z)
    dill.load_session('/home/user/Documents/Master/CY_'+z+'_data.pkl') 
    patients=['PB_T2_3_1', 'PB_T2_3_2', 'PB_T2_3_3', 'PB_T2_4_2', 'PB_T2_5_1', 'PB_T2_6_1', 'PB_T2_6_2', 'PB_T3_23_1', 'PB_T3_23_2', 'PB_T3_23_3', 'PB_T3_24_1', 'PB_T3_24_2', 'PB_T3_24_3']
    pred_rfc_touch = best_model_touch.predict(X_touch_test)
    det_touch_3d=np.where(pred_rfc_touch)
    pred_rfc_drag = best_model_drag.predict(X_drag_test)
    det_drag_3d=np.where(pred_rfc_drag)
    if z=='PB_T2_3_1':
        touch_FI_3d=pd.DataFrame(det_touch_3d, index=[z])
        drag_FI_3d=pd.DataFrame(det_drag_3d, index=[z])
    else:
        touch_FI_3d_temp=pd.DataFrame(det_touch_3d, index=[z])
        touch_FI_3d=touch_FI_3d.append(touch_FI_3d_temp)
        drag_FI_3d_temp=pd.DataFrame(det_drag_3d, index=[z])
        drag_FI_3d=drag_FI_3d.append(drag_FI_3d_temp)