import argparse
import logging


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


arg_lists = []
parser = argparse.ArgumentParser(
    description="Sub-Pixel Accurate Keypoint Refinement.", formatter_class=argparse.ArgumentDefaultsHelpFormatter
)
parser.add_argument("experiment", type=str)


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


common_arg = add_argument_group("Common")
common_arg.add_argument(
    "--fmat", "-fmat", action="store_true", help="estimate the fundamental matrix, instead of the essential matrix"
)

common_arg.add_argument(
    "--datasets", "-ds", default="megadepth", help="which datasets to use, separate multiple datasets by comma"
)

common_arg.add_argument(
    "--nfeatures",
    "-nf",
    type=int,
    default=2048,
    help="fixes number of features by clamping/replicating, set to -1 for dynamic feature count but then batchsize (-bs) has to be set to 1",
)

common_arg.add_argument(
    "--sideinfo",
    action="store_true",
    help="Do not provide side information (matching ratios) to the network. The network should be trained and tested consistently.",
)

common_arg.add_argument("--detect_anomaly", "-da", action="store_true", help="Anomaly detection of PyTorch")

common_arg.add_argument(
    "--ransac_thr", "-rt", type=float, default=1.0, help="GCRANSAC inlier threshold. Recommended value is 1.0px"
)

common_arg.add_argument(
    "--train_thr", "-tt", type=float, default=1.5, help="Train inlier threshold. Recommended value is 1.5px"
)

common_arg.add_argument(
    "--model", "-m", default="", help="load a model to contuinue training or leave empty to create a new model"
)

common_arg.add_argument(
    "--detector",
    "-detc",
    choices=["spnn", "splg", "aliked", "dedode", "xfeat"],
    default="sp",
    help='Type of detector; "sp" means SuperPoint, "aliked" means ALIKED',
)

# Training parameters
train_arg = add_argument_group("Training")
train_arg.add_argument("--resume", action="store_true", help="Resume from existing experiment")

train_arg.add_argument("--batchsize", "-bs", type=int, default=8, help="batch size")

train_arg.add_argument("--learning_rate", "-lr", type=float, default=0.0001, help="learning rate")

train_arg.add_argument("--weight_decay", "-wd", type=float, default=0.0, help="weight decay")

train_arg.add_argument("--train_iter", "-it", type=int, default=200000, help="number of iterations to train")

train_arg.add_argument("--save_intv", type=int, default=500, help="model saving interval")

train_arg.add_argument("--val_intv", type=int, default=5000, help="evaluation interval")

# Testing parameters
test_arg = add_argument_group("Testing")
test_arg.add_argument("--total_run", "-tr", type=int, default=10, help="Total number of runs on validation set")

test_arg.add_argument("--test", action="store_true", help="Testing mode")

test_arg.add_argument("--total_split", "-ts", type=int, default=1, help="total split size")

test_arg.add_argument("--current_split", "-cs", type=int, default=0, help="current split")

test_arg.add_argument("--vanilla", action="store_true", help="Run vanilla pose estimation pipeline (without K2S model)")


def get_config():
    config, unparsed = parser.parse_known_args()
    return config, unparsed


def print_usage():
    parser.print_usage()


def get_logger(args):
    formatter = logging.Formatter(fmt="[%(asctime)s %(name)s %(levelname)s] %(message)s", datefmt="%m/%d/%Y %H:%M:%S")
    handler = logging.StreamHandler()
    handler.setFormatter(formatter)
    handler.setLevel(logging.INFO)

    logger = logging.getLogger(f"K2S")
    logger.setLevel(logging.INFO)
    logger.addHandler(handler)
    logger.propagate = False

    __module_name__ = f"K2S"
    return logger
