import numpy as np

import pandas as pd

import glob

def setup():

participant_dirs = glob.glob("./data/p*")

participant_dirs = [path for path in participant_dirs if not path.endswith(".meta")]

# preparing and cleaning all data

allData = pd.DataFrame()

for id, folder in enumerate(participant_dirs):

df = pd.read_csv(f'{folder}/main.csv')

df['id'] = id

df['folder'] = folder

allData = pd.concat([allData, df])

allData = allData.astype({'trialNum': 'int64', 'soundAngleR': 'float64', 'soundAngleT': 'float64', 'soundAngleP': 'float64', 'headAngleX': 'float64', 'headAngleY': 'float64', 'headAngleZ': 'float64', 'initialHeadAngleX': 'float64', 'initialHeadAngleY': 'float64', 'initialHeadAngleZ': 'float64'})

# clean up "error" trials

errors = allData.id[allData['id'] == "ERROR"]

for errorNum in errors.index:

allData = allData.drop(errorNum)

allData = allData.drop(errorNum-1)

allData['soundAngleTDeg'] = 360 - np.degrees(allData['soundAngleT']) + 90

allData['soundAngleTDeg'] = np.where(allData['soundAngleTDeg'] < 0, allData['soundAngleTDeg'] + 360, allData['soundAngleTDeg'])

allData['soundAngleTDeg'] = np.where(allData['soundAngleTDeg'] >= 360, allData['soundAngleTDeg'] - 360, allData['soundAngleTDeg'])

allData['soundAnglePDeg'] = np.degrees(allData['soundAngleP'])

allData['headAngleXAdjust'] = (360 - allData['headAngleX'] + 90) % 360 - 90

allData['diffY'] = abs(allData['soundAngleTDeg'] - allData['headAngleY'])

allData['diffY'] = np.where(allData['diffY'] >= 180, allData['diffY'] - 180, allData['diffY'])

allData['diffX'] = abs(allData['soundAnglePDeg'] - allData['headAngleXAdjust'])

allData['diffX'] = np.where(allData['diffX'] >= 180, allData['diffX'] - 180, allData['diffX'])

return allData

def rectangular_to_spherical(x, y, z):

r = np.sqrt(x**2 + y**2 + z**2)

theta = np.arccos(z / r) # Inclination angle (0 to pi)

phi = np.arctan2(y, x) # Azimuthal angle (-pi to pi)

return r, theta, phi

def spherical_to_rectangular(r, theta, phi):

theta = np.deg2rad(theta)

phi = np.deg2rad(phi)

# phi: -pi/3 to pi/3 (-60 to 60)

phi = -(phi - np.pi/2)

z = r * np.sin(phi) * np.cos(theta)

x = r * np.sin(phi) * np.sin(theta)

y = r * np.cos(phi)

# x, y, z in Unity coordinate system

return x, y, z