# -*- coding: utf-8 -*-
"""
@author: Ameya Shringi (ameya.shringi@kitware.com)

Wrapper for VIRAT R-C3D: generating pickle files, training, testing, evaluation
Assumed that all frames, truth train/test json, activity list files are already created
"""
from __future__ import print_function
import os
import numpy as np
import argparse
import pickle as pkl
import sys
import datetime
import json

# DIVA support scripts
from tdcnn.exp_config import expcfg_from_file, experiment_config
from roidb_generation import generate_roidb_pkl
from log_to_nist import convert_log_to_json
import pprint
# R-C3D (https://gitlab.kitware.com/kwiver/R-C3D)
import caffe
from tdcnn.train import train_net
from tdcnn.config import cfg_from_file, cfg
from tdcnn.test import test_net

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

    if args.model_cfg is not None:
        cfg_from_file(args.model_cfg)

    window_length = cfg.TRAIN.LENGTH[0]
    cfg.GPU_ID = experiment_config.gpu

    if not os.path.exists(os.path.join(experiment_config.experiment_root,
                            experiment_config.results_path)):
      os.makedirs(os.path.join(experiment_config.experiment_root,
                        experiment_config.results_path))



    #Step 4: Train R-C3D using output pkl file from step 2

    # if the training weights does not exist
    weights_path = None
    if not experiment_config.train.weights:
      print('No pretrained model provided, skipping training')
      args.skip_train = True
    else:
      weights_path = os.path.join(experiment_config.model_root,
                                    experiment_config.train.weights)
      if not os.path.exists(weights_path):
          print('Could not find the provided weights file... skipping training')
          args.skip_train = True

    if not args.skip_train:
        if not os.path.exists(os.path.dirname(experiment_config.train.pickle_output)):
          os.makedirs(os.path.dirname(experiment_config.train.pickle_output))
        if experiment_config.kpf:
          annotation_dirs = experiment_config.train.kpf_annotation_dirs
        elif experiment_config.json:
          annotation_dirs = experiment_config.train.json_annotations
        else:
          annotation_dirs = experiment_config.train.density_annotation_dirs

        if experiment_config.train.pickle_output:
          generate_roidb_pkl(test_mode=False,
                             kpf_mode=experiment_config.kpf,
                             json_mode=experiment_config.json,
                             window_length=window_length,
                             window_divisor=4,
                             min_length=3,
                             sampling_frequency=8,
                             use_flipped=experiment_config.use_flipped,
                             frame_roots=experiment_config.frame_roots,
                             annotation_dirs=annotation_dirs,
                             class_index=experiment_config.class_index,
                             pickle_file=experiment_config.train.pickle_output)
        train_roidb = pkl.load(open(experiment_config.train.pickle_output, 'rb'))
        print("Training the network")
        if weights_path is not None and os.path.exists(weights_path):
            pprint.pprint(experiment_config.train)
            np.random.seed(experiment_config.train.seed)
            caffe.set_random_seed(experiment_config.train.seed)
            caffe.set_device(experiment_config.gpu)
            caffe.set_mode_gpu()
            if not os.path.exists(experiment_config.model_root):
                os.makedirs(experiment_config.model_root)
            train_net(experiment_config.train.solver,
                      train_roidb, None,
                      pretrained_model=weights_path,
                      max_iters=experiment_config.train.iterations,
                      restore_iter=experiment_config.train.restore_iteration)
        else:
            print("Using an older model")

    #Step 5: Test R-C3D using output from Step 3
    if not args.skip_test:
        # use all truth for now regardless of whether it is the only activity on a given set of frames
        if experiment_config.test.pickle_output:
          if not os.path.exists(os.path.dirname(experiment_config.test.pickle_output)):
            os.makedirs(os.path.dirname(experiment_config.test.pickle_output))
          generate_roidb_pkl(test_mode=True,
                             kpf_mode=experiment_config.kpf,
                             json_mode=experiment_config.json,
                             window_length=window_length,
                             window_divisor=4,
                             min_length=3,
                             sampling_frequency=8,
                             use_flipped=experiment_config.use_flipped,
                             frame_roots=experiment_config.frame_roots,
                             class_index=experiment_config.class_index,
                             pickle_file=experiment_config.test.pickle_output,
                             file_index=experiment_config.test.file_index,
                             chunk_file=experiment_config.test.chunk_file,
                             chunk_id=args.chunk_id)
        test_roidb = pkl.load(open(experiment_config.test.pickle_output))
        print('Testing the network')
        caffe.set_device(experiment_config.gpu)
        caffe.set_mode_gpu()
        net = caffe.Net( 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]
        if not cfg.TEST.HAS_RPN:
            test_roidb.set_proposal_method(cfg.TEST.PROPOSAL_METHOD)
        test_logs = test_net(net, test_roidb, 
                        max_per_image=experiment_config.test.max_detections, 
                        vis=experiment_config.test.visualize)

        if args.chunk_id == -1:
            output_prefix = ""
        else:
            output_prefix = "Chunk" + str(args.chunk_id) + "_"

        if args.dump_individual:
            json_path = os.path.join(experiment_config.cache_root,
                                     output_prefix + 'sysfile_')
        else:
            json_path = os.path.join(experiment_config.cache_root,
                                     output_prefix + 'sysfile.json')
        convert_log_to_json(test_logs, experiment_config.class_index,
                            experiment_config.test.file_index,
                            json_path, use_json_class_index=True,
                            dump_individual=args.dump_individual,
                            test_activity_index=args.test_activity_index)



if __name__ == "__main__":
    parser = argparse.ArgumentParser("Training, testing and evaluation wrapper for R-C3D")
    parser.add_argument("--exp", help="Experiment configuration file")
    parser.add_argument("--model-cfg", help="Model configuration file")
    parser.add_argument("--skip-train", help="Flag to skip training",
                        action="store_true", dest="skip_train")
    parser.add_argument("--skip-test", help="Flag to skip testing", \
            action="store_true", dest="skip_test")
    parser.add_argument("--use-json-class-index", help="use json class index",
            action="store_true", dest="use_json_class_index")
    parser.add_argument("--dump-individual",
            help="Create json files for every video",
            action="store_true", dest="dump_individual")
    parser.add_argument("--test-activity-index", help="Activity index used for testing")
    parser.set_defaults(dump_individual=False)
    parser.add_argument("--chunk_id", type=int, default=-1,
                        help="Chunk id used for testing (default=-1 which signifies no chunks)")
    parser.set_defaults(skip_train=False)
    parser.set_defaults(skip_test=False)
    parser.set_defaults(use_json_class_index=False)
    args = parser.parse_args()
    main(args)
