from __future__ import division
import glob
import os
import sys
import shutil
import numpy as np
import csv
import math

# We use PupilLabs Invisible glasses for eye tracking, which provides gaze positions as a csv file with the following columns:
## 'gaze_timestamp': the timestamp of the gaze position, starting at 0 (start of recording).
## 'world_index': number of the frame on the scene camera video
## 'confidence': a quality measure not yet (June 2022) implemented by PupilLabs and therefore contains integers equal to 0. If you don't have this column, create a column with this header and set all values equal to 0.
## 'norm_pos_x': normalized x-position of the gaze
## 'norm_pos_y': normalized y-position of the gaze

##############helper functions################
# Function to calculate the dispersion of points
def get_dispersion(points):
    dispersion = 0
    if len(points) == 0:
        return dispersion
    xx = []
    yy = []
    for i in points:
        xx.append(i[3])
        yy.append(i[4])
    argxmin = np.min(xx)
    argxmax = np.max(xx)

    argymin = np.min(yy)
    argymax = np.max(yy)

    dispersion = ((argxmax - argxmin) + (argymax - argymin)) / 2

    return dispersion

# Function to calculate the mean position of points
def get_sum_x_y(points):
    x = 0
    y = 0
    if len(points) == 0:
        return x, y
    for i in points:
        x += i['norm_pos_x']
        y += i['norm_pos_y']
    return x / len(points), y / len(points)

# Function to calculate the average distance from the mean position
def get_distance(points, x, y):
    if len(points) == 0:
        return 0
    res = 0
    for i in points:
        res += math.sqrt(((i['norm_pos_x'] - x) ** 2) +
                         ((i['norm_pos_y'] - y) ** 2))
    return res / len(points)

# Function to calculate the time duration of a fixation
def get_time(points):
    time = 0
    if len(points) == 0:
        return time
    first = points[0]['gaze_timestamp']
    last = points[len(points) - 1]['gaze_timestamp']

    time = (last - first) / 1000000
    return time
##############helper functions################

# Function to calculate fixation and saccades based on the algorithm proposed by Dario D. Salvucci (https://doi.org/10.1145/355017.355028)
def idt(data, dis_threshold, dur_threshold):
    export = [] # List to store fixation data
    export_line = []  # List to store saccade data
    x = []  # List to store individual fixations
    nano_dur_threshold = dur_threshold * 1000
    first = data[0]['gaze_timestamp'] # First gaze timestamp
    index = 0
    for val in data:
        second_time = val['gaze_timestamp']
        vl = second_time - first  # Duration between consecutive gaze timestamps
        # Check if the duration or dispersion exceeds the thresholds
        if nano_dur_threshold > vl or get_dispersion(x) < dis_threshold:
            x.append(val)
        else:
            new_first = x[0]
            new_last = x[len(x) - 1]
            xx, yy = get_sum_x_y(x)  # Calculate the mean position of the fixations
            xx_d = get_distance(x, xx, yy)  # Calculate the average distance from the mean position
            # Append fixation data to the export list
            export.append(
                [index, get_time(x), int(new_first['world_index']), xx, yy, int(new_first['world_index']),
                 int(new_last['world_index']), xx_d])
            # Append saccade data to the export_line list
            export_line.append([index, new_first['gaze_timestamp'] / 1000000, int(new_first['world_index']),
                                new_first['confidence'], new_first['norm_pos_x'], new_first['norm_pos_y'],
                                new_last['gaze_timestamp'] / 1000000, int(new_last['world_index']),
                                new_last['confidence'], new_last['norm_pos_x'], new_last['norm_pos_y']])
            index += 1
            x = [] # Reset the current fixation
    # Write fixation data to a CSV file
    filename = "fixation_" + gaze_raw_data_file
    with open(filename, 'w', newline='') as csvfile:
        csvwriter = csv.writer(csvfile)
        csvwriter.writerow(fields)
        for i in export:
            if len(i) > 0:
                csvwriter.writerow(i)
    # Write saccade data to a CSV file
    filename = "saccades_" + gaze_raw_data_file
    with open(filename, 'w', newline='') as csvfile:
        csvwriter = csv.writer(csvfile)
        csvwriter.writerow(fields_line)
        for i in export_line:
            if len(i) > 0:
                csvwriter.writerow(i)
    print('file created')

# The output of the idt function are two csv files:
# fixations_....csv:['id', 'time', 'world_index', 'x_mean','y_mean', 'start_frame', 'end_frame', 'dispersion']
## 'time': fixation duration
## 'world_index': frame number of the fixation
## 'x_mean': x coordinate of the fixation
## 'y_mean': y coordinate of the fixation
## 'start_frame': frame number of the first gaze point that formed a fixation
## 'end_frame': frame number of the last gaze point that formed a fixation
##'dispersion': fixation dispersion (radius)
# saccades_....csv:['id', 'first_gaze_timestamp', 'first_world_index', 'first_confidence', 'first_norm_pos_x','first_norm_pos_y', 'last_gaze_timestamp', 'last_world_index', 'last_confidence', 'last_norm_pos_x', 'last_norm_pos_y']
##'first_gaze_timestamp': timestamp of the first gaze point in the first fixation
##'first_world_index': frame number of the first fixation
##'first_confidence': confidence of the first fixation (=0)
##'first_norm_pos_x': x coordinate of the first fixation
##'first_norm_pos_y': y coordinate of the first fixation
##'last_gaze_timestamp': timestamp of the last gaze point in the second fixation
##'last_world_index': frame number of the second fixation
##'last_confidence': confidence of the second fixation (=0)
##'last_norm_pos_x': x coordinate of the second fixation
##'last_norm_pos_y': y coordinate of the second fixation
# With this information you can easily calculate the saccade length, the azimuth, the duration and even the velocity (if the frequency is high enough).


# Iterate over all CSV files in the current directory where you have your gaze positions files and calculate the saccades and fixations for each recording.
for file in glob.glob("*.csv"):
    gaze_raw_data_file = file
    # Load the CSV file into a NumPy array
    gaze_np = np.genfromtxt('{}'.format(gaze_raw_data_file), delimiter=',', skip_header=1,
                            names=['gaze_timestamp', 'world_index', 'confidence', 'norm_pos_x', 'norm_pos_y'])
    # Round the 'gaze_timestamp' values in the array
    for i in gaze_np:
        i['gaze_timestamp'] = round((i['gaze_timestamp'] * 1000000), 2)

    # Define the field names for the output CSV files
    fields = ['id', 'time', 'world_index', 'x_mean',
              'y_mean', 'start_frame', 'end_frame', 'dispersion']
    fields_line = ['id', 'first_gaze_timestamp', 'first_world_index', 'first_confidence', 'first_norm_pos_x',
                   'first_norm_pos_y',
                   'last_gaze_timestamp', 'last_world_index', 'last_confidence', 'last_norm_pos_x', 'last_norm_pos_y']
    # Call the idt function with the specified parameters
    idt(gaze_np, 0.02, 100)
