import copy
import json
import os
import sys
from typing import Dict, Any, Optional

from allennlp import commands
from allennlp.common.params import parse_overrides


def maybe_convert_to_numeric(v: Any):
    try:
        return int(v)  # this will execute if string v is integer
    except ValueError:
        pass
    try:
        return float(v)  # this will execute if string v is float
    except ValueError:
        pass

    return v


def integrate_arguments_and_run_allennlp(
        serialization_dir: Optional[str] = None, overrides_default: Optional[Dict] = None
):
    """
    This script processes the arguments generated by WandB sweep and then execute allennlp train command with updated
    arguments. When a wandb agent is started, it runs a script (SweepScript) provided as value of 'program' key
    in SweepConfig. SweepScript takes arguments from 'command' key in the SweepConfig along with the arguments
    generated by the sweep. These extra arguments are hyper parameters selected by wandb for that particular run.
    SweepScript then calls main method of this script with some meta data related to sweep_run and some additional
    arguments required for training with cross validation. This script processes the arguments in a format that allennlp
    train command can parse. As an example of this whole process,
    WandB agent is started as:
    $wandb agent PROJECT/SWEEP_ID
    (This is autogenerated command on WandB web interface when we create sweep from
    wandb.ai/TEAM_NAME/PROJECT_NAME/create-sweep)

    This command in return generate another command by reading sweep config file, e.g:
    $python SweepScript.py train
    -s experiments/adu/sciarg_scibert_split_FIXEDOFFSETS_HPT
    -f training_scripts/hpt/adu_sciarg_scibert_split.jsonnet
    -o {data_loader.batch_size:16}
    --trainer.optimizer.lr=0.0005

    Here 'trainer.optimizer.lr' is generated by sweep and needs to be processed before passing to allennlp train command.
    It is added to overwrite argument of allennlp train command such that final command becomes
    $allennlp train
    -s experiments/adu/sciarg_scibert_split_FIXEDOFFSETS_HPT
    -f training_scripts/hpt/adu_sciarg_scibert_split.jsonnet
    -o {data_loader.batch_size:16, trainer.optimizer.lr:0.0005}
    In addition to this, some additional arguments passed as parameter to this method, can also be added to overwrite
    argument. General description of these parameters is given below:

    :param serialization_dir: path to a directory where meta data and runs will be stored, if provided then value is
                            passed to -s (--serialization) argument of allennlp train command
    :param overrides_default: overrides default parameters of allennlp config, this parameter expects a dictionary with
                            keys related to cross validation and is appended to 'dataset_reader' key of allennlp config
    :return: None
    """

    args = sys.argv
    overrides = {}
    if serialization_dir is not None:
        args.extend(['-s', serialization_dir])
    if overrides_default is not None:
        overrides = copy.deepcopy(overrides_default)
    keys = ["--model", "--trainer", "--dataset_reader", "--data_loader"]
    print(f'original arguments:{args}')
    rem_keys = []
    i = 0
    while i < len(sys.argv):
        flag = False
        # Interpret all arguments starting with 'ENV_' as environment variable and delete from sys arguments
        if sys.argv[i].startswith('--ENV_'):
            env_name, env_value = sys.argv[i].split('=')
            env_name = env_name[2:]
            del sys.argv[i]
            os.environ[env_name] = env_value
            continue
        # Interpret all arguments starting with '-o' or '--overrides' as overrides for allennlp command and delete
        # from sys arguments
        if sys.argv[i] == '-o' or sys.argv[i] == '--overrides':
            overrides.update(parse_overrides(sys.argv[i + 1]))
            del sys.argv[i]
            del sys.argv[i]
            continue
        # Save all the arguments passed by the sweep and delete from sys arguments
        for prefix in keys:
            if sys.argv[i].startswith(prefix):
                rem_keys.append(sys.argv[i])
                del sys.argv[i]
                flag = True
        if not flag:
            i += 1

    for arg in rem_keys:
        # remove leading dashes "--"
        arg = arg[2:]
        # split key and value
        arg_k, arg_v = arg.split("=")
        # numerics have to be converted to be recognized correctly by AllenNLP
        arg_v = maybe_convert_to_numeric(arg_v)
        overrides[arg_k] = arg_v

    sys.argv.append("-o")
    sys.argv.append(json.dumps(overrides))
    print(f'updated arguments:{sys.argv}')
    commands.main()


if __name__ == '__main__':
    integrate_arguments_and_run_allennlp()
