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

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)
    pathway='/home/user/owncloud/3D_videos/CY_2D/'+z+'_P3_CY_camera-1DLC_resnet50_Cylinder_camera-1Jun15shuffle1_250000_filtered.csv'
    data=prep_dlc(pathway, 0.9, 1080)
    names=[i[0] for i in data.columns[::3]]
    for j in names:
        data[j, 'x']=data[j, 'x'].interpolate(method='polynomial', order=1)
        data[j, 'y']=data[j, 'y'].interpolate(method='polynomial', order=1)
    dill.load_session('/home/user/Documents/Master/CY_2D_'+z+'_data_save.pkl')
    touch_reset=0
    drag_reset=0
    touch=[]
    drag=[]
    length=[len(data), len(df_vel_acc1)]
    for i in range(0, np.min(length)):
        if data['snout', 'y'][i]>data['cylinder_top', 'y'][i]:
            #touch right
            if (touch_reset==0) and (df_vel_acc1['angle_shoulder_right_elbow_right_wrist_right'][i]>150):
                touch.append(i)
                touch_reset=1
                print('touch at ', i)
    
            #drag right
            if (drag_reset==0) and (df_vel_acc1['angle_acceleration_tailbase_shoulder_right_wrist_right'][i]>30):
                drag.append(i)
                drag_reset=1
                print('drag at ', i)
    
        if (touch_reset==1) and (data['snout', 'y'][i]<data['cylinder_top', 'y'][i]):
            touch_reset=0
            drag_reset=0
    if z=='PB_T2_3_1':
        touch_2d=pd.DataFrame([touch], index=[z])
        drag_2d=pd.DataFrame([drag], index=[z])
    else:
        touch_2d_temp=pd.DataFrame([touch], index=[z])
        touch_2d=touch_2d.append(touch_2d_temp)
        drag_2d_temp=pd.DataFrame([drag], index=[z])
        drag_2d=drag_2d.append(drag_2d_temp)

touch_2d.columns=['Frame #']*touch_2d.shape[1]
dfi.export(touch_2d, '/home/user/owncloud/thesis_figures/DLC_CY_touch_2d.png', max_cols=touch_2d.shape[1])
drag_2d.columns=['Frame #']*drag_2d.shape[1]
dfi.export(drag_2d, '/home/user/owncloud/thesis_figures/DLC_CY_drag_2d.png', max_cols=drag_2d.shape[1])

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_2D_'+z+'_data.pkl') 
    pred_rfc_touch_2d = best_model_touch.predict(X_touch_test)
    det_touch_2d=np.where(pred_rfc_touch_2d)
    pred_rfc_drag_2d = best_model_drag.predict(X_drag_test)
    det_drag_2d=np.where(pred_rfc_drag_2d)
    if z=='PB_T2_3_1':
        touch_FI_2d=pd.DataFrame(det_touch_2d, index=[z])
        drag_FI_2d=pd.DataFrame(det_drag_2d, index=[z])
    else:
        touch_FI_temp_2d=pd.DataFrame(det_touch_2d, index=[z])
        touch_FI_2d=touch_FI_2d.append(touch_FI_temp_2d)
        drag_FI_temp_2d=pd.DataFrame(det_drag_2d, index=[z])
        drag_FI_2d=drag_FI_2d.append(drag_FI_temp_2d)