# -*- coding: utf-8 -*-
"""
Created on Mon May 28 13:41:56 2018

@author: Ameya Shringi

Wrapper without roidb 
"""

from __future__ import print_function
import os
import cv2
import numpy as np
import argparse
import pprint
import cPickle 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 
from roidb_generation import generate_testing_video_db

# R-C3D (https://gitlab.kitware.com/kwiver/R-C3D)
os.environ['GLOG_minloglevel'] = '2'
import caffe
from tdcnn.config import cfg_from_file, cfg
from tdcnn.test_online import test_net_online

from diva.utils import experiment


def get_video_prefix(file_name):
    """
    Create relative path from the directory where v1 frames are stored
    Uses first 4 numerical values in the kpf filename
    :param file_name: name of the kpf file
    :return relative path from the frame's root directory
    """
    video_name = file_name.split(".")[0]
    try:
        _, _, video_number, _ = video_name.split("_", 3)
    except ValueError:
        _, _, video_number = video_name.split("_")
    video_rel_path = os.path.join(video_number[:4], video_name)
    return video_rel_path


def main(args):
    # Config paramters for the experiments as specified by the experiment yaml file
    
    if args.exp_config is not None:
        expcfg_from_file(args.exp_config)

    if experiment_config.cnn_config is not None:
        cfg_from_file(os.path.join(experiment_config.experiment_root, 
                                    experiment_config.cnn_config))

    window_length = cfg.TRAIN.LENGTH[0]
    cfg.GPU_ID = experiment_config.gpu
    print('Testing the network')
    caffe.set_device(experiment_config.gpu)
    caffe.set_mode_gpu()
    net = caffe.Net(os.path.join(experiment_config.experiment_root, 
                                 experiment_config.test.network), 
                    os.path.join(experiment_config.experiment_root,
                        experiment_config.results_path,
                        experiment_config.test.model), caffe.TEST)
    net.name = os.path.splitext(os.path.basename(experiment_config.test.model))[0]
    all_logs = [] 
    diva_config = experiment()
    diva_config.read_experiment(args.diva_config)
    frame_iterator = diva_config.get_input()
    previous_buffer = None
    frame_count = 0
    while (frame_iterator.has_next_frame()):
        if frame_count % args.stride == 0:
            current_image = frame_iterator.get_next_frame().asarray()
            image_copy = np.array(current_image, copy=True)
            current_test_log, previous_buffer = test_net_online(net, current_image,
                               frame_count, args.stride,
                               max_per_image=experiment_config.test.max_detections, 
                               vis=experiment_config.test.visualize, 
                               previous_buffer=previous_buffer,
                               use_running_frames=args.use_running_frames,
                               dataset_id = frame_iterator.get_dataset_id())
    
            if args.postprocess_logs and len(all_logs) > 1:
                current_test_log, last_log = \
                    current_test_log.combine_logs(all_logs[len(all_logs)-1], 
                            args.temporal_threshold)
                if len(last_log.activities) > 0:
                    all_logs[len(all_logs)-1] = last_log
                else:
                    all_logs.pop()
            all_logs.append(current_test_log)
        else:
            frame_iterator.get_next_frame()
        frame_count += 1

    print("Creating json output")
    sys_out = sys_to_res(all_logs, generate_classes(os.path.join(
                                        experiment_config.data_root, 
                                        experiment_config.class_index)), 
                                        thresh = 0.05, act_id=0, fps=25.0)

    vid_res = {'filesProcessed': [], 'activities': [] }
    for vid, vinfo in sys_out.iteritems():
        vid_res['filesProcessed'].extend( [vid])
        vid_res['activities'].extend(vinfo)
    with open( args.json_path, 'w') as outfile:
        json.dump(vid_res, outfile, indent=4)

if __name__ == "__main__":
    parser = argparse.ArgumentParser("Training, testing and evaluation wrapper for R-C3D")
    parser.add_argument("--exp_config", help="Experiment configuration file")
    parser.add_argument("--diva_config", help="Configuration file for diva_utils")
    parser.add_argument("--stride", type=int, help="stride for the model")
    parser.add_argument("--json_path", help="Json path for results")
    parser.add_argument("--use_running_frames", action="store_true", \
                        help="test with running frames")
    parser.add_argument("--postprocess_logs", action="store_true", \
                        help="flag to merge logs")
    parser.add_argument("--temporal_threshold", type=float, default=0.5, \
                        help="temporal threshold to merge the logs")
    parser.set_defaults(use_running_frames=False)
    parser.set_defaults(postprocess_logs=False)
    args = parser.parse_args()
    main(args)
