import argparse

from superpoint_pruning.evaluation import eval as eval_mod
from superpoint_pruning import export as export_mod
from superpoint_pruning.distillation import setup as setup_mod
from superpoint_pruning.distillation import lightning_trainer as train_mod
from superpoint_pruning.evaluation import benchmark_sp_tensorrt as benchmark_mod
from superpoint_pruning.evaluation import plot_keypoints as plot_keypoints_mod


def main():
    parser = argparse.ArgumentParser(prog="superpoint-pruning")
    subparsers = parser.add_subparsers(dest="command", required=True)

    eval_parser = subparsers.add_parser("evaluate", help="Run SuperPoint evaluation")
    eval_mod.add_parser_args(eval_parser)
    eval_parser.set_defaults(func=eval_mod.main)

    export_parser = subparsers.add_parser("export", help="Export SuperPoint to ONNX")
    export_mod.add_parser_args(export_parser)
    export_parser.set_defaults(func=export_mod.main)

    setup_parser = subparsers.add_parser(
        "setup", help="Setup the dataset for distillation and evaluation"
    )
    setup_mod.add_parser_args(setup_parser)
    setup_parser.set_defaults(func=setup_mod.main)

    train_parser = subparsers.add_parser("train", help="Train a pruned SuperPoint")
    train_mod.add_parser_args(train_parser)
    train_parser.set_defaults(func=train_mod.main)

    benchmark_parser = subparsers.add_parser(
        "benchmark", help="Benchmark a TensorRT SuperPoint engine"
    )
    benchmark_mod.add_parser_args(benchmark_parser)
    benchmark_parser.set_defaults(func=benchmark_mod.main)

    plot_parser = subparsers.add_parser(
        "plot", help="Plot original and pruned SuperPoint keypoints"
    )
    plot_keypoints_mod.add_parser_args(plot_parser)
    plot_parser.set_defaults(func=plot_keypoints_mod.main)

    args = parser.parse_args()
    args.func(args)