#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Fri Aug  6 16:14:05 2021

@author: titan
"""


# -*- coding: utf-8 -*-
"""
Created on Wed Jul 28 10:47:14 2021
@author: Dimitrios
"""

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


def read1(name1):
#k4 = k4 != 100
    pkl_file = open(name1, 'rb')
    dot = pickle.load(pkl_file)
    dot = pd.DataFrame(dot)
    return dot

def read2(name2):
    stimuli = np.load(name2, allow_pickle=True)
    stimuli = pd.DataFrame(stimuli)
    return stimuli

def sizefilter(df,length):
    newtotal3=pd.DataFrame()
    for i in range(df['trial'].max()):
       # print('trial',i)
        b=df.loc[df['trial']==i]
        c=len(b)
        #print('length',c)
        if c<length:
            #print('adding this one')
            newtotal3= pd.concat([newtotal3,b])
    return newtotal3


def helpingdf(df):
    dist=df[['trial','distance']]
    dist2=dist.pivot(columns='trial',values='distance')
    distmin=dist2.min()
    test2=df.loc[df['distance'].isin(distmin)]
    test2['mid_x']=(test2['right_ear_x_c']+test2['left_ear_x_c'])/2
    test2['mid_y']=(test2['right_ear_y_c']+test2['left_ear_y_c'])/2
    test2['m']=(test2.y_snout-test2.mid_y)/(test2.x_snout-test2.mid_y)
    test2['b']=test2.y_snout-test2.m*test2.x_snout
    test2['x_infront']=(test2['y']-test2['b'])/test2['m']
    test2['testcol2']=test2.x_infront-test2.x
    test2['where']=['right' if i>0 else 'left' for i in  test2['testcol2'] ]
   
    test2['dist_left_ear']=np.sqrt((test2['left_ear_x_c']-test2['x'])**2+(test2['left_ear_y_c']-test2['y'])**2)
    test2['dist_right_ear']=np.sqrt((test2['right_ear_x_c']-test2['x'])**2+(test2['right_ear_y_c']-test2['y'])**2)
    test2['dist_l-r']=test2['dist_left_ear']-test2['dist_right_ear']
    test2['where_dist']=['right' if i>0 else 'left' for i in  test2['dist_l-r'] ]
    
    return test2


'''Begin by opening the stimulus data'''

k = read1('/home/titan/Desktop/SIlencing/375529_4/thedata96.pkl') #open the pickle file with stimuli data

k.columns=['x','y','edges','opacity','radius','orientation','delete','x','y','edges','opacity','radius','orientation','laser','frame_time']
k = k.drop([0])
k.frame_time = k.frame_time.astype(float) #change the frame time into float 


k1=k.iloc[:,0:6] # first stimuli

k2=k.iloc[:,7:12] # second stimuli

k3 =k.iloc[:,13] # laser

k4 = k.iloc[:,14] # frame time


k = pd.concat([k1, k3, k4], axis=1) #stimulus we are analysing

stim2 = pd.concat([k2, k3], axis=1) #send stimuli to analyse

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

''' Open the Pose file'''

df = pd.read_hdf('/home/titan/Desktop/SIlencing/375529_4/blackfly_375529_2021-07-22_4_DLC.hdf5') #open the hdf5 file from DLC live

df.columns=['snout_x','snout_y','snout_likelihood','left_ear_x','left_ear_y',
            'left_ear_likelihood','right_ear_x','right_ear_y','right_ear_likelihood',
            'tail_x','tail_y','tail_likelihood','frame_time','pose_time'] #rename the columns from the DLC live file 

df.frame_time = df.frame_time.astype(float)

k.frame_time = k.frame_time.round(1) #round the decimal points of the frame time to match the hdf5 file
 
df.frame_time = df.frame_time.round(1) #round the decimal points of the frame time to match the pickle file

'''merge the two files '''

merged = df.merge(k, how='outer', on='frame_time') # merge the hdf5 file with the pickle file dependent on the time frame


''' find the head body angle '''

[x1,y1] = [merged.snout_x,merged.snout_y]
[x2,y2] = [merged.left_ear_x,merged.left_ear_y]
[x3,y3] = [merged.right_ear_x,merged.right_ear_y]
[x4,y4] = [(x3+x2)/2,(y3+y2)/2]
angles = []
    
for i in range(merged.shape[0]):
    a = math.sqrt((x4[i]-x2[i])**2+(y4[i]-y2[i])**2)
    b = math.sqrt((x4[i]-x1[i])**2+(y4[i]-y1[i])**2)
    c = math.sqrt((x1[i]-x2[i])**2+(y1[i]-y2[i])**2)
    if (a**2+b**2-c**2)/(2*b*a) > 1:
        angle = math.acos(1)
    else:
        angle = math.acos(np.round((a**2+b**2-c**2)/(2*b*a),4))
    angle = math.degrees(angle)
    angles.append(angle)

right_head_angle = [180 - Value for Value in (angles)]  


[x1a,y1a] = [merged.tail_x,merged.tail_y]
[x2,y2] = [merged.left_ear_x,merged.left_ear_y]
[x3,y3] = [merged.right_ear_x,merged.right_ear_y]
[x4a,y4a] = [(x3+x2)/2,(y3+y2)/2]
angles1 = []
    
for i in range(merged.shape[0]):
    d = math.sqrt((x4a[i]-x2[i])**2+(y4a[i]-y2[i])**2)
    e = math.sqrt((x4a[i]-x1a[i])**2+(y4a[i]-y1a[i])**2)
    f = math.sqrt((x1a[i]-x2[i])**2+(y1a[i]-y2[i])**2)
    
    angle1 = math.acos(np.round((d**2+e**2-f**2)/(2*e*d),4))
    angle1 = math.degrees(angle1)
    angles1.append(angle1)
    
result = [Value + Value for Value, Value in zip(angles, angles1)]


data= df = pd.DataFrame(result) 

total = pd.concat([data, merged], axis=1).reindex(merged.index) #merge angles into the total dataframe


''' Change the name of the columns to fit the new dataframe'''

total.columns=['left_head_body_angle', 'snout_x','snout_y','snout_likelihood','left_ear_x','left_ear_y',
            'left_ear_likelihood','right_ear_x','right_ear_y','right_ear_likelihood',
            'tail_x','tail_y','tail_likelihood','frame_time','pose_time','x','y','edges'
            ,'opacity','radius','orientation','laser']


'''translate the frame time into real time'''

total.frame_time.astype(float)
first_frame = total.iloc[1,13]

translated_time=total.frame_time-first_frame

total = pd.concat([total, translated_time], axis=1).reindex(total.index)


total.left_head_body_angle.astype(float)


total.columns=['left_head_body_angle', 'snout_x','snout_y','snout_likelihood','left_ear_x','left_ear_y',
            'left_ear_likelihood','right_ear_x','right_ear_y','right_ear_likelihood',
            'tail_x','tail_y','tail_likelihood','frame_time','pose_time','x','y','edges'
            ,'opacity','radius','orientation','laser','time']

#Correct the data to only show the apperance of the stimuli

total = total.fillna(method='ffill')

total = total.dropna()



'''make laser true or false'''

laser = total['laser']

laser1 = laser%2==0


'''convert coordinates into psychopy'''

snout = ((total.iloc[0:total['snout_x'].shape[0]-1,1:3] - [45, 15]) / 380) * .93 + .035
snout['snout_y'] = 1-snout['snout_y']
snout = (snout-0.5)*2

total = pd.concat([total, snout], axis=1).reindex(total.index)

leftear = ((total.iloc[0:total['left_ear_x'].shape[0]-1,4:6] - [45, 15]) / 380) * .93 + .035
#snout = ((total.iloc[764,2] - [45, 15]) / 380) * .93 + .035
leftear['left_ear_y'] = 1-leftear['left_ear_y']
leftear = (leftear-0.5)*2

total = pd.concat([total, leftear], axis=1).reindex(total.index)

rightear = ((total.iloc[0:total['right_ear_x'].shape[0]-1,7:9] - [45, 15]) / 380) * .93 + .035
#snout = ((total.iloc[764,2] - [45, 15]) / 380) * .93 + .035
rightear['right_ear_y'] = 1-rightear['right_ear_y']
rightear = (rightear-0.5)*2

total = pd.concat([total, rightear], axis=1).reindex(total.index)

tail = ((total.iloc[0:total['tail_x'].shape[0]-1,10:12] - [45, 15]) / 380) * .93 + .035
#snout = ((total.iloc[764,2] - [45, 15]) / 380) * .93 + .035
tail['tail_y'] = 1-tail['tail_y']
tail = (tail-0.5)*2


total = pd.concat([total, tail], axis=1).reindex(total.index)


total.columns=['left_head_body_angle', 'snout_x_w','snout_y_w','snout_likelihood_w','left_ear_x_w','left_ear_y_w',
            'left_ear_likelihood','right_ear_x','right_ear_y','right_ear_likelihood',
            'tail_x','tail_y','tail_likelihood','frame_time','pose_time','x','y','edges'
            ,'opacity','radius','orientation','laser','time','x_snout','y_snout','left_ear_x_c','left_ear_y_c','right_ear_x_c','right_ear_y_c','tail_x_c','tail_y_c']

total = pd.concat([total, laser1], axis=1).reindex(total.index)


total = total.reset_index()
"find the angle of the dot from x axis and invert head angle in pose so the subtraction with dot angle is correct"""


#Make snout psychopy coordinates fit with the stimulus psychopy coordinates
total['x_snout'] = total['x_snout'] * 0.4
total['y_snout'] = total['y_snout'] * 0.5

total['left_ear_x_c'] = total['left_ear_x_c'] * 0.4
total['left_ear_y_c'] = total['left_ear_y_c'] * 0.5

total['right_ear_x_c'] = total['right_ear_x_c'] * 0.4
total['right_ear_y_c'] = total['right_ear_y_c'] * 0.5

total['tail_x_c'] = total['tail_x_c'] * 0.4
total['tail_y_c'] = total['tail_y_c'] * 0.5


dangle = []
for i in total.index:
    dangle.append(math.atan2(total['y'][i], total['x'][i]))
total['dangle'] = dangle
total['dangle'] = np.rad2deg(total['dangle'])

total['dx'] = total['x_snout'] - total['x']
total['dy'] = total['y_snout'] - total['y']

total['dxsq'] = total['dx']**2
total['dysq'] = total['dy']**2
total['dis'] = total['dxsq'] + total['dysq']
total['dis'] = np.sqrt(total['dis'])


#calculate the line between the snout and midpoint between ears

[x1,y1] = [total.x_snout,total.y_snout]
[x2,y2] = [total.left_ear_x_c,total.left_ear_y_c]
[x3,y3] = [total.right_ear_x_c,total.right_ear_y_c]
[x4,y4] = [(x3+x2)/2,(y3+y2)/2]

m = (y1 - y4) /(x1 - x4)
b = y1 - m * x1


ysnout = total['y_snout']
xsnout = total['x_snout']
xstimulus = total['x']
ystimulus = total['y']
y22 = m*(xstimulus) + b

mouseangles = []

for i in range(len(total)):
    a = y22[i]-ystimulus[i]  

    b = math.sqrt((xstimulus[i]-xsnout[i])**2 + (y22[i]-ysnout[i])**2)
    c = math.sqrt((xstimulus[i]-xsnout[i])**2 + (ystimulus[i]-ysnout[i])**2)
    angle = math.acos((b**2+c**2-a**2)/(2*b*c))
    angle = math.degrees(angle)
    mouseangles.append(angle)


df = pd.DataFrame(mouseangles)
total = pd.concat([total, df], axis=1).reindex(total.index)

total = total.drop(columns=['index'])


total.columns=['left_head_body_angle', 'snout_x','snout_y','snout_likelihood','left_ear_x','left_ear_y',
            'left_ear_likelihood','right_ear_x','right_ear_y','right_ear_likelihood',
            'tail_x','tail_y','tail_likelihood','frame_time','pose_time','x','y','edges'
            ,'opacity','radius','orientation','laser','time','x_snout','y_snout','left_ear_x_c','left_ear_y_c',
            'right_ear_x_c','right_ear_y_c','tail_x_c','tail_y_c','laser1','dangle','dx','dy','dxsq','dysq','distance','mouse_angle']


new_total = total

new_total['looking']=['right' if i>180 else 'left' for i in  new_total['left_head_body_angle'] ]


# drop the nan values

new_total = new_total.dropna()
new_total = new_total[new_total['opacity'] != 100]
new_total = new_total[new_total['y_snout'] < 0.4]
new_total = new_total[new_total['y_snout'] > 0].reset_index()


#find the start, end and the number of trials

new_total.insert(40, 'trial', 0)
x = 0
idx = 0

for i in range(len(new_total)):
    if x == new_total['x'][i]:
        new_total.iloc[i,40] = idx

    else:
        x = new_total['x'][i]
        idx = idx + 1
        new_total.iloc[i,40] = idx

'''
dimitris = helpingdf(new_total)

Side = helpingdf(new_total)
Side1 = Side['trial']
    Side2 = Side['where_dist']
    Side = pd.concat([Side1, Side2], axis=1)
    Side = Side.reset_index()
    Side = Side.drop(columns=['index'])
trials = new_total['trial']
t = trials.value_counts()
t = t.reset_index()
t.columns=['trial','frame']
'''
        
headangle = new_total['mouse_angle']
trials = new_total['trial']               

headangletrials = pd.concat([trials, headangle], axis=1)
headangletrials = headangletrials.pivot(columns='trial', values='mouse_angle')
difference_angle = headangletrials.diff();
difference_angle = difference_angle.min()

difference_angle = difference_angle.reset_index()
difference_angle.columns=['trial','angle']

normal = difference_angle.merge(new_total,how='outer',on='trial')
nostimuli = normal[normal['opacity'] == 0]
stimuli = normal[normal['opacity'] != 0]

nostimuli1 = nostimuli['trial']
nostimuli1 = nostimuli1.drop_duplicates()
nostimuli1 = nostimuli1.count()
stimuli1 = stimuli['trial']
stimuli1 = stimuli1.drop_duplicates()
stimuli1 = stimuli1.count()

nostimuli2 = nostimuli[nostimuli['angle'] < -100]
nostimuli2 = nostimuli2['trial']
nostimuli2 = nostimuli2.drop_duplicates()
nostimuli2 = nostimuli2.count()
stimuli2 = stimuli[stimuli['angle'] < -100]
stimuli2 = stimuli2['trial']
stimuli2 = stimuli2.drop_duplicates()
stimuli2 = stimuli2.count()

percent_nostim = nostimuli2/nostimuli1
percent_stim = stimuli2/stimuli1




test=new_total.pivot(columns='trial',values='distance')
min_dist_per_trial=test.min()
test2=new_total.loc[new_total.distance.isin(min_dist_per_trial)]
test3=test2.drop_duplicates(subset=['frame_time'])
frametimes=test3['frame_time']

indices=test3['index']
two_frames_before=indices-2
new_total_two_before=new_total.loc[new_total['index'].isin(two_frames_before)]
test3['mouse_angle']-new_total_two_before['mouse_angle']
