"""
Data reader for ACT tublet for VIRAT V1 dataset
Author: Ameya Shringi (ameya.shringi@kitware.com)
"""
#import _init_paths

import os
import sys
import cPickle as pkl
import cv2
import json
import argparse
import diva_python_utils
from PIL import Image
import numpy as np

class ViratDataset(object):
    def __init__(self, frame_roots, flow_roots, train_dirs, test_dirs, 
                    class_index_file, save_directory, kpf_mode, json_mode, 
                    save_prefix='virat', is_test=False):
        """
        Constructor for the data reader
        :param frame_roots: list of directories where rgb frames are stored 
        :param flow_roots:  list of directories where  flow frames are stored 
        :param train_dirs: list of directories where annotations for training and validation are stored
        :param test_dirs: list of directories where annotations for testing are stored
        :param class_index_file: path for text file with class labels for virat data
        :param save_directory: directory where groundtruth pickle files would be stored
        :param kpf_mode: flag to use kpf annotations
        :param json_mode: flag to use json annotations
        :param save_prefix: prefix associated with the pickle files generated (default: Virat)
        :return None
        """
        if (kpf_mode and json_mode):
            raise Exception("kpf and json modes are mutually exclusive choose 1 of them")
        if (not kpf_mode and not json_mode):
            raise Exception("select kpf or json mode")
        self.frame_roots = frame_roots
        self.flow_roots = flow_roots
        self.save_directory = save_directory
        if not os.path.exists(save_directory):
            os.system('mkdir -p ' + save_directory)
        self._save_prefix = save_prefix
        self._kpf_mode = kpf_mode
        self._json_mode = json_mode
        self._labels = self._generate_labels(class_index_file)
        if not is_test:
            self._train_video_list, self._train_num_frames, self._train_video_resolution = \
                                                self.video_information(train_dirs,
                                                        save_directory, save_prefix)
            self._train_gttubes = self.extract_tubes(train_dirs, save_directory,
                                                     save_prefix)
        else:
            self._test_video_list, self._test_num_frames, self._test_video_resolution = \
                                                self.video_information(test_dirs,
                                                        save_directory, save_prefix,\
                                                        is_test=is_test)

    def _generate_labels(self, class_index_file):
        """
        Generate dictionary of labels from text file provided by VIRAT (index 0 is 
        background class)
        :param class_index_file: path to the file with class labels 
                                 (eval 1.A uses text file and eval 1.B uses json file)
        :return dict with class name as key and class index as value
        """
        class_index  = 0
        classes = {'Background': 0}
        if os.path.exists(class_index_file):
            if self._json_mode:
                with open(class_index_file, 'r') as f:
                    class_dict = json.load(f)
                    for class_name in class_dict.keys():
                        class_index += 1
                        classes[class_name] = class_index
                return classes
            elif self._kpf_mode:
                with open(class_index_file, 'r') as f:
                    for line in f:
                        class_index += 1
                        line = line.strip()
                        classes[line] = class_index
                return classes
            else:
                raise NotImplementedError("This implementation only supports kpf and json annotations")
        else:
            raise OSError("Class index file does not exists at the specified path")            

    
    def _search_path(self, frame_roots, video_name):
        video_path = None
        for frame_root in frame_roots:
            if video_name in frame_root:
                return frame_root
            else:
                file_names = os.listdir(frame_root)
                for file_name in file_names:
                    file_path = os.path.join(frame_root, file_name)
                    if os.path.isdir(file_path):
                        video_path = self._search_path([file_path], video_name)
                    if video_path is not None:
                        break
                return video_path

    def _get_video_prefix(self, frame_roots, file_name):
        """
        Create relative path from the directory where v1 frames are stored
        Uses first 4 numerical values in the kpf filename
        :param frame_root: root directory for frames
        :param file_name: name of the kpf file
        :return relative path from the frame's root directory
        """
        frame_path = self._search_path( frame_roots, file_name )
        if frame_path is not None:
            return frame_path
        else:
            raise OSError("video not available in the frame roots")
    
    def _find_absolute_path(self, root_directories, relative_video_path): 
        """
        Given a relative path, find the absolute path with respect to the frame roots
        :param relative path: relative path of the video
        return absolute video path if there is one or None is not video path in not 
                found in any video root
        """
        for root_directory in root_directories:
            probable_path = os.path.join(root_directory, relative_video_path)
            if os.path.exists(probable_path):
                return probable_path
        return None 

    def video_information(self, annotation_dirs, save_directory, save_prefix, is_test=False):
        """
        Create list of relative path from frame root and compute resolution and 
        number of frames for training and testing
        Cache the information in a pickle file to avoid recomputation
        :param annotation_dirs: list of directories where annotation are stored
        :param save_directory: path where pickle file is stored
        :param save_prefix: prefix associated with pickle file
        :param is_test: flag to determine is the pickle file is for test (default: False)
        :return Dictionary with list of videos, number of frames and frame resolution
        """
        if is_test:
            save_path = os.path.join(save_directory, save_prefix + '_test_video_information.pkl')
        else:
            save_path = os.path.join(save_directory, save_prefix + '_video_information.pkl')
        if os.path.exists(save_path):
            print "Using existing video information from " + os.path.basename(save_path)
            video_info = pkl.load(open(save_path, 'r'))
            return video_info['video_list'], video_info['num_frames'], \
                    video_info['resolution']
        all_activity_files = []
        video_list = []
        num_frames_dict = {}
        video_resolution_dict = {}
        if self._kpf_mode:
            file_pattern = "activities"
        elif self._json_mode:
            file_pattern = "file-index"
        relevant_files = []
        for annotation_dir in annotation_dirs:
            files = [os.path.join(annotation_dir, f) for f in os.listdir(annotation_dir)\
                        if file_pattern in f]
            relevant_files.extend(files)
        for relevant_file in relevant_files:
            if self._json_mode:
                with open(relevant_file, 'r') as f:
                    file_index = json.load(f)
                    video_names = file_index.keys()
                    video_path = []
                    for video_name in video_names:
                        video_path.append(self._get_video_prefix(self.frame_roots, 
                                                os.path.splitext(video_name)[0]) )
            else:
                video_path = self._get_video_prefix(self.frame_roots,
                                            os.path.basename(relevant_file))
            if isinstance(video_path, list):
                video_list.extend(video_path)
                for path in video_path:
                    if path is not None:
                        all_frames = os.listdir(path)
                        num_frames_dict[path] = len(all_frames)
                        height, width, _ = cv2.imread(os.path.join(path, \
                                                    all_frames[0])).shape
                        video_resolution_dict[path] = (height, width)
                    else:
                        raise OSError("Invalid video path")
            else:
                video_list.append(video_path)
                if os.path.exists(video_path):
                    all_frames = os.listdir(video_path)
                    num_frames_dict[video_rel_path] = len(all_frames)
                    height, width, _ = cv2.imread(os.path.join(video_path, all_frames[0])).shape
                    video_resolution_dict[video_rel_path] = (height, width)
                else:
                    raise OSError("Invalid video path")

        pickle_dict = {}
        pickle_dict['video_list'] = video_list
        pickle_dict['num_frames'] = num_frames_dict
        pickle_dict['resolution'] = video_resolution_dict
        pkl.dump(pickle_dict, open(save_path, 'w'))
        return video_list, num_frames_dict, video_resolution_dict

    def _union_bbox(self, bounding_boxes):
        """
        Compute union of list of bounding boxes
        :param bounding_boxes: list of bounding boxes
        :return: List with 1 bounding box
        """
        largest_box = bounding_boxes[0]
        for bounding_box in bounding_boxes:
            if largest_box[0] > bounding_box[0]:
                largest_box[0] = bounding_box[0]
            if largest_box[1] > bounding_box[1]:
                largest_box[1] = bounding_box[1]
            if largest_box[2] < bounding_box[2]:
                largest_box[2] = bounding_box[2]
            if largest_box[3] < bounding_box[3]:
                largest_box[3] = bounding_box[3]
        return [largest_box]

    def _merge_activity_and_geometry(self, activity_dict, geometry_dict):
        """
        Merge activity and geometry kpf files to extract activity specific object detection
        :param activity_dict: dictionary created from activity kpf annotation
        :param geoemtry_dict: dictionary created from geoemtry kpf annotation
        :return dictionary with class index as key and list of numpy array storing
                    ground truth tubes
        """
        all_tubes = {}
        for _, value in activity_dict.iteritems():
            activity_index, all_actors = value
            frame_dict = {}
            for actor_id, frame_span in all_actors.iteritems():
                start_frame, end_frame = map(int, frame_span)
                if actor_id in geometry_dict.keys():
                    for frame in range(start_frame, end_frame):
                        if frame in geometry_dict[actor_id].keys():
                            if frame in frame_dict.keys():
                                frame_dict[frame].extend([geometry_dict[actor_id][frame]])
                            else:
                                frame_dict[frame] = [geometry_dict[actor_id][frame]]
                        else:
                            print "No bbox in frame " + str(frame) + \
                                    " for actor: " + str(actor_id)
                else:
                    print "No bbox associated with the actor id:" + str(actor_id)
            activity_array = []
            for frame_number, bbox in frame_dict.iteritems():
                if len(bbox) > 1:
                    bbox = self._union_bbox(bbox)
                current_activity = [frame_number]
                current_activity.extend(bbox[0])
                activity_array.append(current_activity)
            activity_array = np.array(activity_array, dtype=np.float32)
            if activity_index in all_tubes.keys():
                all_tubes[activity_index].append(activity_array)
            else:
                all_tubes[activity_index] = [activity_array]
        return all_tubes

    def _get_kpf_bbox_coordinates(self, kpf_bbox):
        return [kpf_bbox.get_x1(), kpf_bbox.get_y1(), kpf_bbox.get_x2(), kpf_bbox.get_y2()] 
   
    def _get_json_bbox_coordinates(self, json_bbox):
        x1 = json_bbox["x"]
        y1 = json_bbox["y"]
        return [x1, y1, x1 + json_bbox["w"], y1 + json_bbox["h"]]

    def _get_actor_ids_and_spans(self, all_actor_span_map):
        actors = {}
        for actor_span_map in all_actor_span_map:
            actor_id, span_vector = actor_span_map
            for span in span_vector:
                start_frame = span[0]
                end_frame = span[1]
            actors[actor_id] = (start_frame, end_frame)
        return actors

    def _generate_kpf_tubes(self, annotation_dirs, is_test):
        all_activity_files = []
        gt_dict = {}
        kpf_activity = diva_python_utils.activity()
        kpf_geometry = diva_python_utils.geometry()
        for kpf_dir in annotation_dirs:
            all_activity_files.extend([os.path.join(kpf_dir, filename) \
                                    for filename in os.listdir(kpf_dir) \
                                    if 'activities' in filename])
        for count, activity_file in enumerate(all_activity_files):
            video_rel_path = self._get_video_prefix(annotation_dirs, 
                                            os.path.basename(activity_file))
            activity_dict = {}
            geometry_dict = {}
            with open(activity_file, 'r') as act_file:
                for line in act_file:
                    if line.startswith("- { act:"):
                        kpf_activity.clear()
                        kpf_activity.from_string(line)
                        label = kpf_activity.get_max_activity_name()
                        if label in self.labels:
                            all_actors = kpf_activity.get_actor_frame_id_span()
                            all_actors_dict = self._get_actor_ids_and_spans(all_actors)
                            activity_id = kpf_activity.get_activity_id()
                            activity_dict[activity_id] = (self.labels[label],\
                                                            all_actors_dict)
            geometry_file = os.path.join(os.path.dirname(activity_file),
                                         os.path.basename(activity_file).split(".")[0]+\
                                        '.geom.yml')
            with open(geometry_file, 'r') as geom_file:
                for line in geom_file:
                    if line.startswith("- { geom:"):
                        kpf_geometry.clear()
                        kpf_geometry.from_string(line)
                        track_id = kpf_geometry.get_track_id()
                        frame_id = kpf_geometry.get_frame_id()
                        bounding_box = self._get_kpf_bbox_coordinates(
                                            kpf_geometry.get_bounding_box_pixels())
                        if track_id not in geometry_dict.keys():
                            geometry_dict[track_id] = {frame_id: bounding_box}
                        else:
                            geometry_dict[track_id][frame_id] = bounding_box
            gtt_tubes = self._merge_activity_and_geometry(activity_dict, geometry_dict)
            gt_dict[video_rel_path] = gtt_tubes
            print ("Completed: " + str(count+1) + "/" + str(len(all_activity_files)))
        return gt_dict

    def _generate_json_tubes(self, annotation_dirs, is_test):
        gt_dict = {}
        for annotation_dir in annotation_dirs:
            with open(os.path.join(annotation_dir, "activities.json"), 'r') as f:
                gt_json = json.load(f)
            all_activities = gt_json["activities"]
            for activities in all_activities:
                activity_label = activities["activity"]
                activity_id = int(activities["activityID"])
                video_name = activities["localization"].keys()[0]
                activity_span = activities["localization"].values()[0]
                start_frame = int(activity_span.keys()[activity_span.values().index(1)])
                end_frame = int(activity_span.keys()[activity_span.values().index(0)])
                video_rel_path = self._get_video_prefix(annotation_dirs, video_name)
                all_objects = activities["objects"]
                frames_dict = {}
                for obj in all_objects:
                    bboxes = obj["localization"].values()[0]
                    for frame_number, bbox_annotation in  bboxes.iteritems():
                        frame_number = int(frame_number)
                        if "boundingBox" in bbox_annotation.keys():
                            bbox_coordinates = self._get_json_bbox_coordinates(
                                                bbox_annotation["boundingBox"])
                        else:
                            continue
                        if frame_number in frames_dict.keys():
                            frames_dict[frame_number].append(bbox_coordinates)
                        else:
                            frames_dict[frame_number] = [bbox_coordinates]
                frame_array = []
                for frame_number, bboxes in frames_dict.iteritems():
                    current_frame_array = [frame_number]
                    current_frame_array.extend(self._union_bbox(bboxes)[0])
                    frame_array.append(current_frame_array)
                frame_array = np.array(frame_array, dtype=np.float32)
                frame_array = frame_array[np.argsort(frame_array[:, 0])]
                activity_index = self.labels[activity_label]
                if video_rel_path in gt_dict.keys():
                    if activity_index in gt_dict[video_rel_path]:
                        gt_dict[video_rel_path][activity_index].append(frame_array)
                    else:
                        gt_dict[video_rel_path][activity_index] = [frame_array]
                else:
                    gt_dict[video_rel_path] = {activity_index: [frame_array]}
        return gt_dict


    def extract_tubes(self, annotation_dirs, save_directory, save_prefix, is_test=False):
        """
        Create ground truth spatio-temporal tubes from annotations
        :param annotation_dirs: list of annotation directory
        :param save_directory: path where generated pickle file is stored
        :param save_prefix: prefix associated with the pickle file
        :param is_test: flag to determine if the dataset used is test set
        :return dictionary with video name as key and groundtruth tubes as value
        """
        if is_test:
            save_path = os.path.join(save_directory, save_prefix + '_test_video_tubes.pkl')
        else:
            save_path = os.path.join(save_directory, save_prefix + '_video_tubes.pkl')
        if os.path.exists(save_path):
            print "Loading ground truth tubes from " + os.path.basename(save_path)
            return pkl.load(open(save_path, 'r'))

        if self._kpf_mode:
            gt_dict = self._generate_kpf_tubes(annotation_dirs, is_test) 
        elif self._json_mode:
            gt_dict = self._generate_json_tubes(annotation_dirs, is_test)
        else:
            raise NotImplementedError("Virat data reader supports only kpf and json files")

        pkl.dump(gt_dict, open(save_path, 'w'))
        return gt_dict

    @property
    def nlabels(self):
        # skip background
        return len(self._labels)

    @property
    def save_prefix(self):
        return self._save_prefix

    @property
    def labels(self):
        return self._labels

    @property
    def train_video_list(self):
        return self._train_video_list

    @property
    def test_video_list(self):
        return self._test_video_list

    def train_num_frames(self, video_name):
        try:
            return self._train_num_frames[video_name]
        except KeyError:
            print "Incorrect video name"
            return None

    def test_num_frames(self, video_name):
        try:
            return self._test_num_frames[video_name]
        except KeyError:
            print "Incorrect video name"
            return None

    def train_video_resolution(self, video_name):
        try:
            return self._train_video_resolution[video_name]
        except KeyError:
            print "Incorrect video name"

    def test_video_resolution(self, video_name):
        try:
            return self._test_video_resolution[video_name]
        except KeyError:
            print "Incorrect video name"
    
    def train_gttubes(self, video_name):
        try:
            return self._train_gttubes[video_name]
        except KeyError:
            print "No tubes associated with: " + str(video_name)
            return None

    def imfile(self, video_name, image_name):
        return os.path.join(video_name, image_name)

    def flowfile(self, video_name, image_name):
        rel_path = None
        for frame_root in self.frame_roots:
            video_index  = video_name.find(frame_root)
            if video_index > -1:
                # need to remove the hash from the beginning in the relative path
                rel_path = video_name[video_index + len(frame_root) + 1:]
                break
        if rel_path is None:
            raise OSError("Video is not available in any frame roots")
        for flow_root in self.flow_roots:
            probable_path = os.path.join(flow_root, rel_path, image_name)
            if os.path.exists(probable_path):
                return probable_path
        raise OSError("Video {0} is not present in flow roots".format(os.path.basename(rel_path)))

def main(args):
    if not args.kpf_mode and not args.json_mode:
        raise Exception("Annotation format not specified")
    virat_dataset = ViratDataset(frame_roots=args.frame_roots,
                                 flow_roots=args.flow_roots,
                                 train_dirs=args.train_dirs,
                                 test_dirs=args.test_dirs,
                                 class_index_file=args.class_index_file, 
                                 save_directory=args.save_directory,
                                 kpf_mode=args.kpf_mode, json_mode=args.json_mode)
    

if __name__ == "__main__":
    parser = argparse.ArgumentParser("Data reader for virat")
    parser.add_argument("--kpf_mode", action="store_true", help="Use kpf annotations")
    parser.add_argument("--json_mode", action="store_true", help="Use json annotations")
    parser.add_argument("--frame_roots", nargs="+", help= "path of RGB images")
    parser.add_argument("--flow_roots", nargs="+", help= "path of flow images")
    parser.add_argument("--train_dirs", nargs="+", help="training annotations")
    parser.add_argument("--test_dirs", nargs="+", help="testing annotations")
    parser.add_argument("--class_index_file", help="index file used to determine class label")
    parser.add_argument("--save_directory", 
            help="absolute path where pickled representation are saved")
    parser.set_defaults(kpf_mode=False)
    parser.set_defaults(json_mode=False)
    args = parser.parse_args()
    main(args)

