import argparse
import torch.nn

def load_args():
    parser = argparse.ArgumentParser()

    # Pre training
    parser.add_argument('--train_num_perclass', type=int, default=40)
    # Pre training
    parser.add_argument('--is_feature_map', type=bool, default=False)
    parser.add_argument('--windowsize', type=int, default=11)
    parser.add_argument('--batch_size', type=int, default=32)
    parser.add_argument('--nclass', type=int, default=15)
    parser.add_argument('--epochs', type=int, default=1)
    parser.add_argument('--fine_tuned_epochs', type=int, default=200)
    parser.add_argument('--lr', type=float, default=1e-4)
    parser.add_argument('--fine_tuned_lr', type=float, default=1e-5)
    parser.add_argument('--init_prompt', type=str, default="a multimodal fusion patch of")
    parser.add_argument('--input_image_size',type=int, default=224)
    parser.add_argument('--type', type=str, default='none')#PCA none EMP
    parser.add_argument('--dataset', type=str, default='2013houston') #Trento  2013houston  Muufl  Berlin Augsburg
    parser.add_argument('--basemodel',type=str, default='ViT-B/16')
    parser.add_argument('--trainer', type=str, default='IPT')
    parser.add_argument('--valid_num', type=float, default=0.1)
    # Network
    parser.add_argument('--mask_ratio', default=0.7, type=float,help='mask_ratio (default: 0.8)')
    parser.add_argument('--mlp_ratio', default=2.0, type=float,
                        help='mlp ratio of encoder/decoder (default: 2.0)')
    parser.add_argument('--hid_chans', default=128, type=int,
                        help='hidden channels for dimension reduction (default: 128)')
    parser.add_argument('--hid_chans_LIDAR', default=128, type=int,
                        help='hidden channels for dimension reduction (default: 128)')
    # Augmentation
    parser.add_argument('--augment', default=False, type=bool,
                        help='either use data augmentation or not (default: False)')
    parser.add_argument('--scale', default=9, type=int,
                        help='the minimum scale for center crop (default: 19)')
    parser.add_argument('--proportion', default=0.05, type=float,)
    # MAE encoder specifics
    parser.add_argument('--encoder_dim', default=128, type=int,
                        help='feature dimension for encoder (default: 64)')
    parser.add_argument('--encoder_depth', default = 4, type=int,
                        help='encoder_depth; number of blocks ')
    parser.add_argument('--encoder_num_heads', default=8, type=int,
                        help='number of heads of encoder (default: 8)')    
    
    # MAE decoder specifics
    parser.add_argument('--decoder_dim', default= 128, type=int,
                        help='feature dimension for decoder (default: 64)')
    parser.add_argument('--deep_prompts_depth', default =12, type=int,
                        help='decoder_depth; number of blocks ')    
    parser.add_argument('--decoder_num_heads', default=8, type=int,
                        help='number of heads of decoder (default: 8)')
    parser.add_argument('--prompt_nctx', default=6, type=int,
                        help='decoder_depth; number of blocks ')
    # options for supervised MAE
    parser.add_argument('--temperature', default=1.0, type=float,
                        help='temperature for classification logits')
    parser.add_argument('--cls_loss_ratio', default=0.005, type=float,
                        help='ratio for classification loss')
    args = parser.parse_args()
    return args
