import argparse
import torch
import os
import pwd
from pytorch_lightning import Trainer

EMPTY_NAME_ERR = 'Name of augmentation or one of its arguments cant be empty\n\
                  Use "name/arg1=value/arg2=value" format'
POSS_VAL_NOT_LIST = (
    "Flag {} has an invalid list of values: {}. Length of list must be >=1"
)


def parse_augmentations(raw_augmentations):
    """
    Parse the list of augmentations, given by configuration, into a list of
    tuple of the augmentations name and a dictionary containing additional args.

    The augmentation is assumed to be of the form 'name/arg1=value/arg2=value'

    :raw_augmentations: list of strings [unparsed augmentations]
    :returns: list of parsed augmentations [list of (name,additional_args)]

    """
    augmentations = []
    for t in raw_augmentations:
        arguments = t.split("/")
        name = arguments[0]
        if name == "":
            raise Exception(EMPTY_NAME_ERR)

        kwargs = {}
        if len(arguments) > 1:
            for a in arguments[1:]:
                splited = a.split("=")
                var = splited[0]
                val = splited[1] if len(splited) > 1 else None
                if var == "":
                    raise Exception(EMPTY_NAME_ERR)

                kwargs[var] = val

        augmentations.append((name, kwargs))

    return augmentations


def parse_dispatcher_config(config):
    """
    Parses an experiment config, and creates jobs. For flags that are expected to be a single item,
    but the config contains a list, this will return one job for each item in the list.
    :config - experiment_config

    returns: jobs - a list of flag strings, each of which encapsulates one job.
        *Example: --train --cuda --dropout=0.1 ...
    returns: experiment_axies - axies that the grid search is searching over
    """
    jobs = [""]
    experiment_axies = []
    search_spaces = config["search_space"]

    # Support a list of search spaces, convert to length one list for backward compatiblity
    if not isinstance(search_spaces, list):
        search_spaces = [search_spaces]

    for search_space in search_spaces:
        # Go through the tree of possible jobs and enumerate into a list of jobs
        for ind, flag in enumerate(search_space):
            possible_values = search_space[flag]
            if len(possible_values) > 1:
                experiment_axies.append(flag)

            children = []
            if len(possible_values) == 0 or type(possible_values) is not list:
                raise Exception(POSS_VAL_NOT_LIST.format(flag, possible_values))
            for value in possible_values:
                for parent_job in jobs:
                    if type(value) is bool:
                        if value:
                            new_job_str = "{} --{}".format(parent_job, flag)
                        else:
                            new_job_str = parent_job
                    elif type(value) is list:
                        val_list_str = " ".join([str(v) for v in value])
                        new_job_str = "{} --{} {}".format(
                            parent_job, flag, val_list_str
                        )
                    else:
                        new_job_str = "{} --{} {}".format(parent_job, flag, value)
                    children.append(new_job_str)
            jobs = children

    return jobs, experiment_axies


def parse_args(args_strings=None):
    parser = argparse.ArgumentParser(
        description="Sandstone research repo. Support Mammograms, CT Scans, Thermal Imaging, Cell Imaging and Chemistry."
    )
    # setup
    parser.add_argument(
        "--train",
        action="store_true",
        default=False,
        help="Whether or not to train model",
    )
    parser.add_argument(
        "--test",
        action="store_true",
        default=False,
        help="Whether or not to run model on test set",
    )
    parser.add_argument(
        "--dev",
        action="store_true",
        default=False,
        help="Whether or not to run model on dev set",
    )
    parser.add_argument(
        "--fine_tune",
        action="store_true",
        default=False,
        help="Whether or not to fine_tune model",
    )
    parser.add_argument(
        "--num_epochs_fine_tune",
        type=int,
        default=1,
        help="Num epochs to finetune model",
    )

    # data
    parser.add_argument(
        "--dataset",
        default="nlst",
        choices=[
            "sybil",
            "nlst",
            "nlst_risk_factors",
            "nlst_for_plco2012",
            "nlst_for_plco2019" "mgh",
        ],
        help="Name of dataset from dataset factory to use [default: nlst]",
    )
    parser.add_argument(
        "--img_size",
        type=int,
        nargs="+",
        default=[256, 256],
        help="Width and height of image in pixels. [default: [256,256]]",
    )
    parser.add_argument(
        "--num_chan", type=int, default=3, help="Number of channels for input image"
    )
    parser.add_argument(
        "--img_mean",
        type=float,
        nargs="+",
        default=[128.1722],
        help="Mean of image per channel",
    )
    parser.add_argument(
        "--img_std",
        type=float,
        nargs="+",
        default=[87.1849],
        help="Standard deviation  of image per channel",
    )
    parser.add_argument(
        "--img_dir",
        type=str,
        default="/data/rsg/mammogram/NLST/nlst-ct-png",
        help="Dir of images. Note, image path in dataset jsons should stem from here",
    )
    parser.add_argument(
        "--img_file_type",
        type=str,
        default="png",
        choices=["png", "dicom"],
        help="Type of image. one of [png, dicom]",
    )
    parser.add_argument(
        "--fix_seed_for_multi_image_augmentations",
        action="store_true",
        default=False,
        help="Use same seed for each slice of volume augmentations",
    )
    parser.add_argument(
        "--dataset_file_path",
        type=str,
        default="/Mounts/rbg-storage1/datasets/NLST/full_nlst_google.json",
        help="Path to dataset file either as json or csv",
    )
    parser.add_argument(
        "--num_classes", type=int, default=6, help="Number of classes to predict"
    )

    # Alternative training/testing schemes
    parser.add_argument(
        "--cross_val_seed",
        type=int,
        default=0,
        help="Seed used to generate the partition.",
    )
    parser.add_argument(
        "--assign_splits",
        action="store_true",
        default=False,
        help="Whether to assign different splits than those predetermined in dataset",
    )
    parser.add_argument(
        "--split_type",
        type=str,
        default="random",
        choices=["random", "institution_split"],
        help="How to split dataset if assign_split = True. Usage: ['random', 'institution_split'].",
    )
    parser.add_argument(
        "--split_probs",
        type=float,
        nargs="+",
        default=[0.6, 0.2, 0.2],
        help="Split probs for datasets without fixed train dev test. ",
    )

    # survival analysis setup
    parser.add_argument(
        "--max_followup", type=int, default=6, help="Max followup to predict over"
    )

    # risk factors
    parser.add_argument(
        "--use_risk_factors",
        action="store_true",
        default=False,
        help="Whether to feed risk factors into last FC of model.",
    )  #
    parser.add_argument(
        "--risk_factor_keys",
        nargs="*",
        default=[],
        help="List of risk factors to include in risk factor vector.",
    )

    # handling CT slices
    parser.add_argument(
        "--resample_pixel_spacing_prob",
        type=float,
        default=1,
        help="Probability of resampling pixel spacing into fixed dimensions. 1 when eval and using resampling",
    )
    parser.add_argument(
        "--num_images",
        type=int,
        default=200,
        help="In multi image setting, the number of images per single sample.",
    )
    parser.add_argument(
        "--min_num_images",
        type=int,
        default=0,
        help="In multi image setting, the min number of images per single sample.",
    )
    parser.add_argument(
        "--slice_thickness_filter",
        type=float,
        help="Slice thickness using, if restricting to specific thickness value.",
    )
    parser.add_argument(
        "--use_only_thin_cuts_for_ct",
        action="store_true",
        default=False,
        help="Wether to use image series with thinnest cuts only.",
    )

    # region annotations
    parser.add_argument(
        "--use_annotations",
        action="store_true",
        default=False,
        help="whether to use image annotations (pixel labels) in modeling",
    )

    parser.add_argument(
        "--region_annotations_filepath", type=str, help="Path to annotations file"
    )
    parser.add_argument(
        "--annotation_loss_lambda",
        type=float,
        default=1,
        help="Weight of annotation losses",
    )
    parser.add_argument(
        "--image_attention_loss_lambda",
        type=float,
        default=1,
        help="Weight of loss for predicting image attention scores",
    )
    parser.add_argument(
        "--volume_attention_loss_lambda",
        type=float,
        default=1,
        help="Weight of loss for predicting volume attention scores",
    )

    # regularization
    parser.add_argument(
        "--primary_loss_lambda",
        type=float,
        default=1.0,
        help="Lambda to weigh the primary loss.",
    )
    parser.add_argument(
        "--adv_loss_lambda",
        type=float,
        default=1.0,
        help="Lambda to weigh the adversary loss.",
    )

    # learning
    parser.add_argument(
        "--batch_size",
        type=int,
        default=32,
        help="Batch size for training [default: 128]",
    )
    parser.add_argument(
        "--init_lr",
        type=float,
        default=0.001,
        help="Initial learning rate [default: 0.001]",
    )
    parser.add_argument(
        "--dropout",
        type=float,
        default=0.25,
        help="Amount of dropout to apply on last hidden layer [default: 0.25]",
    )
    parser.add_argument(
        "--optimizer", type=str, default="adam", help="Optimizer to use [default: adam]"
    )
    parser.add_argument(
        "--momentum", type=float, default=0, help="Momentum to use with SGD"
    )
    parser.add_argument(
        "--lr_decay",
        type=float,
        default=0.1,
        help="Initial learning rate [default: 0.5]",
    )
    parser.add_argument(
        "--weight_decay",
        type=float,
        default=0,
        help="L2 Regularization penaty [default: 0]",
    )
    parser.add_argument(
        "--adv_lr",
        type=float,
        default=0.001,
        help="Initial learning rate for adversary model [default: 0.001]",
    )

    # schedule
    parser.add_argument(
        "--patience",
        type=int,
        default=5,
        help="Number of epochs without improvement on dev before halving learning rate and reloading best model [default: 5]",
    )
    parser.add_argument(
        "--num_adv_steps",
        type=int,
        default=1,
        help="Number of steps for domain adaptation discriminator per one step of encoding model [default: 5]",
    )
    parser.add_argument(
        "--tuning_metric",
        type=str,
        default="c_index",
        help="Criterion based on which model is saved [default: c_index]",
    )

    # model checkpointing
    parser.add_argument(
        "--turn_off_checkpointing",
        action="store_true",
        default=False,
        help="Do not save best model",
    )

    parser.add_argument(
        "--save_dir", type=str, default="snapshot", help="Where to dump the model"
    )

    parser.add_argument(
        "--snapshot",
        type=str,
        default=None,
        help="Filename of model snapshot to load[default: None]",
    )

    # system
    parser.add_argument(
        "--num_workers",
        type=int,
        default=8,
        help="Num workers for each data loader [default: 4]",
    )

    # storing results
    parser.add_argument(
        "--store_hiddens",
        action="store_true",
        default=False,
        help="Save hidden repr from each image to an npz based off results path, git hash and exam name",
    )
    parser.add_argument(
        "--save_predictions",
        action="store_true",
        default=False,
        help="Save hidden repr from each image to an npz based off results path, git hash and exam name",
    )
    parser.add_argument(
        "--hiddens_dir",
        type=str,
        default="hiddens/test_run",
        help='Dir to store hiddens npy"s when store_hiddens is true',
    )
    parser.add_argument(
        "--save_attention_scores",
        action="store_true",
        default=False,
        help="Whether to save attention scores when using attention mechanism",
    )
    parser.add_argument(
        "--results_path",
        type=str,
        default="logs/test.args",
        help="Where to save the result logs",
    )

    # cache
    parser.add_argument(
        "--cache_path", type=str, default=None, help="Dir to cache images."
    )
    parser.add_argument(
        "--cache_full_img",
        action="store_true",
        default=False,
        help="Cache full image locally as well as cachable transforms",
    )

    # run
    parser = Trainer.add_argparse_args(parser)
    if args_strings is None:
        args = parser.parse_args()
    else:
        args = parser.parse_args(args_strings)
    args.lr = args.init_lr

    if (isinstance(args.gpus, str) and len(args.gpus.split(",")) > 1) or (
        isinstance(args.gpus, int) and args.gpus > 1
    ):
        args.accelerator = "ddp"
        args.replace_sampler_ddp = False
    else:
        args.accelerator = None
        args.replace_sampler_ddp = False

    args.unix_username = pwd.getpwuid(os.getuid())[0]

    # learning initial state
    args.step_indx = 1

    return args