import os, sys
from pathlib import Path

import numpy as np
import pandas as pd

import matplotlib.pyplot as plt
import matplotlib.patches as patches

try:
    from data_handle.mid_object import *
except:
    sys.path.append(os.path.join(os.path.dirname(__file__), '..'))
    from data_handle.mid_object import *

'''
There are 3 kinds of folder structure:
1. [1c] data_dir - obj
2. [2c] data_dir - obj - obj&env (each object has its own environment)
3. [2c] data_dir - obj&env       (all objects share one environment)
(2 & 3 is compatible with 1 to some extent)

Functions 'gen_csv_trackers' and 'gather_all_data' are universal.
They only depend on the data structure version.
'''

def gen_csv_trackers(data_dir, channel_per_image=2): # Data structure 2
    # data_dir - obj - obj&env (each object has its own environment)
    cpi = channel_per_image
    obj_folders = os.listdir(data_dir)
    for objf in obj_folders:
        if cpi == 1:
            obj_files = os.listdir(data_dir+objf) # all files/images under this folder
        else:
            obj_files = os.listdir(os.path.join(data_dir,objf,'obj')) # all files/images under this folder
        t_list = []   # time or time step
        x_list = []   # x coordinate
        y_list = []   # y coordinate
        idx_list = [] # more information (e.g. scene index)
        invalid_files = []
        for f in obj_files:
            info = f[:-4] # the last for characters are filename extensions
            try:
                t_list.append(int(info.split('_')[1]))
                x_list.append(float(info.split('_')[2]))
                y_list.append(float(info.split('_')[3]))
                idx_list.append(int(info.split('_')[4]))
            except:
                invalid_files.append(f)
                continue
        for f in invalid_files:
            obj_files.remove(f)
        df = pd.DataFrame({'f':obj_files,'t':t_list,'x':x_list,'y':y_list, 'index':idx_list}).sort_values(by='t', ignore_index=True)
        if cpi == 1:
            df.to_csv(os.path.join(data_dir, objf, '/data.csv'), index=False)
        else:
            df.to_csv(os.path.join(data_dir, objf, 'obj/data.csv'), index=False)

def gather_all_data(data_dir, past, maxT, channel_per_image=2, minT=1, period=1, save_dir=None): # Data structure 2
    # data_dir  -  objf(1,2,...) - obj&env
    if save_dir is None:
        save_dir = data_dir
    cpi = channel_per_image

    column_name = [f'f{i}' for i in range(0,cpi*(past+1))] + ['T', 'x', 'y', 'index']
    df_all = pd.DataFrame(columns=column_name)
    obj_folders = os.listdir(data_dir)
    cnt = 0
    for objf in obj_folders:
        cnt += 1
        print(f'\rProcess {cnt}/{len(obj_folders)}', end='    ')
        if cpi == 1:
            df_obj = pd.read_csv(os.path.join(data_dir, objf, 'data.csv')) # generated by "gen_csv_trackers"
        else:
            df_obj = pd.read_csv(os.path.join(data_dir, objf, 'obj/data.csv')) # generated by "gen_csv_trackers"
        for T in range(minT,maxT+1):
            sample_list = []
            for i in range(len(df_obj)-past*period-T): # each sample
                sample = []
                ################## Sample START ##################
                for j in range(past+1):
                    obj_filename = df_obj.iloc[i+j*period]['f']
                    sample.append(obj_filename)
                    if cpi == 2:
                        sample.append(obj_filename.split('_')[0]+'_'+obj_filename.split('_')[1]+'_'+obj_filename.split('_')[-1])
                sample.append(T)
                sample.append(df_obj.iloc[i+past+T]['x'])
                sample.append(df_obj.iloc[i+past+T]['y'])
                sample.append(df_obj.iloc[i+past+T]['index'])
                ################## Sample E N D ##################
                sample_list.append(sample)
            df_T = pd.DataFrame(sample_list, columns=df_all.columns)
            df_all = pd.concat([df_all, df_T], ignore_index=True)
    df_all.to_csv(os.path.join(save_dir, 'all_data.csv'), index=False)


def save_MID_data(index_list, save_path, sim_time_per_scene:int, channel_per_image=2, dots_per_inch=None):
    # MID - Multiple-scene Interaction Dataset
    cpi = channel_per_image
    dpi = dots_per_inch
    
    cnt = 0
    overall_sim_time = sim_time_per_scene * len(index_list)
    for idx in index_list:
        boundary_coords, obstacle_list, nchoices = return_Map(index=idx) # map parameters
        target_size, ts = (0.5, 0.2) # object parameters
        
        choice_list = list(range(1,nchoices+1))
        for ch in choice_list:
            for i in range(sim_time_per_scene//nchoices):
                cnt += 1
                print(f'\rSimulating: {cnt}/{overall_sim_time}   ', end='')
                inflation = 0
                stagger = 0.4   + (random.randint(0, 20)/10-1) * 0.2
                vmax = 1        + (random.randint(0, 20)/10-1) * 0.3
                if i<((sim_time_per_scene//nchoices)//2):
                    ref_path = get_ref_path(index=idx, choice=ch, reverse=False)
                else:
                    ref_path = get_ref_path(index=idx, choice=ch, reverse=True)
                    
                graph = Graph(boundary_coords, obstacle_list, inflation=inflation)

                obj = MovingObject(ref_path[0], stagger=stagger)
                obj.run(ref_path, ts, vmax)

                ### Generate images
                for j, tr in enumerate(obj.traj):
                    # images containing everything
                    shape = patches.Circle(tr, radius=target_size/2, fc='k')

                    if cpi == 1:
                        # images containing everything
                        _, ax = plt.subplots(dpi=dpi)
                        graph.plot_map(ax, clean=True) ### NOTE change this
                        ax.add_patch(shape)
                        ax.set_aspect('equal', 'box')
                        ax.axis('off')
                        if save_path is None:
                            plt.show()
                        else:
                            folder = os.path.join(save_path,f'{cnt}/')
                            Path(folder).mkdir(parents=True, exist_ok=True)
                            plt.savefig(os.path.join(folder,f'{cnt}_{j}_{round(tr[0],4)}_{round(tr[1],4)}_{idx}.png'), 
                                        bbox_inches='tight', pad_inches=0, dpi=dpi)
                            plt.close()
                    elif cpi == 2:
                        # images containing only object
                        _, ax1 = plt.subplots(dpi=dpi)
                        graph.plot_map(ax1, empty=True) ### NOTE change this
                        ax1.add_patch(shape)
                        ax1.set_aspect('equal', 'box')
                        ax1.axis('off')
                        if save_path is None:
                            plt.show()
                        else:
                            folder = os.path.join(save_path,f'{cnt}/obj')
                            Path(folder).mkdir(parents=True, exist_ok=True)
                            plt.savefig(os.path.join(folder,f'{cnt}_{j}_{round(tr[0],4)}_{round(tr[1],4)}_{idx}.png'), 
                                        bbox_inches='tight', pad_inches=0, dpi=dpi)
                            plt.close()
                        # images containing env
                        _, ax2 = plt.subplots(dpi=dpi)
                        graph.plot_map(ax2, clean=True) ### NOTE change this
                        ax2.set_aspect('equal', 'box')
                        ax2.axis('off')
                        if save_path is None:
                            plt.show()
                        else:
                            folder = os.path.join(save_path,f'{cnt}/env')
                            Path(folder).mkdir(parents=True, exist_ok=True)
                            plt.savefig(os.path.join(folder,f'{cnt}_{j}_{idx}.png'), 
                                        bbox_inches='tight', pad_inches=0, dpi=dpi)
                            plt.close()

                    else:
                        raise(ModuleNotFoundError('CPI must be 1 or 2.'))
    print()

import os, sys
from pathlib import Path

import numpy as np
import pandas as pd

import matplotlib.pyplot as plt
import matplotlib.patches as patches

try:
    from data_handle.mid_object import *
except:
    sys.path.append(os.path.join(os.path.dirname(__file__), '..'))
    from data_handle.mid_object import *

'''
There are 3 kinds of folder structure:
1. [1c] data_dir - obj
2. [2c] data_dir - obj - obj&env (each object has its own environment)
3. [2c] data_dir - obj&env       (all objects share one environment)
(2 & 3 is compatible with 1 to some extent)

Functions 'gen_csv_trackers' and 'gather_all_data' are universal.
They only depend on the data structure version.
'''

def gen_csv_trackers(data_dir, channel_per_image=2): # Data structure 2
    # data_dir - obj - obj&env (each object has its own environment)
    cpi = channel_per_image
    obj_folders = os.listdir(data_dir)
    for objf in obj_folders:
        if cpi == 1:
            obj_files = os.listdir(data_dir+objf) # all files/images under this folder
        else:
            obj_files = os.listdir(os.path.join(data_dir,objf,'obj')) # all files/images under this folder
        t_list = []   # time or time step
        x_list = []   # x coordinate
        y_list = []   # y coordinate
        idx_list = [] # more information (e.g. scene index)
        invalid_files = []
        for f in obj_files:
            info = f[:-4] # the last for characters are filename extensions
            try:
                t_list.append(int(info.split('_')[1]))
                x_list.append(float(info.split('_')[2]))
                y_list.append(float(info.split('_')[3]))
                idx_list.append(int(info.split('_')[4]))
            except:
                invalid_files.append(f)
                continue
        for f in invalid_files:
            obj_files.remove(f)
        df = pd.DataFrame({'f':obj_files,'t':t_list,'x':x_list,'y':y_list, 'index':idx_list}).sort_values(by='t', ignore_index=True)
        if cpi == 1:
            df.to_csv(os.path.join(data_dir, objf, '/data.csv'), index=False)
        else:
            df.to_csv(os.path.join(data_dir, objf, 'obj/data.csv'), index=False)

def gather_all_data(data_dir, past, maxT, channel_per_image=2, minT=1, period=1, save_dir=None): # Data structure 2
    # data_dir  -  objf(1,2,...) - obj&env
    if save_dir is None:
        save_dir = data_dir
    cpi = channel_per_image

    column_name = [f'f{i}' for i in range(0,cpi*(past+1))] + ['T', 'x', 'y', 'index']
    df_all = pd.DataFrame(columns=column_name)
    obj_folders = os.listdir(data_dir)
    cnt = 0
    for objf in obj_folders:
        cnt += 1
        print(f'\rProcess {cnt}/{len(obj_folders)}', end='    ')
        if cpi == 1:
            df_obj = pd.read_csv(os.path.join(data_dir, objf, 'data.csv')) # generated by "gen_csv_trackers"
        else:
            df_obj = pd.read_csv(os.path.join(data_dir, objf, 'obj/data.csv')) # generated by "gen_csv_trackers"
        for T in range(minT,maxT+1):
            sample_list = []
            for i in range(len(df_obj)-past*period-T): # each sample
                sample = []
                ################## Sample START ##################
                for j in range(past+1):
                    obj_filename = df_obj.iloc[i+j*period]['f']
                    sample.append(obj_filename)
                    if cpi == 2:
                        sample.append(obj_filename.split('_')[0]+'_'+obj_filename.split('_')[1]+'_'+obj_filename.split('_')[-1])
                sample.append(T)
                sample.append(df_obj.iloc[i+past+T]['x'])
                sample.append(df_obj.iloc[i+past+T]['y'])
                sample.append(df_obj.iloc[i+past+T]['index'])
                ################## Sample E N D ##################
                sample_list.append(sample)
            df_T = pd.DataFrame(sample_list, columns=df_all.columns)
            df_all = pd.concat([df_all, df_T], ignore_index=True)
    df_all.to_csv(os.path.join(save_dir, 'all_data.csv'), index=False)


def save_MID_data(index_list, save_path, sim_time_per_scene:int, channel_per_image=2, dots_per_inch=None):
    # MID - Multiple-scene Interaction Dataset
    cpi = channel_per_image
    dpi = dots_per_inch
    
    cnt = 0
    overall_sim_time = sim_time_per_scene * len(index_list)
    for idx in index_list:
        boundary_coords, obstacle_list, nchoices = return_Map(index=idx) # map parameters
        target_size, ts = (0.5, 0.2) # object parameters
        
        choice_list = list(range(1,nchoices+1))
        for ch in choice_list:
            for i in range(sim_time_per_scene//nchoices):
                cnt += 1
                print(f'\rSimulating: {cnt}/{overall_sim_time}   ', end='')
                inflation = 0
                stagger = 0.4   + (random.randint(0, 20)/10-1) * 0.2
                vmax = 1        + (random.randint(0, 20)/10-1) * 0.3
                if i<((sim_time_per_scene//nchoices)//2):
                    ref_path = get_ref_path(index=idx, choice=ch, reverse=False)
                else:
                    ref_path = get_ref_path(index=idx, choice=ch, reverse=True)
                    
                graph = Graph(boundary_coords, obstacle_list, inflation=inflation)

                obj = MovingObject(ref_path[0], stagger=stagger)
                obj.run(ref_path, ts, vmax)

                ### Generate images
                for j, tr in enumerate(obj.traj):
                    # images containing everything
                    shape = patches.Circle(tr, radius=target_size/2, fc='k')

                    if cpi == 1:
                        # images containing everything
                        _, ax = plt.subplots(dpi=dpi)
                        graph.plot_map(ax, clean=True) ### NOTE change this
                        ax.add_patch(shape)
                        ax.set_aspect('equal', 'box')
                        ax.axis('off')
                        if save_path is None:
                            plt.show()
                        else:
                            folder = os.path.join(save_path,f'{cnt}/')
                            Path(folder).mkdir(parents=True, exist_ok=True)
                            plt.savefig(os.path.join(folder,f'{cnt}_{j}_{round(tr[0],4)}_{round(tr[1],4)}_{idx}.png'), 
                                        bbox_inches='tight', pad_inches=0, dpi=dpi)
                            plt.close()
                    elif cpi == 2:
                        # images containing only object
                        _, ax1 = plt.subplots(dpi=dpi)
                        graph.plot_map(ax1, empty=True) ### NOTE change this
                        ax1.add_patch(shape)
                        ax1.set_aspect('equal', 'box')
                        ax1.axis('off')
                        if save_path is None:
                            plt.show()
                        else:
                            folder = os.path.join(save_path,f'{cnt}/obj')
                            Path(folder).mkdir(parents=True, exist_ok=True)
                            plt.savefig(os.path.join(folder,f'{cnt}_{j}_{round(tr[0],4)}_{round(tr[1],4)}_{idx}.png'), 
                                        bbox_inches='tight', pad_inches=0, dpi=dpi)
                            plt.close()
                        # images containing env
                        _, ax2 = plt.subplots(dpi=dpi)
                        graph.plot_map(ax2, clean=True) ### NOTE change this
                        ax2.set_aspect('equal', 'box')
                        ax2.axis('off')
                        if save_path is None:
                            plt.show()
                        else:
                            folder = os.path.join(save_path,f'{cnt}/env')
                            Path(folder).mkdir(parents=True, exist_ok=True)
                            plt.savefig(os.path.join(folder,f'{cnt}_{j}_{idx}.png'), 
                                        bbox_inches='tight', pad_inches=0, dpi=dpi)
                            plt.close()

                    else:
                        raise(ModuleNotFoundError('CPI must be 1 or 2.'))
    print()

There are 3 kinds of folder structure: 1. [1c] data_dir - obj 2. [2c] data_dir - obj - obj&env (each object has its own environment) 3. [2c] data_dir - obj&env       (all objects share one environment) (2 & 3 is compatible with 1 to some extent)  Functions 'gen_csv_trackers' and 'gather_all_data' are universal. They only depend on the data structure version. # Data structure 2 # data_dir - obj - obj&env (each object has its own environment) # all files/images under this folder # all files/images under this folder # time or time step # x coordinate # y coordinate # more information (e.g. scene index) # the last for characters are filename extensions # Data structure 2 # data_dir  -  objf(1,2,...) - obj&env # generated by "gen_csv_trackers" # generated by "gen_csv_trackers" # each sample ################## Sample START ################## ################## Sample E N D ################## # MID - Multiple-scene Interaction Dataset # map parameters # object parameters ### Generate images # images containing everything # images containing everything ### NOTE change this # images containing only object ### NOTE change this # images containing env ### NOTE change this