# -*- coding: utf-8 -*-
"""
Kwiver process to write ``track_object_set`` from RC3D into NIST specified JSON format

@author: Ameya Shringi
"""

#kwiver/sprokit import
from kwiver.sprokit.pipeline import process, datum, DatumType
from kwiver.sprokit.processes.kwiver_process import KwiverProcess
from kwiver.vital.types import DetectedObjectSet, BoundingBox, DetectedObject, DetectedObjectType
from kwiver.vital.types import Track, ObjectTrackState
import os
import numpy
import json
import math
import threading
import copy

# RC3D imports
from log_to_nist import generate_classes, generate_classes_from_json
from tdcnn.exp_config import expcfg_from_file, experiment_config
from tdcnn.config import cfg_from_file, cfg

class RC3DJsonWriter(KwiverProcess):
    """
    Write ``object_track_set`` from RC3D into NIST specified JSON format
    * Input Ports:
        * ``object_track_set`` Detected object set obtained from RC3D (Required)
        * ``timestamp`` Timestamp associated with the input from which the detected object set was obtained from (Required)
        * ``file_name`` Name of the input source (Required)

    * Output Ports:
        * None

    * Configuration:
        * ``experiment_file_name`` Experiment configuration used by RC3D (Eg. `experiment.yml`_)
        * ``model_cfg`` Model configuration used by RC3D (Eg. `td_cnn_end2end.yml`_)
        * ``stride`` Temporal stride for RC3D (default=8)
        * ``temporal_buffer`` Buffer frames when merging frames (default=8)
        * ``json_path`` Path where output json file is saved (default=sysfile.json)
        * ``confidence_threshold`` Minimum confidence threshold used by RC3D to  accept an activity detection (default=0.05)

    .. Repo links:

    .. _td_cnn_end2end.yml: https://gitlab.kitware.com/kwiver/R-C3D/blob/master/experiments/virat/td_cnn_end2end.yml
    .. _experiment.yml: https://gitlab.kitware.com/kwiver/R-C3D/tree/master/experiments/virat/experiment.yml
    """
    def __init__(self, conf):
        """
        Constructor for RC3D json writer
        :param conf: configuration parameters for RC3D json writer
        :return None
        """
        KwiverProcess.__init__(self, conf)
        # Experiment configuration
        self.add_config_trait("experiment_file_name", "experiment_file_name", \
                                "experiment.yml",
                                "experiment configuration")
        self.declare_config_using_trait('experiment_file_name')
        self.add_config_trait("model_cfg", "model_cfg", \
                                "td_cnn_end2end.yml",
                                "model configuration")
        self.declare_config_using_trait('model_cfg')
        # Temporal Buffer used for merging frames (Not used yet)
        self.add_config_trait("temporal_buffer", "temporal_buffer", \
                                "1536", "Allow for some buffer when merging frames")
        self.declare_config_using_trait('temporal_buffer')
        # Match the stride with RC3D's temporal stride
        self.add_config_trait("stride", "stride", \
                                "8", "Temporal stride for RC3D")
        self.declare_config_using_trait("stride")
        # Save path for json
        self.add_config_trait("json_path", "json_path", \
                                "sysfile.json", "Path to save the json file")
        self.declare_config_using_trait("json_path")
        # Threshold to add an activity instance to the log
        self.add_config_trait("confidence_threshold", "confidence_threshold", \
                                "0.05", "confidence threshold for the scorer")
        self.declare_config_using_trait("confidence_threshold")
        # Activity index provided by NIST
        self.add_config_trait("activity_index", "activity_index", \
                                "", "Activity index that can have a subset of activities")
        self.declare_config_using_trait("activity_index")
        # set up required flags
        required = process.PortFlags()
        required.add(self.flag_required)
        self.declare_input_port_using_trait("object_track_set", required)
        self.declare_input_port_using_trait('timestamp', required )
        self.declare_input_port_using_trait('file_name', required )

    def _configure(self):
        """
        Configure RC3D Json Writer
        """
        # look for 'experiment_file_name' key in the config
        expcfg_from_file(self.config_value('experiment_file_name'))
        if os.path.exists(self.config_value("json_path")):
            print ("Removing old system output file")
            os.remove(self.config_value("json_path"))
        if len(self.config_value("activity_index")) > 0:
            if os.path.exists(self.config_value("activity_index")):
                self.pruned_activity_index = json.load(open(self.config_value("activity_index"), "r")).keys()
            else:
                raise OSError("Invalid path {0} specified for the activity index".\
                                format(self.config_value("activity_index")))
        else:
            self.pruned_activity_index = None
        # experiment configuration
        cfg_from_file(self.config_value('model_cfg'))
        self.activity_id = 0
        if experiment_config.json:
            self.classes = generate_classes_from_json(experiment_config.class_index)
        else:
            self.classes = generate_classes(experiment_config.class_index)

        self.current_activity_frames = [-1]*len(self.classes)
        self.start_frames = [0]*len(self.classes)

        self.current_activities = []
        self.active_tracks = []
        self.video_name = None

    def _compute_temporal_intersection( self, start_frame_1, end_frame_1, start_frame_2,
                               end_frame_2 ):
        start_frame = max(start_frame_1, start_frame_2)
        end_frame = min(end_frame_1, end_frame_2)
        intersection = max(0, end_frame - start_frame)
        return intersection

    def merge_active_tracks(self, object_track, matched_track_ids):
        merged_track = Track()
        object_frame_ids = list(object_track.all_frame_ids())
        matched_frame_ids = {}
        for track_id in matched_track_ids:
            matched_frame_ids[track_id] = self.active_tracks[track_id].all_frame_ids()
        all_frame_ids = copy.deepcopy(object_frame_ids)
        for frame_ids in matched_frame_ids.values():
            all_frame_ids.extend(frame_ids)
        all_frame_ids = list(set(all_frame_ids))
        for frame_id in sorted(all_frame_ids):
            largest_class_score = -1.
            largest_matched_track_id = -1
            for matched_track_id in matched_track_ids:
                if frame_id in matched_frame_ids[matched_track_id]:
                    active_track_score = self.active_tracks[matched_track_id].find_state(frame_id).\
                                            detection().type().get_most_likely_score()
                    if active_track_score > largest_class_score:
                        largest_matched_track_id = matched_track_id
                        largest_class_score = active_track_score
            if largest_matched_track_id >= 0:
                # A match was found the in active tracks
                if frame_id in object_frame_ids:
                    # The frame is between object track and an active track
                    object_track_score = object_track.find_state(frame_id). \
                                            detection().type().get_most_likely_score()
                    if object_track_score >= largest_class_score:
                        # The frame id must come from object track
                        ots = ObjectTrackState( frame_id, frame_id,
                                                object_track.find_state( frame_id ).detection() )
                        merged_append_successful = merged_track.append(ots)
                    else:
                        ots = ObjectTrackState( frame_id, frame_id,
                                                self.active_tracks[largest_matched_track_id].find_state( frame_id ).detection() )
                        merged_append_successful = merged_track.append(ots)
                else:
                    ots = ObjectTrackState( frame_id, frame_id,
                                            self.active_tracks[largest_matched_track_id].find_state( frame_id ).detection() )
                    merged_append_successful = merged_track.append(ots)
            else:
                # The frame id must come from object track
                ots = ObjectTrackState( frame_id, frame_id,
                                        object_track.find_state( frame_id ).detection() )
                merged_append_successful = merged_track.append(ots)

        merged_tracks = []
        for track_id, track in enumerate(self.active_tracks):
            if track_id not in matched_track_ids:
                merged_tracks.append(track)
        merged_tracks.append(merged_track)
        return merged_tracks

    def merge_object_tracks(self, object_tracks):
        # Merge object tracks which overlap
        is_objs_matched = [False]*len(object_tracks)
        for obj_id, object_track in enumerate(object_tracks):
            object_do = object_track.find_state(object_track.first_frame).detection()
            object_class_name = object_do.type().get_most_likely_class()
            active_matched = []
            for active_id, active_track in enumerate(self.active_tracks):
                active_do = active_track.find_state(active_track.first_frame).detection()
                active_class_name = active_do.type().get_most_likely_class()
                if active_class_name != object_class_name:
                    continue
                temporal_intersection = \
                        self._compute_temporal_intersection(active_track.first_frame,
                                    active_track.last_frame, object_track.first_frame,
                                    object_track.last_frame)
                if temporal_intersection <= 0:
                    continue
                else:
                    active_matched.append(active_id)
            if len(active_matched) > 0:
                self.active_tracks = self.merge_active_tracks( object_track, active_matched )
                is_objs_matched[obj_id] = True

        # Add the remaining objects to active tracks
        for obj_id, is_obj_matched in enumerate(is_objs_matched):
            if not is_obj_matched:
                self.active_tracks.append(object_tracks[obj_id])

    def get_json_activities(self, finished_activities):
        segment = []
        for track_id, track in enumerate(self.active_tracks):
            start_frame = track.first_frame
            last_frame = track.last_frame
            if start_frame == last_frame:
                continue
            det = track.find_state(start_frame).detection()
            class_name = det.type().get_most_likely_class()
            scores = [t.detection().type().get_most_likely_score() for t in track]
            segment.append( { "activity": class_name,
                              "activityID": track_id,
                              "presenceConf": max(scores),
                              "alertFrame": start_frame,
                              "localization": {
                                                self.video_name: {
                                                    str(start_frame): 1,
                                                    str(last_frame): 0
                                                }
                                            }
                            } )
        for f_track_id, finished_activity in enumerate(finished_activities):
            finished_activity["activityID"] = f_track_id + len(self.active_tracks)
            segment.append(finished_activity)
        return segment

    def _create_track(self, start_frame, end_frame, class_name, score ):
        track = Track()
        for frame_number in range(start_frame, end_frame+1):
            bbox = BoundingBox(0, 0, 1, 1)
            dot = DetectedObjectType(class_name, score)
            do = DetectedObject(bbox, score, dot)
            ots = ObjectTrackState(frame_number, frame_number, do )
            append_succesful = track.append(ots)
        return track

    def get_tracks_from_file(self, results, timestamp):
        activities = results['activities']
        active_tracks = []
        passive_tracks = []
        for activity in activities:
            loc = activity['localization']
            assert len(loc) == 1, "The json file does not follow NIST specifications"
            activity_duration = list(loc.values())[0]
            start_frame = min(map(int, activity_duration.keys()))
            end_frame = max(map(int, activity_duration.keys()))
            if end_frame < (timestamp.get_frame() - int(self.config_value('temporal_buffer'))):
                passive_tracks.append(activity)
            else:
                class_name = activity['activity']
                score = activity['presenceConf']
                activity_track = self._create_track( start_frame, end_frame,
                                                     class_name, score )
                active_tracks.append(activity_track)
        return active_tracks, passive_tracks

    def _step(self):
        """
        Step function for RC3D JSON Writer
        """
        # Read object track set and timestamp
        if self.peek_at_datum_on_port('object_track_set').type() != DatumType.complete:
            object_track_set = self.grab_input_using_trait('object_track_set')
            object_tracks = object_track_set.tracks()
            ts = self.grab_input_using_trait('timestamp')
            # Only consider track that are spawned with the current timestamp
            object_tracks = [track for track in object_tracks if track.last_frame >= ts.get_frame() \
                              and track.first_frame < ts.get_frame()]
            self.video_name = self.grab_input_using_trait('file_name')
            if ts.get_frame()%int(self.config_value('stride')) == 0:
                print("Number of active tracks: {0}".format(len(self.active_tracks)))
                # If an alternative activity index is specified use it and prune the
                # tracks based on it
                if self.pruned_activity_index is not None:
                    pruned_tracks = []
                    for object_track in object_tracks:
                        object_do = object_track.find_state(object_track.first_frame).detection()
                        object_class_name = object_do.type().get_most_likely_class()
                        if object_class_name in self.pruned_activity_index:
                            pruned_tracks.append(object_track)
                    object_tracks = pruned_tracks
                json_file_path = self.config_value('json_path')
                if os.path.exists(json_file_path):
                    results = json.load(open(json_file_path, 'r'))
                    self.active_tracks, finished_activities = self.get_tracks_from_file(results, ts)
                else:
                    finished_activities = []
                if len(object_tracks) > 0:
                    self.merge_object_tracks(object_tracks)
                # Convert finished tracks in NIST specific format
                segments = self.get_json_activities(finished_activities)
                video_processed = [self.video_name]
                results = {'filesProcessed': video_processed,
                           'activities': segments}
                json.dump(results, open(json_file_path, 'w'), indent=2)

def __sprokit_register__():
    """
    Sprokit registration for the process
    """
    from kwiver.sprokit.pipeline import process_factory

    module_name = 'python:kwiver.RC3DJsonWriter'

    if process_factory.is_process_module_loaded(module_name):
        return
    process_factory.add_process('RC3DJsonWriter', 
                'Write detected_object_set from rc3d to NIST specified JSON format', 
                RC3DJsonWriter)
    process_factory.mark_process_module_as_loaded(module_name)
