import argparse
import time
import os

from dataloader.paths import kitti_dir, fcgf_weights_file
from dataloader.data_loaders import get_dataset_name

arg_lists = []
parser = argparse.ArgumentParser()


def add_argument_group(name):
    arg = parser.add_argument_group(name)
    arg_lists.append(arg)
    return arg


def str2bool(v):
    return v.lower() in ('true', '1')

dataset_arg = add_argument_group('Dataset')
dataset_arg.add_argument('--dataset', type=str, default=None, help='name of dataset for testing')

# snapshot configurations
snapshot_arg = add_argument_group('Snapshot')
snapshot_arg.add_argument('--snapshot_dir', type=str, default=None)
snapshot_arg.add_argument('--tboard_dir', type=str, default=None)
snapshot_arg.add_argument('--snapshot_interval', type=int, default=1)
snapshot_arg.add_argument('--save_dir', type=str, default=None)

# Network configurations
net_arg = add_argument_group('Network')
net_arg.add_argument('--in_dim', type=int, default=6)
net_arg.add_argument('--num_layers', type=int, default=12)
net_arg.add_argument('--num_channels', type=int, default=128)
net_arg.add_argument('--num_iterations', type=int, default=10, help='power iteration algorithm')
net_arg.add_argument('--ratio', type=float, default=0.1, help='max ratio of seeding points')
net_arg.add_argument('--k', type=int, default=40, help='size of local neighborhood')

# Loss configurations
loss_arg = add_argument_group('Loss')
loss_arg.add_argument('--evaluate_interval', type=int, default=10)
loss_arg.add_argument('--balanced', type=str2bool, default=False)
loss_arg.add_argument('--weight_classification', type=float, default=1.0)
loss_arg.add_argument('--weight_spectralmatching', type=float, default=1.0)
loss_arg.add_argument('--weight_transformation', type=float, default=0.0)
loss_arg.add_argument('--transformation_loss_start_epoch', type=int, default=0)

# Optimizer configurations
opt_arg = add_argument_group('Optimizer')
opt_arg.add_argument('--optimizer', type=str, default='ADAM', choices=['SGD', 'ADAM'])
opt_arg.add_argument('--max_epoch', type=int, default=50)
opt_arg.add_argument('--training_max_iter', type=int, default=3500)
opt_arg.add_argument('--val_max_iter', type=int, default=1000)
opt_arg.add_argument('--lr', type=float, default=1e-4)
opt_arg.add_argument('--weight_decay', type=float, default=1e-6)
opt_arg.add_argument('--momentum', type=float, default=0.9)
opt_arg.add_argument('--scheduler', type=str, default='ExpLR')
opt_arg.add_argument('--scheduler_gamma', type=float, default=0.99)
opt_arg.add_argument('--scheduler_interval', type=int, default=1)

# Dataset and dataloader configurations
data_arg = add_argument_group('Data')

data_arg.add_argument('--root', type=str, default=None)
data_arg.add_argument('--descriptor', type=str, default=None, choices=['d3feat', 'fpfh', 'fcgf', None])
data_arg.add_argument('--inlier_threshold', type=float, default=None)
net_arg.add_argument('--sigma_d', type=float, default=None)
data_arg.add_argument('--downsample', type=float, default=None)
data_arg.add_argument('--re_thre', type=float, default=None, help='rotation error thrshold (deg)')
data_arg.add_argument('--te_thre', type=float, default=None, help='translation error thrshold (cm)')

data_arg.add_argument('--fcgf_weights_file', type=str, default=None, help='file containing FCGF network weights')
data_arg.add_argument('--num_node', type=int, default=1000)
data_arg.add_argument('--use_mutual', type=str2bool, default=False)
data_arg.add_argument('--augment_axis', type=int, default=3)
data_arg.add_argument('--augment_rotation', type=float, default=1.0, help='rotation angle = num * 2pi')
data_arg.add_argument('--augment_translation', type=float, default=0.5, help='translation = num (m)')
data_arg.add_argument('--batch_size', type=int, default=16)
data_arg.add_argument('--num_workers', type=int, default=16)

# Other configurations
misc_arg = add_argument_group('Misc')
misc_arg.add_argument('--gpu_mode', type=str2bool, default=True)
misc_arg.add_argument('--verbose', type=str2bool, default=True)
misc_arg.add_argument('--pretrain', type=str, default='')
misc_arg.add_argument('--weights_fixed', type=str2bool, default=False)


def get_config():
    args = parser.parse_args()
    _, dataset = get_dataset_name(args.dataset)
    experiment_id = f"PointDSC_{dataset}_{time.strftime('%m%d%H%M')}"
    if args.snapshot_dir is None:
        args.snapshot_dir = f'snapshot/{experiment_id}'
    if args.tboard_dir is None:
        args.tboard_dir = f'tensorboard/{experiment_id}'
    if args.save_dir is None:
        args.save_dir = os.path.join(f'snapshot/{experiment_id}', 'models/')    
    if dataset == '3DMatch':
        if args.root is None:
            args.root='dummy(root)'        
        if args.descriptor is None:
            args.descriptor='fcgf'
        if args.inlier_threshold is None:
            args.inlier_threshold=0.10
        if args.sigma_d is None:
            args.sigma_d=0.10
        if args.downsample is None:
            args.downsample=0.03
        if args.re_thre is None:
            args.re_thre=15
        if args.te_thre is None:                                        
            args.te_thre=30
    else:
        if args.root is None:
            args.root=kitti_dir
        if args.descriptor is None:
            args.descriptor='fcgf'
        if args.inlier_threshold is None:
            args.inlier_threshold=1.2
        if args.sigma_d is None:
            args.sigma_d=1.2
        if args.downsample is None:
            args.downsample=0.30
        if args.re_thre is None:
            args.re_thre=5
        if args.te_thre is None:                        
            args.te_thre=60        

    return args
