import argparse
from typing import IO, List, Optional, Union

import pytorch_lightning as pl
from loguru import logger
from omegaconf import DictConfig, ListConfig

from slp.config.omegaconf import OmegaConfExtended as OmegaConf
from slp.plbind import add_optimizer_args, add_trainer_args, add_tune_args


def make_cli_parser(
    parser: argparse.ArgumentParser, datamodule_cls: pl.LightningDataModule
) -> argparse.ArgumentParser:
    """make_cli_parser Augment an argument parser for slp with the default arguments

    Default arguments for training, logging, optimization etc. are added to the input {parser}.
    If you use make_cli_parser, the following command line arguments will be included

        usage: my_script.py [-h] [--hidden MODEL.INTERMEDIATE_HIDDEN]
                                        [--optimizer {Adam,AdamW,SGD,Adadelta,Adagrad,Adamax,ASGD,RMSprop}]
                                        [--lr OPTIM.LR] [--weight-decay OPTIM.WEIGHT_DECAY]
                                        [--lr-scheduler] [--lr-factor LR_SCHEDULE.FACTOR]
                                        [--lr-patience LR_SCHEDULE.PATIENCE]
                                        [--lr-cooldown LR_SCHEDULE.COOLDOWN]
                                        [--min-lr LR_SCHEDULE.MIN_LR] [--seed SEED] [--config CONFIG]
                                        [--experiment-name TRAINER.EXPERIMENT_NAME]
                                        [--run-id TRAINER.RUN_ID]
                                        [--experiment-group TRAINER.EXPERIMENT_GROUP]
                                        [--experiments-folder TRAINER.EXPERIMENTS_FOLDER]
                                        [--save-top-k TRAINER.SAVE_TOP_K]
                                        [--patience TRAINER.PATIENCE]
                                        [--wandb-project TRAINER.WANDB_PROJECT]
                                        [--tags [TRAINER.TAGS [TRAINER.TAGS ...]]]
                                        [--stochastic_weight_avg] [--gpus TRAINER.GPUS]
                                        [--val-interval TRAINER.CHECK_VAL_EVERY_N_EPOCH]
                                        [--clip-grad-norm TRAINER.GRADIENT_CLIP_VAL]
                                        [--epochs TRAINER.MAX_EPOCHS] [--steps TRAINER.MAX_STEPS]
                                        [--tbtt_steps TRAINER.TRUNCATED_BPTT_STEPS] [--debug]
                                        [--offline] [--early-stop-on TRAINER.EARLY_STOP_ON]
                                        [--early-stop-mode {min,max}] [--num-trials TUNE.NUM_TRIALS]
                                        [--gpus-per-trial TUNE.GPUS_PER_TRIAL]
                                        [--cpus-per-trial TUNE.CPUS_PER_TRIAL]
                                        [--tune-metric TUNE.METRIC] [--tune-mode {max,min}]
                                        [--val-percent DATA.VAL_PERCENT]
                                        [--test-percent DATA.TEST_PERCENT] [--bsz DATA.BATCH_SIZE]
                                        [--bsz-eval DATA.BATCH_SIZE_EVAL]
                                        [--num-workers DATA.NUM_WORKERS] [--no-pin-memory]
                                        [--drop-last] [--no-shuffle-eval]

        optional arguments:
          -h, --help            show this help message and exit
          --hidden MODEL.INTERMEDIATE_HIDDEN
                                                        Intermediate hidden layers for linear module
          --optimizer {Adam,AdamW,SGD,Adadelta,Adagrad,Adamax,ASGD,RMSprop}
                                                        Which optimizer to use
          --lr OPTIM.LR         Learning rate
          --weight-decay OPTIM.WEIGHT_DECAY
                                                        Learning rate
          --lr-scheduler        Use learning rate scheduling. Currently only
                                                        ReduceLROnPlateau is supported out of the box
          --lr-factor LR_SCHEDULE.FACTOR
                                                        Multiplicative factor by which LR is reduced. Used if
                                                        --lr-scheduler is provided.
          --lr-patience LR_SCHEDULE.PATIENCE
                                                        Number of epochs with no improvement after which
                                                        learning rate will be reduced. Used if --lr-scheduler
                                                        is provided.
          --lr-cooldown LR_SCHEDULE.COOLDOWN
                                                        Number of epochs to wait before resuming normal
                                                        operation after lr has been reduced. Used if --lr-
                                                        scheduler is provided.
          --min-lr LR_SCHEDULE.MIN_LR
                                                        Minimum lr for LR scheduling. Used if --lr-scheduler
                                                        is provided.
          --seed SEED           Seed for reproducibility
          --config CONFIG       Path to YAML configuration file
          --experiment-name TRAINER.EXPERIMENT_NAME
                                                        Name of the running experiment
          --run-id TRAINER.RUN_ID
                                                        Unique identifier for the current run. If not provided
                                                        it is inferred from datetime.now()
          --experiment-group TRAINER.EXPERIMENT_GROUP
                                                        Group of current experiment. Useful when evaluating
                                                        for different seeds / cross-validation etc.
          --experiments-folder TRAINER.EXPERIMENTS_FOLDER
                                                        Top-level folder where experiment results &
                                                        checkpoints are saved
          --save-top-k TRAINER.SAVE_TOP_K
                                                        Save checkpoints for top k models
          --patience TRAINER.PATIENCE
                                                        Number of epochs to wait before early stopping
          --wandb-project TRAINER.WANDB_PROJECT
                                                        Wandb project under which results are saved
          --tags [TRAINER.TAGS [TRAINER.TAGS ...]]
                                                        Tags for current run to make results searchable.
          --stochastic_weight_avg
                                                        Use Stochastic weight averaging.
          --gpus TRAINER.GPUS   Number of GPUs to use
          --val-interval TRAINER.CHECK_VAL_EVERY_N_EPOCH
                                                        Run validation every n epochs
          --clip-grad-norm TRAINER.GRADIENT_CLIP_VAL
                                                        Clip gradients with ||grad(w)|| >= args.clip_grad_norm
          --epochs TRAINER.MAX_EPOCHS
                                                        Maximum number of training epochs
          --steps TRAINER.MAX_STEPS
                                                        Maximum number of training steps
          --tbtt_steps TRAINER.TRUNCATED_BPTT_STEPS
                                                        Truncated Back-propagation-through-time steps.
          --debug               If true, we run a full run on a small subset of the
                                                        input data and overfit 10 training batches
          --offline             If true, forces offline execution of wandb logger
          --early-stop-on TRAINER.EARLY_STOP_ON
                                                        Metric for early stopping
          --early-stop-mode {min,max}
                                                        Minimize or maximize early stopping metric
          --num-trials TUNE.NUM_TRIALS
                                                        Number of trials to run for hyperparameter tuning
          --gpus-per-trial TUNE.GPUS_PER_TRIAL
                                                        How many gpus to use for each trial. If gpus_per_trial
                                                        < 1 multiple trials are packed in the same gpu
          --cpus-per-trial TUNE.CPUS_PER_TRIAL
                                                        How many cpus to use for each trial.
          --tune-metric TUNE.METRIC
                                                        Tune this metric. Need to be one of the keys of
                                                        metrics_map passed into make_trainer_for_ray_tune.
          --tune-mode {max,min}
                                                        Maximize or minimize metric
          --val-percent DATA.VAL_PERCENT
                                                        Percent of validation data to be randomly split from
                                                        the training set, if no validation set is provided
          --test-percent DATA.TEST_PERCENT
                                                        Percent of test data to be randomly split from the
                                                        training set, if no test set is provided
          --bsz DATA.BATCH_SIZE
                                                        Training batch size
          --bsz-eval DATA.BATCH_SIZE_EVAL
                                                        Evaluation batch size
          --num-workers DATA.NUM_WORKERS
                                                        Number of workers to be used in the DataLoader
          --no-pin-memory       Don't pin data to GPU memory when transferring
          --drop-last           Drop last incomplete batch
          --no-shuffle-eval     Don't shuffle val & test sets

    Args:
        parser (argparse.ArgumentParser): A parent argument to be augmented
        datamodule_cls (pytorch_lightning.LightningDataModule): A data module class that injects arguments through the add_argparse_args method

    Returns:
        argparse.ArgumentParser: The augmented command line parser

    Examples:
        >>> import argparse
        >>> from slp.plbind.dm import PLDataModuleFromDatasets
        >>> parser = argparse.ArgumentParser("My cool model")
        >>> parser.add_argument("--hidden", dest="model.hidden", type=int)  # Create parser with model arguments and anything else you need
        >>> parser = make_cli_parser(parser, PLDataModuleFromDatasets)
        >>> args = parser.parse_args(args=["--bsz", "64", "--lr", "0.01"])
        >>> args.data.batch_size
        64
        >>> args.optim.lr
        0.01
    """
    parser = add_optimizer_args(parser)
    parser = add_trainer_args(parser)
    parser = add_tune_args(parser)
    parser = datamodule_cls.add_argparse_args(parser)

    return parser


def parse_config(
    parser: argparse.ArgumentParser,
    config_file: Optional[Union[str, IO]],
    args: Optional[List[str]] = None,
    include_none: bool = False,
) -> Union[ListConfig, DictConfig]:
    """parse_config Parse a provided YAML config file and command line args and merge them

    During experimentation we want ideally to have a configuration file with the model and training configuration,
    but also be able to run quick experiments using command line args.
    This function allows you to double dip, by overriding values in a YAML config file through user provided command line arguments.

    The precedence for merging is as follows
       * default cli args values < config file values < user provided cli args

    E.g.:

       * if you don't include a value in your configuration it will take the default value from the argparse arguments
       * if you provide a cli arg (e.g. run the script with --bsz 64) it will override the value in the config file

    Note we use an extended OmegaConf istance to achieve this (see slp.config.omegaconf.OmegaConf)

    Args:
        parser (argparse.ArgumentParser): The argument parser you want to use
        config_file (Union[str, IO]): Configuration file name or file descriptor
        args (Optional[List[str]]): Optional input sys.argv style args. Useful for testing.
            Use this only for testing. By default it uses sys.argv[1:]

    Returns:
        OmegaConf.DictConfig: The parsed configuration as an OmegaConf DictConfig object

    Examples:
        >>> import io
        >>> from slp.config.config_parser import parse_config
        >>> mock_config_file = io.StringIO('''
        model:
          hidden: 100
        ''')
        >>> parser = argparse.ArgumentParser("My cool model")
        >>> parser.add_argument("--hidden", dest="model.hidden", type=int, default=20)
        >>> cfg = parse_config(parser, mock_config_file)
        {'model': {'hidden': 100}}
        >>> type(cfg)
        <class 'omegaconf.dictconfig.DictConfig'>
        >>> cfg = parse_config(parser, mock_config_file, args=["--hidden", "200"])
        {'model': {'hidden': 200}}
        >>> mock_config_file = io.StringIO('''
        random_value: hello
        ''')
        >>> cfg = parse_config(parser, mock_config_file)
        {'model': {'hidden': 20}, 'random_value': 'hello'}
    """
    # Merge Configurations Precedence: default kwarg values < default argparse values < config file values < user provided CLI args values

    if config_file is not None:
        dict_config = OmegaConf.from_yaml(config_file)  # type: ignore
    else:
        dict_config = OmegaConf.create({})

    user_cli, default_cli = OmegaConf.from_argparse(parser, include_none=include_none)
    config = OmegaConf.merge(default_cli, dict_config, user_cli)

    logger.info("Running with the following configuration")
    logger.info(f"\n{OmegaConf.to_yaml(config)}")

    return config


def generate_example_config(
    parser: argparse.ArgumentParser,
    output_file: str,
    args: Optional[List[str]] = None,
) -> None:
    """parse_config Parse a provided YAML config file and command line args and merge them

    During experimentation we want ideally to have a configuration file with the model and training configuration,
    but also be able to run quick experiments using command line args.
    This function allows you to double dip, by overriding values in a YAML config file through user provided command line arguments.

    The precedence for merging is as follows
       * default cli args values < config file values < user provided cli args

    E.g.:

       * if you don't include a value in your configuration it will take the default value from the argparse arguments
       * if you provide a cli arg (e.g. run the script with --bsz 64) it will override the value in the config file

    Note we use an extended OmegaConf istance to achieve this (see slp.config.omegaconf.OmegaConf)

    Args:
        parser (argparse.ArgumentParser): The argument parser you want to use
        output_file (Union[str, IO]): Configuration file name or file descriptor to save example configuration
        args (Optional[List[str]]): Optional input sys.argv style args. Useful for testing.
            Use this only for testing. By default it uses sys.argv[1:]
    """
    config = parse_config(parser, None, include_none=True)
    OmegaConf.save(config, output_file)
