import argparse
def parse_args_and_update_config():
    parser = argparse.ArgumentParser()

    # In-N-Out Path
    parser.add_argument('--fig_out_dir', type=str, default='./figs', help="The output directory of visualization.")
    parser.add_argument('--res_in_dir', type=str, default='./results', help="The input directory of results.")
    parser.add_argument('--data_in_dir', type=str, default='./datasets/coco', help="Please refer to the specific dataset directory.")

    # Model
    parser.add_argument('--sdm_version', type=str, default='sdm_2_1_base', help="Select between: 'sdm_2_0_base' | 'sdm_2_1_base'.")
    parser.add_argument('--int_mode', type=str, default='logistic', help="Select between 'uniform' | 'logistic'.")
    parser.add_argument('--n_samples_per_point', type=int, default=1, help="If 'res_type == nll_XD', set it as 1; elif 'res_type == mse_XD'set it as 50.")
    parser.add_argument('--num_steps', type=int, default=100, help="If 'res_type == mses_2D', set it as 10; else set it as 100.")

    # Dataset
    parser.add_argument('--dataset_type', type=str, default="COCO-IT", help="Select among: 'COCO-IT' | 'coco_ours'.")
    parser.add_argument('--csv_name', type=str, default="COCO100-IT", help="Select between: 'COCO100-IT' | 'COCO-WL'.")

    # Random seed
    parser.add_argument('--seed', type=int, default=42, help="101010")

    # Others
    parser.add_argument('--eval_metrics', type=str, default='pid', help="Select between 'mi' | 'cmi'.")
    parser.add_argument('--visual_type', type=str, default='mmse_curve', help="Select between: 'mmse_curve' | 'scatter_plot' | 'denoising_diffusion' | 'mi_cmi_attn' | 'image_edit'.")
    parser.add_argument('--norm_over_all_data', action='store_true', help="")
    parser.add_argument('--norm_over_all_maps', action='store_true', help="")
    parser.add_argument('--exp', action='store_true', help="")
    parser.add_argument('--cmap', type=str, default='jet', help="Select between: 'jet' | 'gray'")

    args = parser.parse_args()
    return args

