"""
Goal - to produce a csv file with frame, temp, gs, individual, date of exp, trial, subtrial, time, position, speed, acceleration
""" 

import os
import pathlib
from pprint import pprint

import numpy as np
from scipy import stats
from scipy.spatial import distance
import matplotlib.pyplot as plt
from matplotlib.pyplot import figure

import trajectorytools as tt
import trajectorytools.plot as ttplot
import trajectorytools.socialcontext as ttsocial
from trajectorytools.constants import dir_of_data
import csv
import pickle
import argparse

rows1 = []
with open('../../data/temp_collective/roi/all_data.csv', 'r') as csvfile:
    looms = csv.reader(csvfile)
    for row in looms:
        rows1.append(row)




def position(tr):
    return(tr.s)


def speed(tr):
    v = (position(tr)[2:] - position(tr)[:-2]) / 2
    b = np.linalg.norm(v, axis=-1)
    return(b*60)

def acceleration(tr):
    a = position(tr)[2:] - 2 * position(tr)[1:-1] + position(tr)[:-2]
    aa = np.linalg.norm(a, axis=-1)  
    return(aa*3600)
        

def filter(tr, roi = 5): #ind (for individual) starts from 0, roi - edge of region of interest
    position_mask0 = np.ma.masked_where((tr.distance_to_origin[1:-1] > roi)|(tr.distance_to_origin[0:-2] > roi)|(tr.distance_to_origin[2:] > roi), position(tr)[1:-1,:,0],copy=False)
    position_mask1 = np.ma.masked_where((tr.distance_to_origin[1:-1] > roi)|(tr.distance_to_origin[0:-2] > roi)|(tr.distance_to_origin[2:] > roi), position(tr)[1:-1,:,1],copy=False)
    return(position_mask0,position_mask1)  
    
def filter_speed(tr, roi = 5): 
    speed_mask = np.ma.masked_where((tr.distance_to_origin[1:-1] > roi)|(tr.distance_to_origin[0:-2] > roi)|(tr.distance_to_origin[2:] > roi), speed(tr),copy=False)
    
    return(speed_mask)         



def filter_acc(tr, roi = 5): 
    acc_mask = np.ma.masked_where((tr.distance_to_origin[1:-1] > roi)|(tr.distance_to_origin[0:-2] > roi)|(tr.distance_to_origin[2:] > roi), acceleration(tr),copy=False)
    
    return(acc_mask)#[~acc_mask.mask].data)   

temperature = range(9,30,4)

group = [1,2,4,8,16,32]

replication = range(10) # number of replicates per treatment

#output parent directory
parent_dir = '../../output/temp_collective/roi'

ii = 0 # to keep count of temperature

#frames = 5000 #number of frames for which annd is calculated
with open('../../data/temp_collective/roi/masked_data.csv', mode='w') as stats_speed:
    writer = csv.writer(stats_speed, delimiter=',', quotechar='"', quoting=csv.QUOTE_MINIMAL)

    writer.writerow(['Temperature', 'Groupsize', 'Replicate', 'Trial','Date','Subtrial','Time_fish_in','Time_start_record','Frame','Individual','x','y','speed','acceleration'])
    for i in temperature:
    
        jj = 0 # to keep count of groups
        for j in group:
        
            for k in replication:

                if j == 1:
                    trajectories_file_path = '../../data/temp_collective/roi/'+str(i)+'/' +str(j)+'/GS_'+str(j)+'_T_'+str(i)+'_roi_'+str(k+1)+'/trajectories.npy'
                else:
                    trajectories_file_path = '../../data/temp_collective/roi/'+str(i)+'/' +str(j)+'/GS_'+str(j)+'_T_'+str(i)+'_roi_'+str(k+1)+'/trajectories_wo_gaps.npy'
                try:
                    tr = tt.Trajectories.from_idtrackerai(trajectories_file_path,center=True).normalise_by('body_length')
                    tr.new_time_unit(tr.params['frame_rate'], 'seconds')
                except FileNotFoundError:
                    print(i,j,k)
                    print('File not found')
                    continue
                for l in range(1,len(rows1)):
                    if i == int(rows1[l][0]) and j == int(rows1[l][1]) and (k+1) == int(rows1[l][2]):
                        print(i,j,k)
                        for m in range(tr.s.shape[0]-2):
                            for n in range(tr.s.shape[1]):
                                if np.isnan(filter_speed(tr,5)[m][n]) != True:
                                    writer.writerow([i,j,k+1,int(rows1[l][3]),rows1[l][4],int(rows1[l][5]),rows1[l][6],rows1[l][7], m, n+1, filter(tr,5)[0][m][n], filter(tr,5)[1][m][n], filter_speed(tr,5)[m,n], filter_acc(tr,5)[m,n]])


"""
Goal - to produce a csv file with frame, temp, gs, individual, date of exp, trial, subtrial, time, position, speed, acceleration
""" 

import os
import pathlib
from pprint import pprint

import numpy as np
from scipy import stats
from scipy.spatial import distance
import matplotlib.pyplot as plt
from matplotlib.pyplot import figure

import trajectorytools as tt
import trajectorytools.plot as ttplot
import trajectorytools.socialcontext as ttsocial
from trajectorytools.constants import dir_of_data
import csv
import pickle
import argparse

rows1 = []
with open('../../data/temp_collective/roi/all_data.csv', 'r') as csvfile:
    looms = csv.reader(csvfile)
    for row in looms:
        rows1.append(row)




def position(tr):
    return(tr.s)


def speed(tr):
    v = (position(tr)[2:] - position(tr)[:-2]) / 2
    b = np.linalg.norm(v, axis=-1)
    return(b*60)

def acceleration(tr):
    a = position(tr)[2:] - 2 * position(tr)[1:-1] + position(tr)[:-2]
    aa = np.linalg.norm(a, axis=-1)  
    return(aa*3600)
        

def filter(tr, roi = 5): #ind (for individual) starts from 0, roi - edge of region of interest
    position_mask0 = np.ma.masked_where((tr.distance_to_origin[1:-1] > roi)|(tr.distance_to_origin[0:-2] > roi)|(tr.distance_to_origin[2:] > roi), position(tr)[1:-1,:,0],copy=False)
    position_mask1 = np.ma.masked_where((tr.distance_to_origin[1:-1] > roi)|(tr.distance_to_origin[0:-2] > roi)|(tr.distance_to_origin[2:] > roi), position(tr)[1:-1,:,1],copy=False)
    return(position_mask0,position_mask1)  
    
def filter_speed(tr, roi = 5): 
    speed_mask = np.ma.masked_where((tr.distance_to_origin[1:-1] > roi)|(tr.distance_to_origin[0:-2] > roi)|(tr.distance_to_origin[2:] > roi), speed(tr),copy=False)
    
    return(speed_mask)         



def filter_acc(tr, roi = 5): 
    acc_mask = np.ma.masked_where((tr.distance_to_origin[1:-1] > roi)|(tr.distance_to_origin[0:-2] > roi)|(tr.distance_to_origin[2:] > roi), acceleration(tr),copy=False)
    
    return(acc_mask)#[~acc_mask.mask].data)   

temperature = range(9,30,4)

group = [1,2,4,8,16,32]

replication = range(10) # number of replicates per treatment

#output parent directory
parent_dir = '../../output/temp_collective/roi'

ii = 0 # to keep count of temperature

#frames = 5000 #number of frames for which annd is calculated
with open('../../data/temp_collective/roi/masked_data.csv', mode='w') as stats_speed:
    writer = csv.writer(stats_speed, delimiter=',', quotechar='"', quoting=csv.QUOTE_MINIMAL)

    writer.writerow(['Temperature', 'Groupsize', 'Replicate', 'Trial','Date','Subtrial','Time_fish_in','Time_start_record','Frame','Individual','x','y','speed','acceleration'])
    for i in temperature:
    
        jj = 0 # to keep count of groups
        for j in group:
        
            for k in replication:

                if j == 1:
                    trajectories_file_path = '../../data/temp_collective/roi/'+str(i)+'/' +str(j)+'/GS_'+str(j)+'_T_'+str(i)+'_roi_'+str(k+1)+'/trajectories.npy'
                else:
                    trajectories_file_path = '../../data/temp_collective/roi/'+str(i)+'/' +str(j)+'/GS_'+str(j)+'_T_'+str(i)+'_roi_'+str(k+1)+'/trajectories_wo_gaps.npy'
                try:
                    tr = tt.Trajectories.from_idtrackerai(trajectories_file_path,center=True).normalise_by('body_length')
                    tr.new_time_unit(tr.params['frame_rate'], 'seconds')
                except FileNotFoundError:
                    print(i,j,k)
                    print('File not found')
                    continue
                for l in range(1,len(rows1)):
                    if i == int(rows1[l][0]) and j == int(rows1[l][1]) and (k+1) == int(rows1[l][2]):
                        print(i,j,k)
                        for m in range(tr.s.shape[0]-2):
                            for n in range(tr.s.shape[1]):
                                if np.isnan(filter_speed(tr,5)[m][n]) != True:
                                    writer.writerow([i,j,k+1,int(rows1[l][3]),rows1[l][4],int(rows1[l][5]),rows1[l][6],rows1[l][7], m, n+1, filter(tr,5)[0][m][n], filter(tr,5)[1][m][n], filter_speed(tr,5)[m,n], filter_acc(tr,5)[m,n]])


Goal - to produce a csv file with frame, temp, gs, individual, date of exp, trial, subtrial, time, position, speed, acceleration #ind (for individual) starts from 0, roi - edge of region of interest #[~acc_mask.mask].data) # number of replicates per treatment #output parent directory # to keep count of temperature #frames = 5000 #number of frames for which annd is calculated # to keep count of groups