import os 
import pandas as pd 
import pickle
from tqdm import tqdm


def check_append(obj_name, m1, mind_name, obj_frame, flags, label):
    if label:
        if not obj_name in m1:
            m1[obj_name] = []
            m1[obj_name].append(
                [obj_frame.frame, [obj_frame.x_min, obj_frame.y_min, obj_frame.x_max, obj_frame.y_max], 0])
            flags[mind_name] = 1
        elif not flags[mind_name]:
            m1[obj_name].append(
                [obj_frame.frame, [obj_frame.x_min, obj_frame.y_min, obj_frame.x_max, obj_frame.y_max], 0])
            flags[mind_name] = 1
    else: # false belief 
        if obj_name in m1:
            if flags[mind_name]:
                m1[obj_name].append(
                    [obj_frame.frame, [obj_frame.x_min, obj_frame.y_min, obj_frame.x_max, obj_frame.y_max], 1])
                flags[mind_name] = 0
    return flags, m1


def store_mind_set(clip, annotation_path, save_path):
    if not os.path.exists(save_path):
        os.makedirs(save_path)
    annt = pd.read_csv(annotation_path + clip, sep=",", header=None)
    annt.columns = ["obj_id", "x_min", "y_min", "x_max", "y_max", "frame", "lost", "occluded", "generated", "name",
                    "label"]
    obj_names = annt.name.unique()
    m1, m2, m12, m21, mc = {}, {}, {}, {}, {}
    flags = {'m1':0, 'm2':0, 'm12':0, 'm21':0, 'mc':0}
    for obj_name in obj_names:
        if obj_name == 'P1' or obj_name == 'P2':
            continue
        obj_frames = annt.loc[annt.name == obj_name]
        for index, obj_frame in obj_frames.iterrows():
            if type(obj_frame.label) == float:
                continue
            labels = obj_frame.label.split()
            for label in labels:
                if label == 'in_m1' or label == '"in_m1"':
                    flags, m1 = check_append(obj_name, m1, 'm1', obj_frame, flags, 1)
                elif label == 'in_m2' or label == '"in_m2"':
                    flags, m2 = check_append(obj_name, m2, 'm2', obj_frame, flags, 1)
                elif label == 'in_m12'or label == '"in_m12"':
                    flags, m12 = check_append(obj_name, m12, 'm12', obj_frame, flags, 1)
                elif label == 'in_m21' or label == '"in_m21"':
                    flags, m21 = check_append(obj_name, m21, 'm21', obj_frame, flags, 1)
                elif label == 'in_mc'or label == '"in_mc"':
                    flags, mc = check_append(obj_name, mc, 'mc', obj_frame, flags, 1)
                elif label == 'm1_false' or label == '"m1_false"':
                    flags, m1 = check_append(obj_name, m1, 'm1', obj_frame, flags, 0)
                    flags, m12 = check_append(obj_name, m12, 'm12', obj_frame, flags, 0)
                    flags, m21 = check_append(obj_name, m21, 'm21', obj_frame, flags, 0)
                    false_belief = 'm1_false'
                    with open(save_path + clip.split('.')[0] + '.txt', "a") as file:
                        file.write(f"{obj_frame['frame']},{obj_name},{false_belief}\n")
                elif label == 'm2_false' or label == '"m2_false"':
                    flags, m2 = check_append(obj_name, m2, 'm2', obj_frame, flags, 0)
                    flags, m12 = check_append(obj_name, m12, 'm12', obj_frame, flags, 0)
                    flags, m21 = check_append(obj_name, m21, 'm21', obj_frame, flags, 0)
                    false_belief = 'm2_false'
                    with open(save_path + clip.split('.')[0] + '.txt', "a") as file:
                        file.write(f"{obj_frame['frame']},{obj_name},{false_belief}\n")
                elif label == 'm12_false' or label == '"m12_false"':
                    flags, m12 = check_append(obj_name, m12, 'm12', obj_frame, flags, 0)
                    flags, mc = check_append(obj_name, mc, 'mc', obj_frame, flags, 0)
                    false_belief = 'm12_false'
                    with open(save_path + clip.split('.')[0] + '.txt', "a") as file:
                        file.write(f"{obj_frame['frame']},{obj_name},{false_belief}\n")
                elif label == 'm21_false' or label == '"m21_false"':
                    flags, m21 = check_append(obj_name, m2, 'm21', obj_frame, flags, 0)
                    flags, mc = check_append(obj_name, mc, 'mc', obj_frame, flags, 0)
                    false_belief = 'm21_false'
                    with open(save_path + clip.split('.')[0] + '.txt', "a") as file:
                        file.write(f"{obj_frame['frame']},{obj_name},{false_belief}\n")
    # print('m1', m1)
    # print('m2', m2)
    # print('m12', m12)
    # print('m21', m21)
    # print('mc', mc)
    #with open(save_path + clip.split('.')[0] + '.p', 'wb') as f:
    #    pickle.dump([m1, m2, m12, m21, mc], f)


if __name__ == "__main__":

    annotation_path = '/scratch/bortoletto/data/tbd/reformat_annotation/'
    save_path = '/scratch/bortoletto/data/tbd/store_mind_set/'

    for clip in tqdm(os.listdir(annotation_path), desc="Processing videos", unit="item"):
        store_mind_set(clip, annotation_path, save_path)