# This files creates all the needed data structures that relab needs
# User gives the option for training and testing mode

import numpy as np
import os
import xml.etree.ElementTree as ET
from numpy import linalg as LA
import matplotlib.pyplot as plt
import cv2
from faster_rcnn import network
from faster_rcnn.faster_rcnn import FasterRCNN
from faster_rcnn.utils.timer import Timer
import pickle
from numpy import unravel_index
from test2 import *
from faster_rcnn.nms.py_cpu_nms import py_cpu_nms

repo_of_ground_truth = '/home/revan/VOCdevkit/VOC2007/Annotations'
repo_of_images = '/home/revan/VOCdevkit/VOC2007/JPEGImages'
train_set = '/home/revan/VOCdevkit/VOC2007/ImageSets/Main/trainval.txt'
test_set = '/home/revan/VOCdevkit/VOC2007/ImageSets/Main/test.txt'


classes_dict = {'__background__': 0, 'aeroplane': 1, 'bicycle': 2, 'bird': 3, 'boat': 4, 'bottle': 5, 'bus': 6, 'car': 7,
           'cat': 8, 'chair': 9, 'cow': 10, 'diningtable': 11, 'dog': 12, 'horse': 13, 'motorbike': 14,
           'person': 15, 'pottedplant': 16, 'sheep': 17, 'sofa': 18, 'train': 19, 'tvmonitor': 20}

imdb_name = 'voc_2007_test'
cfg_file = 'experiments/cfgs/faster_rcnn_end2end.yml'
trained_model = 'VGGnet_fast_rcnn_iter_70000.h5'
rand_seed = 1024
save_name = 'faster_rcnn_100000'
max_per_image = 300
thresh = 0.05
vis = False


def main(train=1, serialize=1):
    distribution_mode = 'softmax'  # the mode can be 'peak', k_peak or 'softmax'
    if train:
        if serialize:
            content = get_filenames_of_set(train_set)
            p, gt, width_and_height, all_boxes_all_images, p_binary = create_p_and_gt(content, repo_of_images, classes_dict, distribution_mode)

            # dump data structures into pickle files
            with open('p_trainval.pickle', 'wb') as f:
                pickle.dump(p, f, pickle.HIGHEST_PROTOCOL)
            with open('gt_trainval.pickle', 'wb') as f:
                pickle.dump(gt, f, pickle.HIGHEST_PROTOCOL)
            with open('width_and_height_traival.pickle', 'wb') as f:
                pickle.dump(width_and_height, f, pickle.HIGHEST_PROTOCOL)
            with open('all_boxes_all_images_trainval.pickle', 'wb') as f:
                pickle.dump(all_boxes_all_images, f, pickle.HIGHEST_PROTOCOL)
            with open('p_binary_trainval.pickle', 'wb') as f:
                pickle.dump(p_binary, f, pickle.HIGHEST_PROTOCOL)

        # load the pickle files
        with open('p_trainval.pickle', 'rb') as f:
            p = pickle.load(f)
        with open('gt_trainval.pickle', 'rb') as f:
            gt = pickle.load(f)
        with open('width_and_height_traival.pickle', 'rb') as f:
            width_and_height = pickle.load(f)
        with open('all_boxes_all_images_trainval.pickle', 'rb') as f:
            all_boxes_all_images = pickle.load(f)
        with open('p_binary_trainval.pickle', 'rb') as f:
            p_binary = pickle.load(f)


        if serialize:
            # do the matching between rcnn results and the ground truth
            info_all_images = postprocess_all_images(p, gt, width_and_height, all_boxes_all_images, p_binary, false_positives=False, false_negatives=False)

            # pickle the final data structure
            with open('train_k_peak_softmax_nofp_nofn_including_binary_thresholded_iou.pickle', 'wb') as f:
                pickle.dump(info_all_images, f, pickle.HIGHEST_PROTOCOL)

        # load the final data structure
        with open('train_k_peak_softmax_nofp_nofn_including_binary_thresholded_iou.pickle', 'rb') as f:
            info_all_images = pickle.load(f)

    else:
        if serialize:
            content = get_filenames_of_set(test_set)
            p, gt, width_and_height, all_boxes_all_images, p_binary = create_p_and_gt(content, repo_of_images, classes_dict, distribution_mode)