import pandas as pd
import os
import glob
import pickle

DATASET_LOCATION = "YOUR_PATH_HERE"

def reframe_annotation():
    annotation_path = f'{DATASET_LOCATION}/retrieve_annotation/all/'
    save_path = f'{DATASET_LOCATION}/reformat_annotation/'
    if not os.path.exists(save_path):
        os.makedirs(save_path)
    tasks = glob.glob(annotation_path + '*.txt')
    id_map = pd.read_csv('id_map.csv')
    for task in tasks:
        if not task.split('/')[-1].split('_')[2] == '1.txt':
            continue
        with open(task, 'r') as f:
            lines = f.readlines()
        task_id = int(task.split('/')[-1].split('_')[1]) + 1
        clip = id_map.loc[id_map['ID'] == task_id].folder
        print(task_id, len(clip))
        if len(clip) == 0:
            continue
        with open(save_path + clip.item() + '.txt', 'w') as f:
            for line in lines:
                words = line.split()
                f.write(words[0] + ',' + words[1] + ',' + words[2] + ',' + words[3] + ',' + words[4] + ',' + words[5] +
                        ',' + words[6] + ',' + words[7] + ',' + words[8] + ',' + words[9] + ',' + ' '.join(words[10:]) + '\n')
        f.close()

def get_grid_location(obj_frame):
    x_min = obj_frame['x_min']#.item()
    y_min = obj_frame['y_min']#.item()
    x_max = obj_frame['x_max']#.item()
    y_max = obj_frame['y_max']#.item()
    gridLW = 1280 / 25.
    gridLH = 720 / 15.
    center_x, center_y = (x_min + x_max)/2, (y_min + y_max)/2
    X, Y = int(center_x / gridLW), int(center_y / gridLH)
    return X, Y

def regenerate_annotation():
    annotation_path = f'{DATASET_LOCATION}/reformat_annotation/'
    save_path=f'{DATASET_LOCATION}/regenerate_annotation/'
    if not os.path.exists(save_path):
        os.makedirs(save_path)
    tasks = glob.glob(annotation_path + '*.txt')
    for task in tasks:
        print(task)
        annt = pd.read_csv(task, sep=",", header=None)
        annt.columns = ["obj_id", "x_min", "y_min", "x_max", "y_max", "frame", "lost", "occluded", "generated", "name", "label"]
        obj_records = {}
        for index, obj_frame in annt.iterrows():
            if obj_frame['name'].startswith('P'):
                continue
            else:
                assert obj_frame['name'].startswith('O')
            obj_name = obj_frame['name']
            # 0: enter 1: disappear 2: update 3: unchange
            frame_id = obj_frame['frame']
            curr_loc = get_grid_location(obj_frame)
            mind_dict = {'m1': {'fluent': 3, 'loc': None}, 'm2': {'fluent': 3, 'loc': None},
                         'm12': {'fluent': 3, 'loc': None},
                         'm21': {'fluent': 3, 'loc': None}, 'mc': {'fluent': 3, 'loc': None},
                         'mg': {'fluent': 3, 'loc': curr_loc}}
            mind_dict['mg']['loc'] = curr_loc
            if not type(obj_frame['label']) == float:
                mind_labels = obj_frame['label'].split()
                for mind_label in mind_labels:
                    if mind_label == 'in_m1' or mind_label == 'in_m2' or mind_label == 'in_m12' \
                        or mind_label == 'in_m21' or mind_label == 'in_mc' or mind_label == '"in_m1"' or mind_label == '"in_m2"'\
                            or mind_label == '"in_m12"' or mind_label == '"in_m21"' or mind_label == '"in_mc"':
                        mind_name = mind_label.split('_')[1].split('"')[0]
                        mind_dict[mind_name]['loc'] = curr_loc
                    else:
                        mind_name = mind_label.split('_')[0].split('"')
                        if len(mind_name) > 1:
                            mind_name = mind_name[1]
                        else:
                            mind_name = mind_name[0]
                        last_loc = obj_records[obj_name][frame_id - 1][mind_name]['loc']
                        mind_dict[mind_name]['loc'] = last_loc

            for mind_name in mind_dict.keys():
                if frame_id > 0:
                    curr_loc = mind_dict[mind_name]['loc']
                    last_loc = obj_records[obj_name][frame_id - 1][mind_name]['loc']
                    if last_loc is None and curr_loc is not None:
                        mind_dict[mind_name]['fluent'] = 0
                    elif last_loc is not None and curr_loc is None:
                        mind_dict[mind_name]['fluent'] = 1
                    elif not last_loc ==  curr_loc:
                        mind_dict[mind_name]['fluent'] = 2
            if obj_name not in obj_records:
                obj_records[obj_name] = [mind_dict]
            else:
                obj_records[obj_name].append(mind_dict)

        with open(save_path + task.split('/')[-1].split('.')[0] + '.p', 'wb') as f:
            pickle.dump(obj_records, f)


if __name__ == '__main__':
    reframe_annotation()
    regenerate_annotation()