from kwiver.vital.config import config
from kwiver.vital.modules import load_known_modules
from kwiver.sprokit.pipeline import process_factory, scheduler_factory
from kwiver.sprokit.pipeline import pipeline

import argparse
import json
import os

from multiprocessing import Pool

def create_diva_experiment_config(exp_file):
    exp_config = config.empty_config()
    exp_config.set_value("experiment_file_name", exp_file)
    return exp_config

def create_rc3d_config(exp_file, model_config, gpu, stride):
    rc3d_config = config.empty_config()
    rc3d_config.set_value("stride", str(stride))
    rc3d_config.set_value("experiment_file_name", exp_file)
    rc3d_config.set_value("model_cfg", model_config)
    rc3d_config.set_value("gpu", str(gpu))
    return rc3d_config

def create_json_writer_config(exp_file, model_config, activity_index,
                              json_path):
    json_writer_config = config.empty_config()
    json_writer_config.set_value("experiment_file_name", exp_file)
    json_writer_config.set_value("model_cfg", model_config)
    json_writer_config.set_value("json_path", json_path)
    json_writer_config.set_value("activity_index", activity_index)
    return json_writer_config

def create_pipeline(exp_file, rc3d_experiment_file, model_config, stride,
                    gpu, activity_index, json_path):
    load_known_modules()
    exp_config = create_diva_experiment_config(exp_file)
    exp_process = process_factory.create_process("diva_experiment",
                                                 "exp",
                                                  exp_config)
    rc3d_config = create_rc3d_config(rc3d_experiment_file, model_config, gpu,
                                     stride)
    rc3d_process = process_factory.create_process("RC3DDetector",
                                                  "rc3d",
                                                  rc3d_config)
    json_writer_config = create_json_writer_config(rc3d_experiment_file,
                                                   model_config,
                                                   activity_index, json_path)
    json_writer_process = process_factory.create_process("RC3DJsonWriter",
                                                         "json_writer",
                                                         json_writer_config)
    rc3d_pipeline = pipeline.Pipeline()
    rc3d_pipeline.add_process(exp_process)
    rc3d_pipeline.add_process(rc3d_process)
    rc3d_pipeline.add_process(json_writer_process)
    rc3d_pipeline.connect( "exp", "image",
                           "rc3d", "image" )
    rc3d_pipeline.connect( "exp", "timestamp",
                           "rc3d", "timestamp" )
    rc3d_pipeline.connect( "exp", "file_name",
                           "rc3d", "file_name" )
    rc3d_pipeline.connect( "rc3d", "object_track_set",
                           "json_writer", "object_track_set" )
    rc3d_pipeline.connect( "exp", "timestamp",
                           "json_writer", "timestamp" )
    rc3d_pipeline.connect( "exp", "file_name",
                           "json_writer", "file_name" )
    rc3d_pipeline.setup_pipeline()
    pipeline_status = False
    if rc3d_pipeline.setup_successful():
       scheduler = scheduler_factory.create_scheduler( "pythread_per_process",
                                                       rc3d_pipeline )
       scheduler.start()
       scheduler.wait()
       scheduler.shutdown()
       print("Experiment associated with {0} complete".\
               format(os.path.basename(exp_file)))
       pipeline_status = True
    else:
       print("Failed to create a pipeline for {0}".\
               format(os.path.basename(exp_file)))
    del rc3d_pipeline
    return pipeline_status

def pipeline_parameter_generator(experiment_root, json_root, chunk_files,
                                 chunk_gpus):
    for index, video_name in enumerate(chunk_files):
        video_folder, _ = os.path.splitext(video_name)
        experiment_file_name = video_folder + "_experiment.yml"
        experiment_path = os.path.join(experiment_root,
                                        experiment_file_name)
        json_path = os.path.join(json_root,
                                 video_folder + ".json")
        gpu = chunk_gpus[index]
        yield [experiment_path, json_path, gpu]

def main(args):
    if not os.path.exists(args.chunk_json):
        raise OSError("Invalid path {0} provided for Chunk.json".format(args.chunk_json))
    if not os.path.exists(args.experiment_root):
        raise OSError("Invalid path {0} provided for experiment_root"\
                .format(args.experiment_root))
    if not os.path.exists(args.out_video_root):
        raise OSError("Invalid path {0} provided for json_root"\
                .format(args.out_video_root))

    num_gpus = len(args.gpus)
    assert num_gpus>0, "Chunk runner requires at least 1 gpu"

    chunks = json.load(open(args.chunk_json, 'r'))
    assert args.chunk_id in chunks.keys(), \
            "Invalid chunk id {0} provided".format(args.chunk_id)
    chunk_files = chunks[args.chunk_id]["files"]
    chunk_gpus = [ args.gpus[chunk_index%num_gpus]
                   for chunk_index in range(len(chunk_files)) ]
    pipeline_parameters_seq = \
            pipeline_parameter_generator(args.experiment_root,
                                         args.out_video_root,
                                         chunk_files,
                                         chunk_gpus)
    chunk_complete = False
    while not chunk_complete :
        pipeline_parameters = []
        for _ in range(num_gpus):
            try:
                experiment_path, json_path, gpu = next(pipeline_parameters_seq)
                if not os.path.exists(experiment_path):
                    print("Invalid path {0} provided for experiment_path".\
                          format(experiment_path))
                    continue
                print("Json path: {0}".format(json_path))
                pipeline_parameter = (experiment_path,
                                      args.rc3d_experiment_file,
                                      args.model_cfg,
                                      args.stride,
                                      gpu,
                                      args.activity_index,
                                      json_path)
                pipeline_parameters.append(pipeline_parameter)
            except StopIteration:
                chunk_complete = True
                break
        if len(pipeline_parameters)>0:
            with Pool(len(pipeline_parameters)) as p:
                pipeline_pool_completion = p.starmap(create_pipeline,
                                                     pipeline_parameters)
                p.close()
                for index, pipeline_completion in enumerate(pipeline_pool_completion):
                    if not pipeline_completion:
                        print("Pipeline {0} with the following parameters {1} failed".
                              format(index, pipeline_parameters[index]))
                    else:
                        print("Pipeline {0} with the following parameters {1} completed".
                              format(index, pipeline_parameters[index]))
    print("Finished processing Chunk: {0}".format(args.chunk_id))


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--gpus", nargs="+", type=int, default=[0,1,2,3],
                        help="gpu ids for rc3d")
    parser.add_argument("--chunk-id",
                        help="Chunk id over which the pipeline would run")
    parser.add_argument("--stride", type=int, default=8,
                        help="Temporal stride for rc3d")
    parser.add_argument("--chunk-json", help="JSON file with all the chunks")
    parser.add_argument("--experiment-root",
                        help="Root folder where experiment.yml are stored")
    parser.add_argument("--out-video-root",
                        help="Root folder where json file for a video is saved")
    parser.add_argument("--activity-index", default="",
                        help="Activity index provided by NIST")
    parser.add_argument("--rc3d-experiment-file",
                        default="/R-C3D/experiments/virat/experiment-docker.yml",
                        help="experiment file used by RC3D")
    parser.add_argument("--model-cfg",
                        default="/R-C3D/experiments/virat/td_cnn_end2end.yml",
                        help="Second experiment file used by RC3D")
    args = parser.parse_args()
    main(args)
