import argparse
from datetime import datetime
import logging
import pprint
import os
import torch
from src.config import load_config, init_logging
from src.Q_SLAM import Q_SLAM
from src.common import setup_seed



def main():
    parser = argparse.ArgumentParser(
        description='Arguments for running the Point-SLAM.'
    )
    parser.add_argument('config', type=str, help='Path to config file.')
    parser.add_argument('--input_folder', type=str,
                        help='input folder, this have higher priority, can overwrite the one in config file')
    parser.add_argument('--output', type=str,
                        help='output folder, this have higher priority, can overwrite the one in config file')
    parser.add_argument('--wandb', action='store_true')
    parser.add_argument('--no_wandb', action='store_true')

    def optional_int(string):
        return None if string == "None" else int(string)
    parser.add_argument("--stop", type=optional_int,
                        help="stop after n frames")

    args = parser.parse_args()

    cfg = load_config(args.config, 'configs/q_slam.yaml')

    setup_seed(cfg["setup_seed"])

    if args.stop:
        cfg["mapping"]["ckpt_freq"] = args.stop
        cfg["mapping"]["keyframe_every"] = 10

    # torch.cuda.set_device("cuda")
    # cuda debugging
    # os.environ['CUDA_LAUNCH_BLOCKING'] = "1"
    now = datetime.now()
    time_string = now.strftime("%Y%m%d_%H%M%S") if args.stop is None else None
    print("🚀 Start at: {}".format(now.strftime("%Y-%m-%d %H:%M:%S")))
    print("⭐️ PID: {}".format(os.getpid())) 
    print("cfgs: {}".format(pprint.pformat(cfg)))
    init_logging()
    slam = Q_SLAM(cfg, args, time_string=time_string)
    slam.run()
    print("🎉 Done! {}".format(datetime.now().strftime("%Y-%m-%d %H:%M:%S")))

if __name__ == '__main__':
    main()
