"""
@author: Ameya Shringi, Linus Sherrill

Kwiver process that encapsulates forward pass of RC3D 
"""

# kwiver/sprokit imports
from kwiver.sprokit.pipeline import process
from kwiver.sprokit.processes.kwiver_process import KwiverProcess
from kwiver.vital.util.VitalPIL import from_pil, get_pil_image
from kwiver.vital.types import DetectedObjectSet, BoundingBox, DetectedObject, DetectedObjectType
from kwiver.vital.types import Track, ObjectTrackState, ObjectTrackSet
from kwiver.sprokit.pipeline import datum, DatumType

import os
import cv2
import numpy as np
import argparse
import pprint
import pickle as pkl
import sys
import datetime
import json

# DIVA support scripts
from tdcnn.exp_config import expcfg_from_file, experiment_config
from log_to_nist import sys_to_res, generate_classes, generate_classes_from_json
# R-C3D (https://gitlab.kitware.com/kwiver/R-C3D)
import caffe
from tdcnn.config import cfg_from_file, cfg
from tdcnn.test_online import test_net_online

import time

# Now for the process
class RC3DDetector(KwiverProcess):
    """
    Forward pass for RC3D

    * Input Ports:
        * ``image`` RGB image (Required)
        * ``timestamp`` Timestamp for the image (Required)
        * ``file_name`` Name of the input source (Optional)

    * Output Ports:
        * ``track_object_set`` Temporal detections obtained from forward pass

    * 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)
        * ``gpu`` Gpu index on which RC3D is executed (default=8)

    .. 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):
        """
        RC3DDetector constructor
        :param conf: Configuration file for the process
        :return: None
        """
        KwiverProcess.__init__(self, conf)
        self.add_config_trait("experiment_file_name", "experiment_file_name",
                            '.', 'experiment configuration for RC3D')
        self.declare_config_using_trait('experiment_file_name')
        self.add_config_trait("model_cfg", "model_cfg",
                            '.', 'model configuration for RC3D')
        self.declare_config_using_trait('model_cfg')
        self.add_config_trait("stride", "stride",
                            '8', 'Temporal Stride for RC3D')
        self.declare_config_using_trait('stride')
        self.add_config_trait("gpu", "gpu", "0", "gpu for RC3D")
        self.declare_config_using_trait("gpu")

        # set up required flags
        required = process.PortFlags()
        required.add(self.flag_required)
        #  declare our ports ( port-name, flags)
        self.declare_input_port_using_trait('image', required)
        self.declare_input_port_using_trait('timestamp', required )
        self.declare_input_port_using_trait('file_name', required )
        self.declare_output_port_using_trait('object_track_set', process.PortFlags() )
        self.video_name = None

    # ---------------------------------------------
    def _configure(self):
        """
        Configure RC3D Detector

        :return: None
        """
        # look for 'experiment_file_name' key in the config
        expcfg_from_file(self.config_value('experiment_file_name'))
        cfg_from_file(self.config_value('model_cfg'))
        # merge experiment configuration and network configuration
        if experiment_config.json:
            self.classes = generate_classes_from_json(experiment_config.class_index)
        else:
            self.classes = generate_classes(experiment_config.class_index)
        window_length = cfg.TRAIN.LENGTH[0]
        self.gpu_id = int(self.config_value('gpu'))
        cfg.GPU_ID = self.gpu_id
        experiment_config.gpu = self.gpu_id
        print('Testing the network')
        # Set device and load the network
        caffe.set_mode_gpu()
        caffe.set_device(self.gpu_id)
        self.net = caffe.Net(os.path.join(experiment_config.model_root,
                                     experiment_config.test.network),
                            os.path.join(experiment_config.experiment_root,
                                     experiment_config.results_path,
                                     experiment_config.test.model), caffe.TEST)
        self.net.name = os.path.splitext(os.path.basename(experiment_config.test.model))[0]
        # Initialize buffer for RC3D
        self.previous_buffer = None

    def compute_absolute_frame(self, frame, buffer_start):
        """
        Compute frame number associated with an image based on relative frame
        number associated with the current buffer and the starting point of the
        buffer

        :param frame: Relative frame number
        :param buffer_start: Starting point of the buffer

        :return: Absolute frame number
        """
        return buffer_start + frame * int(self.config_value("stride"))

    def get_framewise_scores(self, bboxes, ts, buffer_start ):
        framewise_scores = {}
        for bbox in bboxes:
            start_frame, end_frame, conf = bbox
            # Ensure that the predictions are between 0 and current timestamp
            start_frame = int(max(0, self.compute_absolute_frame(start_frame, buffer_start)))
            end_frame = int(min(ts.get_frame(),
                            self.compute_absolute_frame(end_frame, buffer_start)))
            if start_frame == end_frame:
                continue
            elif start_frame > end_frame:
                continue
            for frame_id in range(start_frame, end_frame+1):
                if frame_id not in framewise_scores.keys():
                    framewise_scores[frame_id] = conf
                else:
                    framewise_scores[frame_id] = max(conf, framewise_scores[frame_id])

        return framewise_scores

    # ----------------------------------------------
    def _step(self):
        """
        Step function for the process
        """
        if self.peek_at_datum_on_port('image').type() != DatumType.complete:
            # grab image container from port using traits
            in_img_c = self.grab_input_using_trait('image')
            ts = self.grab_input_using_trait('timestamp')
            video_name = self.grab_input_using_trait('file_name')
            # Set device configuration for the thread
            caffe.set_mode_gpu()
            caffe.set_device(self.gpu_id)
            # Get numpy array from the image container
            image = in_img_c.asarray()
            # Swap channels as we need bgr images
            image = image[..., ::-1]

            # Strided execution (temporal stride of 8)
            stride = int(self.config_value('stride'))
            if ts.get_frame()%stride == 0:
                print("Detection at frame: {0}".format(ts.get_frame()))
                logs, self.previous_buffer = test_net_online(self.net,
                                    image, ts.get_frame(),
                                    int(self.config_value('stride')),
                                    max_per_image=experiment_config.test.max_detections,
                                    vis=experiment_config.test.visualize,
                                    previous_buffer=self.previous_buffer,
                                    use_running_frames=True,
                                    dataset_id="pipeline-streamint-input")
                buffer_start = ts.get_frame() - \
                                int(self.config_value('stride')) * cfg.TEST.LENGTH[0]
                if buffer_start <= 0:
                    buffer_start = 0

                # Add temporal annotation to detected object set
                classes = []
                scores = []
                tracks = []
                for activity_id, bboxes in logs.activities.items():
                    class_name = self.classes[activity_id]
                    # Find the largest score for every class
                    framewise_scores = self.get_framewise_scores( bboxes, ts, buffer_start )
                    obj_track = Track()
                    all_frame_ids = sorted(list(framewise_scores.keys()))
                    for frame_id in all_frame_ids:
                        frame_score = framewise_scores[frame_id]
                        # Redundant field in temporal localization
                        bbox = BoundingBox(0, 0, 1, 1)
                        dot = DetectedObjectType( class_name, frame_score )
                        do = DetectedObject( bbox, frame_score, dot )
                        obj_track_state = ObjectTrackState( frame_id, frame_id, do )
                        successfully_appended = obj_track.append(obj_track_state)
                    if len(obj_track) > 0:
                        tracks.append( obj_track )
                track_set = ObjectTrackSet( tracks )
            else:
                track_set = ObjectTrackSet()
            self.push_to_port_using_trait('object_track_set', track_set)


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

    module_name = 'python:kwiver.RC3DDetector'

    if process_factory.is_process_module_loaded(module_name):
        return

    process_factory.add_process('RC3DDetector', 'Apply R-C3D detector to an image',
                                RC3DDetector)

    process_factory.mark_process_module_as_loaded(module_name)
