"""
A python wrapper around ChemProp that trains separate models using features
generated by different CP3D models.
"""


import argparse

from nff.utils import parse_args, CHEMPROP_METRICS
from nff.io.cprop import (cp_train, cp_hyperopt,
                             modify_config, modify_hyp_config)


def main(base_config_path,
         hyp_config_path,
         train_folder,
         metric,
         train_feat_path,
         val_feat_path,
         test_feat_path,
         cp_folder,
         features_only,
         use_hyperopt,
         rerun_hyperopt,
         no_features,
         **kwargs):
    """
    Load pre-set features to train a ChemProp model.
    Args:
      base_config_path (str): where your basic job config file
        is, with parameters that may or may not be changed depending
        on the given run
      hyp_config_path (str): where your basic hyperopt job config file
        is, with parameters that may or may not be changed depending
        on the given run
      train_folder (str): where you want to store your trained models
      metric (str): what metric you want to optimize in this run
      train_feat_path (str): where the features of your training set are
      val_feat_path (str): where the features of your validation set are
      test_feat_path (str): where the features of your test set are
      cp_folder (str): path to the chemprop folder on your computer
      features_only (bool): whether to just train with the features and no
        MPNN
      use_hyperopt (bool): do a hyperparameter optimization before training
        the model
      rerun_hyperopt (bool): whether to rerun hyperparameter optimization if
        `hyp_folder` already exists and has the completion file
        `best_params.json`.
      no_features (bool): Don't use external features when training model.
    Returns:
      None
    """

    # if doing a hyperparameter optimization, run the optimization and get
    # the best parameters

    if use_hyperopt:

        hyp_feat_path = train_feat_path.replace("train", "hyperopt")
        hyp_folder = train_folder + "_hyp"

        modify_hyp_config(hyp_config_path=hyp_config_path,
                          metric=metric,
                          hyp_feat_path=hyp_feat_path,
                          hyp_folder=hyp_folder,
                          features_only=features_only,
                          no_features=no_features)

        hyp_params = cp_hyperopt(cp_folder=cp_folder,
                                 hyp_folder=hyp_folder,
                                 rerun=rerun_hyperopt)

    else:
        hyp_params = {}

    # modify the base config with the parameters specific to this
    # run -- optimized hyperparameters, train/val/test features paths,
    # metric, etc.

    modify_config(base_config_path=base_config_path,
                  metric=metric,
                  train_feat_path=train_feat_path,
                  val_feat_path=val_feat_path,
                  test_feat_path=test_feat_path,
                  train_folder=train_folder,
                  features_only=features_only,
                  hyp_params=hyp_params,
                  no_features=no_features)

    # train the model

    cp_train(cp_folder=cp_folder,
             train_folder=train_folder)


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--base_config_path", type=str,
                        help=("Path to the reference config file "
                              "used to train a ChemProp model. "
                              "This file will be modified with "
                              "the arguments specified below "
                              "(metric and features paths)."
                              "If they are not specified then "
                              "the config file will not be "
                              "modified."))
    parser.add_argument("--hyp_config_path", type=str, default=None,
                        help=("Same as `base_config_path`, but "
                              "for the hyperparameter optimization "
                              "stage"))
    parser.add_argument("--use_hyperopt", action='store_true',
                        help=("Do hyperparameter optimization before "
                              "training "))
    parser.add_argument("--rerun_hyperopt", action='store_true',
                        help=("Rerun hyperparameter optimization even if "
                              "it has been done already. "))

    parser.add_argument("--metric", type=str,
                        choices=CHEMPROP_METRICS,
                        help=("Metric for which to evaluate "
                              "the model performance"),
                        default=None)
    parser.add_argument("--train_feat_path", type=str,
                        help=("Path to features file for training set"),
                        default=None)
    parser.add_argument("--val_feat_path", type=str,
                        help=("Path to features file for validation set"),
                        default=None)
    parser.add_argument("--test_feat_path", type=str,
                        help=("Path to features file for test set"),
                        default=None)
    parser.add_argument("--train_folder", type=str,
                        help=("Folder in which you will store the "
                              "ChemProp model."),
                        default=None)
    parser.add_argument("--features_only", action='store_true',
                        help=("Train model with only the stored features"))
    parser.add_argument("--cp_folder", type=str,
                        help=("Path to ChemProp folder."))
    parser.add_argument("--no_features", action="store_true",
                        help=("Don't use external features when training "
                              "model."))
    parser.add_argument('--this_config_file', type=str,
                        help=("Path to JSON file with arguments "
                              "for this script. If given, any "
                              "arguments in the file override the "
                              "command line arguments."))

    args = parse_args(parser, config_flag="this_config_file")
    main(**args.__dict__)
