# Copyright (c) Microsoft Corporation

import argparse
import torch


class BaseOptions:
    def __init__(self):
        self.parser = argparse.ArgumentParser()
        self.initialized = False

    def initialize(self):
        # experiment specifics
        self.parser.add_argument(
            "--name",
            type=str,
            default="label2city",
            help="name of the experiment. It decides where to store samples and models",
        )
        self.parser.add_argument(
            "--gpu_ids",
            type=str,
            default="0",
            help="gpu ids: e.g. 0  0,1,2, 0,2. use -1 for CPU",
        )
        self.parser.add_argument(
            "--checkpoints_dir",
            type=str,
            default="./checkpoints",
            help="models are saved here",
        )  ## note: to add this param when using philly
        # self.parser.add_argument('--project_dir', type=str, default='./', help='the project is saved here')  ################### This is necessary for philly
        self.parser.add_argument(
            "--outputs_dir", type=str, default="./outputs", help="models are saved here"
        )  ## note: to add this param when using philly  Please end with '/'
        self.parser.add_argument(
            "--model", type=str, default="pix2pixHD", help="which model to use"
        )
        self.parser.add_argument(
            "--norm",
            type=str,
            default="instance",
            help="instance normalization or batch normalization",
        )
        self.parser.add_argument(
            "--use_dropout", action="store_true", help="use dropout for the generator"
        )
        self.parser.add_argument(
            "--data_type",
            default=32,
            type=int,
            choices=[8, 16, 32],
            help="Supported data type i.e. 8, 16, 32 bit",
        )
        self.parser.add_argument(
            "--verbose", action="store_true", default=False, help="toggles verbose"
        )

        # input/output sizes
        self.parser.add_argument(
            "--batchSize", type=int, default=1, help="input batch size"
        )
        self.parser.add_argument(
            "--loadSize", type=int, default=1024, help="scale images to this size"
        )
        self.parser.add_argument(
            "--fineSize", type=int, default=512, help="then crop to this size"
        )
        self.parser.add_argument(
            "--label_nc", type=int, default=35, help="# of input label channels"
        )
        self.parser.add_argument(
            "--input_nc", type=int, default=3, help="# of input image channels"
        )
        self.parser.add_argument(
            "--output_nc", type=int, default=3, help="# of output image channels"
        )

        # for setting inputs
        self.parser.add_argument(
            "--dataroot", type=str, default="./datasets/cityscapes/"
        )
        self.parser.add_argument(
            "--resize_or_crop",
            type=str,
            default="scale_width",
            help="scaling and cropping of images at load time [resize_and_crop|crop|scale_width|scale_width_and_crop]",
        )
        self.parser.add_argument(
            "--serial_batches",
            action="store_true",
            help="if true, takes images in order to make batches, otherwise takes them randomly",
        )
        self.parser.add_argument(
            "--no_flip",
            action="store_true",
            help="if specified, do not flip the images for data argumentation",
        )
        self.parser.add_argument(
            "--nThreads", default=2, type=int, help="# threads for loading data"
        )
        self.parser.add_argument(
            "--max_dataset_size",
            type=int,
            default=float("inf"),
            help="Maximum number of samples allowed per dataset. If the dataset directory contains more than max_dataset_size, only a subset is loaded.",
        )

        # for displays
        self.parser.add_argument(
            "--display_winsize", type=int, default=512, help="display window size"
        )
        self.parser.add_argument(
            "--tf_log",
            action="store_true",
            help="if specified, use tensorboard logging. Requires tensorflow installed",
        )

        # for generator
        self.parser.add_argument(
            "--netG", type=str, default="global", help="selects model to use for netG"
        )
        self.parser.add_argument(
            "--ngf", type=int, default=64, help="# of gen filters in first conv layer"
        )
        self.parser.add_argument(
            "--k_size", type=int, default=3, help="# kernel size conv layer"
        )
        self.parser.add_argument("--use_v2", action="store_true", help="use DCDCv2")
        self.parser.add_argument("--mc", type=int, default=1024, help="# max channel")
        self.parser.add_argument(
            "--start_r", type=int, default=3, help="start layer to use resblock"
        )
        self.parser.add_argument(
            "--n_downsample_global",
            type=int,
            default=4,
            help="number of downsampling layers in netG",
        )
        self.parser.add_argument(
            "--n_blocks_global",
            type=int,
            default=9,
            help="number of residual blocks in the global generator network",
        )
        self.parser.add_argument(
            "--n_blocks_local",
            type=int,
            default=3,
            help="number of residual blocks in the local enhancer network",
        )
        self.parser.add_argument(
            "--n_local_enhancers",
            type=int,
            default=1,
            help="number of local enhancers to use",
        )
        self.parser.add_argument(
            "--niter_fix_global",
            type=int,
            default=0,
            help="number of epochs that we only train the outmost local enhancer",
        )

        self.parser.add_argument(
            "--load_pretrain",
            type=str,
            default="",
            help="load the pretrained model from the specified location",
        )

        # for instance-wise features
        self.parser.add_argument(
            "--no_instance",
            action="store_true",
            help="if specified, do *not* add instance map as input",
        )
        self.parser.add_argument(
            "--instance_feat",
            action="store_true",
            help="if specified, add encoded instance features as input",
        )
        self.parser.add_argument(
            "--label_feat",
            action="store_true",
            help="if specified, add encoded label features as input",
        )
        self.parser.add_argument(
            "--feat_num", type=int, default=3, help="vector length for encoded features"
        )
        self.parser.add_argument(
            "--load_features",
            action="store_true",
            help="if specified, load precomputed feature maps",
        )
        self.parser.add_argument(
            "--n_downsample_E",
            type=int,
            default=4,
            help="# of downsampling layers in encoder",
        )
        self.parser.add_argument(
            "--nef",
            type=int,
            default=16,
            help="# of encoder filters in the first conv layer",
        )
        self.parser.add_argument(
            "--n_clusters", type=int, default=10, help="number of clusters for features"
        )

        # diy
        self.parser.add_argument(
            "--self_gen", action="store_true", help="self generate"
        )
        self.parser.add_argument(
            "--mapping_n_block",
            type=int,
            default=3,
            help="number of resblock in mapping",
        )
        self.parser.add_argument(
            "--map_mc", type=int, default=64, help="max channel of mapping"
        )
        self.parser.add_argument("--kl", type=float, default=0, help="KL Loss")
        self.parser.add_argument(
            "--load_pretrainA",
            type=str,
            default="",
            help="load the pretrained model from the specified location",
        )
        self.parser.add_argument(
            "--load_pretrainB",
            type=str,
            default="",
            help="load the pretrained model from the specified location",
        )
        self.parser.add_argument("--feat_gan", action="store_true")
        self.parser.add_argument("--no_cgan", action="store_true")
        self.parser.add_argument("--map_unet", action="store_true")
        self.parser.add_argument("--map_densenet", action="store_true")
        self.parser.add_argument("--fcn", action="store_true")
        self.parser.add_argument(
            "--is_image", action="store_true", help="train image recon only pair data"
        )
        self.parser.add_argument("--label_unpair", action="store_true")
        self.parser.add_argument("--mapping_unpair", action="store_true")
        self.parser.add_argument("--unpair_w", type=float, default=1.0)
        self.parser.add_argument("--pair_num", type=int, default=-1)
        self.parser.add_argument("--Gan_w", type=float, default=1)
        self.parser.add_argument("--feat_dim", type=int, default=-1)
        self.parser.add_argument("--abalation_vae_len", type=int, default=-1)

        ######################### useless, just to cooperate with docker
        self.parser.add_argument("--gpu", type=str)
        self.parser.add_argument("--dataDir", type=str)
        self.parser.add_argument("--modelDir", type=str)
        self.parser.add_argument("--logDir", type=str)
        self.parser.add_argument("--data_dir", type=str)

        self.parser.add_argument("--use_skip_model", action="store_true")
        self.parser.add_argument("--use_segmentation_model", action="store_true")

        self.parser.add_argument("--spatio_size", type=int, default=64)
        self.parser.add_argument("--test_random_crop", action="store_true")
        ##########################

        self.parser.add_argument("--contain_scratch_L", action="store_true")
        self.parser.add_argument(
            "--mask_dilation", type=int, default=0
        )  ## Don't change the input, only dilation the mask

        self.parser.add_argument(
            "--irregular_mask",
            type=str,
            default="",
            help="This is the root of the mask",
        )
        self.parser.add_argument(
            "--mapping_net_dilation",
            type=int,
            default=1,
            help="This parameter is the dilation size of the translation net",
        )

        self.parser.add_argument(
            "--VOC",
            type=str,
            default="VOC_RGB_JPEGImages.bigfile",
            help="The root of VOC dataset",
        )

        self.parser.add_argument(
            "--non_local", type=str, default="", help="which non_local setting"
        )
        self.parser.add_argument(
            "--NL_fusion_method",
            type=str,
            default="add",
            help="how to fuse the origin feature and nl feature",
        )
        self.parser.add_argument(
            "--NL_use_mask",
            action="store_true",
            help="If use mask while using Non-local mapping model",
        )
        self.parser.add_argument(
            "--correlation_renormalize",
            action="store_true",
            help="Since after mask out the correlation matrix(which is softmaxed), the sum is not 1 any more, enable this param to re-weight",
        )

        self.parser.add_argument(
            "--Smooth_L1", action="store_true", help="Use L1 Loss in image level"
        )

        self.parser.add_argument(
            "--face_restore_setting",
            type=int,
            default=1,
            help="This is for the aligned face restoration",
        )
        self.parser.add_argument("--face_clean_url", type=str, default="")
        self.parser.add_argument("--syn_input_url", type=str, default="")
        self.parser.add_argument("--syn_gt_url", type=str, default="")

        self.parser.add_argument(
            "--test_on_synthetic",
            action="store_true",
            help="If you want to test on the synthetic data, enable this parameter",
        )

        self.parser.add_argument(
            "--use_SN", action="store_true", help="Add SN to every parametric layer"
        )

        self.parser.add_argument(
            "--use_two_stage_mapping",
            action="store_true",
            help="choose the model which uses two stage",
        )

        self.parser.add_argument("--L1_weight", type=float, default=10.0)
        self.parser.add_argument("--softmax_temperature", type=float, default=1.0)
        self.parser.add_argument(
            "--patch_similarity",
            action="store_true",
            help="Enable this denotes using 3*3 patch to calculate similarity",
        )
        self.parser.add_argument(
            "--use_self",
            action="store_true",
            help="Enable this denotes that while constructing the new feature maps, using original feature (diagonal == 1)",
        )

        self.parser.add_argument("--use_own_dataset", action="store_true")

        self.parser.add_argument(
            "--test_hole_two_folders",
            action="store_true",
            help="Enable this parameter means test the restoration with inpainting given twp folders which are mask and old respectively",
        )

        self.parser.add_argument(
            "--no_hole",
            action="store_true",
            help="While test the full_model on non_scratch data, do not add random mask into the real old photos",
        )  ## Only for testing
        self.parser.add_argument(
            "--random_hole",
            action="store_true",
            help="While training the full model, 50% probability add hole",
        )

        self.parser.add_argument(
            "--NL_res", action="store_true", help="NL+Resdual Block"
        )

        self.parser.add_argument(
            "--image_L1", action="store_true", help="Image level loss: L1"
        )
        self.parser.add_argument(
            "--hole_image_no_mask",
            action="store_true",
            help="while testing, give hole image but not give the mask",
        )

        self.parser.add_argument(
            "--down_sample_degradation",
            action="store_true",
            help="down_sample the image only, corresponds to [down_sample_face]",
        )

        self.parser.add_argument(
            "--norm_G",
            type=str,
            default="spectralinstance",
            help="The norm type of Generator",
        )
        self.parser.add_argument(
            "--init_G",
            type=str,
            default="xavier",
            help="normal|xavier|xavier_uniform|kaiming|orthogonal|none",
        )

        self.parser.add_argument("--use_new_G", action="store_true")
        self.parser.add_argument("--use_new_D", action="store_true")

        self.parser.add_argument(
            "--only_voc",
            action="store_true",
            help="test the trianed celebA face model using VOC face",
        )

        self.parser.add_argument(
            "--cosin_similarity",
            action="store_true",
            help="For non-local, using cosin to calculate the similarity",
        )

        self.parser.add_argument(
            "--downsample_mode",
            type=str,
            default="nearest",
            help="For partial non-local, choose how to downsample the mask",
        )

        self.parser.add_argument(
            "--mapping_exp",
            type=int,
            default=0,
            help="Default 0: original PNL|1: Multi-Scale Patch Attention",
        )
        self.parser.add_argument(
            "--inference_optimize", action="store_true", help="optimize the memory cost"
        )

        self.initialized = True

    def parse(self, custom_args: list):

        assert len(custom_args) > 0, "Manually Pass Arguments!"

        if not self.initialized:
            self.initialize()

        self.opt = self.parser.parse_args(custom_args)
        self.opt.isTrain = self.isTrain  # train or test

        str_ids = self.opt.gpu_ids.split(",")
        self.opt.gpu_ids = []
        for str_id in str_ids:
            int_id = int(str_id)
            if int_id >= 0:
                self.opt.gpu_ids.append(int_id)

        # set gpu ids
        if torch.cuda.is_available() and len(self.opt.gpu_ids) > 0:
            if torch.cuda.device_count() > self.opt.gpu_ids[0]:
                try:
                    torch.cuda.set_device(self.opt.gpu_ids[0])
                except:
                    print("Failed to set GPU device. Using CPU...")

            else:
                print("Invalid GPU ID. Using CPU...")

        return self.opt