import pickle as pkl
from exp_config import experiment_config, expcfg_from_file
from virat_dataset import ViratDataset
import numpy as np
import cv2
import argparse

def generate_bias_maps(virat_dataset, is_test=False):
    if is_test:
        gttubes = virat_dataset._test_gttubes
    else:
        gttubes = virat_dataset._train_gttubes

    dataset_heatmaps = np.zeros([virat_dataset.nlabels-1, 1080, 1920])
    for video_name, frames in gttubes.iteritems():
        for class_index, activities in frames.iteritems():
            for activity in activities:
                for bbox in activity:
                    _, x1, y1, x2, y2 = bbox.astype(int)
                    dataset_heatmaps[class_index-1, y1:y2, x1:x2] += 1

    overall_heatmap = np.zeros([1080, 1920])
    overall_frame_count = 0
    for index in range(dataset_heatmaps.shape[0]):
        dataset_heatmap = dataset_heatmaps[index, :, :]
        dataset_heatmap = (dataset_heatmap - np.min(dataset_heatmap))/(np.max(dataset_heatmap) - np.min(dataset_heatmap))
        dataset_heatmap = dataset_heatmap * 255.0
        class_name = virat_dataset.labels.keys()[virat_dataset.labels.values().index(index+1)]
        if is_test:
            image_name = class_name + "_test.png" 
        else:
            image_name = class_name + ".png"
        cv2.imwrite(image_name, dataset_heatmap.astype(np.uint8))
        overall_heatmap += dataset_heatmap
    overall_heatmap = (overall_heatmap - np.min(overall_heatmap))/(np.max(overall_heatmap) - np.min(overall_heatmap))
    overall_heatmap *= 255.0
    if is_test:
        image_name = "overall_test.png"
    else:
        image_name = "overall.png"
    cv2.imwrite(image_name, overall_heatmap.astype(np.uint8))

def main(args):
    if args.exp is not None:
        expcfg_from_file(args.exp)

    virat_dataset = ViratDataset(data_root=experiment_config.data.data_root,
            frames_root=experiment_config.data.frames_root,
            flow_root=experiment_config.data.flow_root,
            train_kpf_dirs=experiment_config.data.kpf_train_annotation_dirs,
            test_kpf_dirs=experiment_config.data.kpf_test_annotation_dirs,
            class_index_file=experiment_config.data.class_index,
            save_directory=experiment_config.data.save_directory)
    generate_bias_maps(virat_dataset)
    generate_bias_maps(virat_dataset, True)

if __name__ == "__main__":
    argument_parser = argparse.ArgumentParser("Command line options to detect bias")
    argument_parser.add_argument("--exp", help="experiment configuration for virat")
    argument_parser.add_argument("--save_path", 
                            help="directory where all the images are stored")
    args = argument_parser.parse_args()
    main(args)
