# main imports
import sys, os, argparse
import math
import numpy as np
import pickle
import time

# processing imports
from PIL import Image

import matplotlib.pyplot as plt
import scipy.stats as stats

# modules imports
sys.path.insert(0, '') # trick to enable import of main folder module

import custom_config as cfg
import utils as utils_functions

# variables
data_expe_folder          = cfg.data_expe_folder
position_file_pattern     = cfg.position_file_pattern
click_line_pattern        = cfg.click_line_pattern

min_x                     = cfg.min_x_coordinate
min_y                     = cfg.min_y_coordinate
image_scene_size          = cfg.image_scene_size
scene_width, scene_height = image_scene_size


def main():

    parser = argparse.ArgumentParser(description="Compute expe data into output file")

    parser.add_argument('--n', type=int, help="`n` first clicks", required=True)
    parser.add_argument('--folder', type=str, help="output folder expected", required=True)
    parser.add_argument('--reverse', type=int, help="reverse or not y axis clicks", default=False)

    args = parser.parse_args()

    p_n        = args.n
    p_folder   = args.folder
    p_reverse  = bool(args.reverse)

    # list all folders
    subjects = os.listdir(data_expe_folder)

    print('Number of subjects', len(subjects))

    output_folder_path = os.path.join(cfg.media_data_folder, p_folder)

    if not os.path.exists(output_folder_path):
        os.makedirs(output_folder_path)

    # keep scene_data in memory
    scenes_data = {}   

    for scene in cfg.scenes_names:
        
        scenes_data[scene] = {}
        # construct for each scene
        scenes_data[scene]['x'] = []
        scenes_data[scene]['y'] = []   

        
    for index, subject in enumerate(subjects):
        
        subject_folder = os.path.join(data_expe_folder, subject)
        data_files = os.listdir(subject_folder)

        pos_file = [f for f in data_files if position_file_pattern in f][0]
        
        pos_filepath = os.path.join(subject_folder, pos_file)

        previous_path_scene = ""
        path_scene          = ""
        new_scene           = True
        number_of_scenes    = 0
        counter             = 0
        scene_name          = ""

        print('Extract images clicks of subject', subject)

        # open pos file and extract click information 
        with open(pos_filepath, 'r') as f:
            
            x_points = []
            y_points = []

            for line in f.readlines():
                
                if click_line_pattern in line and scene_name in cfg.scenes_names:
                    
                    x, y = utils_functions.extract_click_coordinate(line)

                    p_x = x - min_x
                    p_y = y - min_y

                    # only accept valid coordinates
                    if utils_functions.check_coordinates(p_x, p_y):
                        
                        if p_reverse:
                            # add reversed points here
                            p_y = scene_height - p_y

                        if counter < p_n:

                            # add points here
                            x_points.append(p_x)
                            y_points.append(p_y)

                            counter += 1
                
                elif click_line_pattern not in line:
                    path_scene = line

                if previous_path_scene != path_scene:

                    previous_path_scene = path_scene
                    new_scene = True
                    scene_name = path_scene.split('/')[4]

                    if scene_name in cfg.scenes_names:
                        number_of_scenes += 1

                        if previous_path_scene != "":

                            subject_path = os.path.join(output_folder_path, subject)
                            
                            if not os.path.exists(subject_path):
                                os.makedirs(subject_path)     

                            output_image_name = subject + '_' + scene_name + '_' + str(p_n) + '.png'
                            img_path = os.path.join(subject_path, output_image_name)

                            title = subject + ' - ' + scene_name + ' (' + str(p_n) + ' clicks)'

                            # save image plot
                            utils_functions.save_img_plot(scene_name, x_points, y_points, title, img_path)

                            # save scene data
                            scenes_data[scene_name]['x'] = scenes_data[scene_name]['x'] + x_points
                            scenes_data[scene_name]['y'] = scenes_data[scene_name]['y'] + y_points

                else:
                    new_scene = False

                if new_scene:
                    counter = 0
                    x_points = []
                    y_points = []

    all_path_folder = os.path.join(output_folder_path, cfg.all_subjects_data_folder)

    print('Merge images clicks of subjects into', all_path_folder)
    if not os.path.exists(all_path_folder):
        os.makedirs(all_path_folder)

    for k, v in scenes_data.items():

        current_x_points = v['x']
        current_y_points = v['y']

        title = k + ' scene with all subjects (with ' + str(p_n) + ' clicks per subject)'

        img_filename = cfg.all_subjects_data_folder + '_' + k + '_' + str(p_n) + '.png'
        img_path = os.path.join(all_path_folder, img_filename)

        # save figure `all` subjects `p_n` clicks
        utils_functions.save_img_plot(k, current_x_points, current_y_points, title, img_path)


    print('Images are saved into', output_folder_path)
    

if __name__== "__main__":
    main()