import pickle
from os.path import exists

import pandas as pd

import constants


def create_annotations_dict():
    columns = ['filename', 'xmin', 'ymin', 'xmax', 'ymax', 'label']
    annotations_df = pd.read_csv(constants.LABELS_PATH, sep=';', names=columns)
    annotations_dict = {}
    for index, bbox in annotations_df.iterrows():
        image_name = bbox.at['filename'].replace('.ppm', '')
        annotations_df.at[index, 'filename'] = image_name
        if image_name not in annotations_dict:
            annotations_dict[image_name] = []

        annotations_dict[image_name].append(
            list(annotations_df.iloc[index][['xmin', 'ymin', 'xmax', 'ymax', 'label']]))

    pickle.dump(annotations_dict, open(constants.ANNOTATIONS_DICT, "wb"))
    return annotations_dict


def get_annotations_dict():
    if exists(constants.ANNOTATIONS_DICT):
        return pickle.load(open(constants.ANNOTATIONS_DICT, "rb"))
    else:
        return create_annotations_dict()