from argparse import ArgumentParser, Namespace
from distutils.util import strtobool
from pathlib import Path

NEPTUNE_API_KEY = 'eyJhcGlfYWRkcmVzcyI6Imh0dHBzOi8vdWkubmVwdHVuZS5haSIsImFwaV91cmwiOiJodHRwczovL3VpLm5lcHR1bmUu' \
                  'YWkiLCJhcGlfa2V5IjoiZmI0OGMzNzUtOTg1NS00Yzg2LThjMzYtMWFiYjUwMDUyMjVlIn0='

# Parameter Configuration
# =============================================================================
# Argument Parser

_ROOT = Path(__file__).parent

main_arg_parser = ArgumentParser(description="parser for fast-neural-style")

# Main Parameters
main_arg_parser.add_argument("--main_debug", type=strtobool, default=False, help="")
main_arg_parser.add_argument("--main_eval", type=strtobool, default=True, help="")
main_arg_parser.add_argument("--main_seed", type=int, default=69, help="")

# Data Parameters
main_arg_parser.add_argument("--data_class_name", type=str, default='Urban8K', help="")
main_arg_parser.add_argument("--data_worker", type=int, default=6, help="")
main_arg_parser.add_argument("--data_root", type=str, default='data', help="")
main_arg_parser.add_argument("--data_use_preprocessed", type=strtobool, default=True, help="")
main_arg_parser.add_argument("--data_n_mels", type=int, default=64, help="")
main_arg_parser.add_argument("--data_sr", type=int, default=16000, help="")
main_arg_parser.add_argument("--data_hop_length", type=int, default=256, help="")
main_arg_parser.add_argument("--data_n_fft", type=int, default=512, help="")
main_arg_parser.add_argument("--data_stretch", type=strtobool, default=True, help="")

# Transformation Parameters
main_arg_parser.add_argument("--data_loudness_ratio", type=float, default=0.0, help="") # 0.4
main_arg_parser.add_argument("--data_shift_ratio", type=float, default=0.0, help="")    # 0.4
main_arg_parser.add_argument("--data_noise_ratio", type=float, default=0, help="")      # 0.4
main_arg_parser.add_argument("--data_mask_ratio", type=float, default=0.3, help="")     # 0.2
main_arg_parser.add_argument("--data_speed_amount", type=float, default=0, help="")     # 0.4
main_arg_parser.add_argument("--data_speed_min", type=float, default=0, help="")        # 0.7
main_arg_parser.add_argument("--data_speed_max", type=float, default=0, help="")        # 1.7

# Model Parameters
# General
main_arg_parser.add_argument("--model_type", type=str, default="HorizontalVisualTransformer", help="")
main_arg_parser.add_argument("--model_weight_init", type=str, default="xavier_normal_", help="")
main_arg_parser.add_argument("--model_activation", type=str, default="gelu", help="")
main_arg_parser.add_argument("--model_bias", type=strtobool, default=True, help="")
main_arg_parser.add_argument("--model_norm", type=strtobool, default=True, help="")
main_arg_parser.add_argument("--model_dropout", type=float, default=0.2, help="")
main_arg_parser.add_argument("--model_lat_dim", type=int, default=128, help="")
main_arg_parser.add_argument("--model_features", type=int, default=64, help="")

# CNN Specific
main_arg_parser.add_argument("--model_filters", type=str, default="[32, 64, 128, 64]", help="")

# Transformer Specific
main_arg_parser.add_argument("--model_patch_size", type=int, default=3, help="")
main_arg_parser.add_argument("--model_attn_depth", type=int, default=3, help="")
main_arg_parser.add_argument("--model_heads", type=int, default=8, help="")
main_arg_parser.add_argument("--model_embedding_size", type=int, default=64, help="")

# Training Parameters
main_arg_parser.add_argument("--train_outpath", type=str, default="output", help="")
main_arg_parser.add_argument("--train_version", type=strtobool, required=False, help="")
main_arg_parser.add_argument("--train_sto_weight_avg", type=strtobool, default=False, help="")
main_arg_parser.add_argument("--train_weight_decay", type=float, default=0, help="")
main_arg_parser.add_argument("--train_opt_reset_interval", type=int, default=0, help="")
main_arg_parser.add_argument("--train_epochs", type=int, default=200, help="")
main_arg_parser.add_argument("--train_batch_size", type=int, default=200, help="")
main_arg_parser.add_argument("--train_lr", type=float, default=1e-3, help="")
main_arg_parser.add_argument("--train_lr_warmup_steps", type=int, default=10, help="")
main_arg_parser.add_argument("--train_num_sanity_val_steps", type=int, default=0, help="")

# Project Parameters
main_arg_parser.add_argument("--project_name", type=str, default=_ROOT.name, help="")
main_arg_parser.add_argument("--project_owner", type=str, default='si11ium', help="")
main_arg_parser.add_argument("--project_neptune_key", type=str, default=NEPTUNE_API_KEY, help="")

if __name__ == '__main__':
    # Parse it
    args: Namespace = main_arg_parser.parse_args()