import os
import yaml
import argparse

from mmengine.runner import Runner
from mmengine.config import Config

# from data.get_partial_data import create_train_val_dataloaders_config_from_data_path_list
import torch

def argparse_setup():
    parser = argparse.ArgumentParser(description='Train a model')
    parser.add_argument('--config', required=True, type=str, default='config.yaml', help='Path to config file')
    parser.add_argument('--local_rank', '--local-rank', type=int, default=0)
    parser.add_argument("--checkpoint", type=str, required=True, help="Checkpoint file to load")
    parser.add_argument(
        '--launcher',
        choices=['none', 'pytorch', 'slurm', 'mpi'],
        default='none',
        help='job launcher')
    parser.add_argument(
        "--project_dirs",
        type=str,
        default="missing_person_obj_det",
        help="project directories for saving the model and log files")
    parser.add_argument(
        '--exp_name',
        default='',
        type=str,
        help='experiment name, used for saving the model and log files')
    args = parser.parse_args()
    if 'LOCAL_RANK' not in os.environ:
        os.environ['LOCAL_RANK'] = str(args.local_rank)
    return args



# config_names = [
#     # "exp_config/yolox.yaml",
#     # "exp_config/yolov3.yaml",
#     # "exp_config/faster_rcnn.yaml",
#     # "exp_config/retinanet.yaml",
#     "exp_config/ssd.yaml",
#     "exp_config/detr.yaml",
# ]

# checkpoints = [
#     # "/mnt/home/jeongjun/neurips2025/mmdet/missing_person/250503_missing_person_yolox_forestperson_v3/best_coco_bbox_mAP_epoch_17.pth",
#     # "/mnt/home/jeongjun/neurips2025/mmdet/missing_person/250503_missing_person_yolov3_forestperson_v3/best_coco_bbox_mAP_epoch_70.pth",
#     # "/mnt/home/jeongjun/neurips2025/mmdet/missing_person/250503_missing_person_faster_rcnn_forestperson_v3/best_coco_bbox_mAP_epoch_9.pth",
#     # "/mnt/home/jeongjun/neurips2025/mmdet/missing_person/250503_missing_person_retinanet_forestperson_v3/best_coco_bbox_mAP_epoch_9.pth",
#     # "/mnt/home/jeongjun/neurips2025/mmdet/missing_person/250503_missing_person_ssd_forestperson_v3/best_coco_bbox_mAP_epoch_100.pth"
#     "/mnt/home2/etri/ETRI_checkpoints/250503_missing_person_ssd_forestperson_v3/best_coco_bbox_mAP_epoch_100.pth",
#     "/mnt/home2/etri/model_checkpoints/benchmark/250508_missing_person_detr_forestperson_v3/best_coco_bbox_mAP_epoch_126.pth",
# ]


# config_names = [
#     "exp_config/yolov3.yaml",
#     "exp_config/yolox.yaml",
#     "exp_config/faster_rcnn.yaml",
#     "exp_config/retinanet.yaml",
#     "exp_config/ssd.yaml",
#     "exp_config/detr.yaml",
# ]

# checkpoints = [
#     "/mnt/home2/etri/model_checkpoints/standing/250508_missing_person_yolov3_forestperson_v3_standing/best_coco_bbox_mAP_epoch_252.pth",
#     "/mnt/home/jeongjun/neurips2025/mmdet/missing_person/250508_missing_person_yolox_forestperson_v3_standing_all/best_coco_bbox_mAP_epoch_102.pth",
#     "/mnt/home2/etri/model_checkpoints/standing/250508_missing_person_faster_rcnn_forestperson_v3_standing/best_coco_bbox_mAP_epoch_10.pth",
#     "/mnt/home2/etri/model_checkpoints/standing/250508_missing_person_retinanet_forestperson_v3_standing/best_coco_bbox_mAP_epoch_8.pth",
#     "/mnt/home2/etri/model_checkpoints/standing/250509_missing_person_ssd_forestperson_v3_standing/best_coco_bbox_mAP_epoch_100.pth",
#     "/mnt/home2/etri/model_checkpoints/standing/250510_missing_person_detr_forestperson_v3_standing/best_coco_bbox_mAP_epoch_122.pth"
# ]

# config_names = [
#     "exp_config/faster_rcnn.yaml",
#     "exp_config/faster_rcnn.yaml",
#     "exp_config/faster_rcnn.yaml",
# ]

# checkpoints = [
#     # "/mnt/home/jeongjun/neurips2025/mmdet/missing_person/250501_missing_person_faster_rcnn_annotation_wisardz_90k/best_coco_bbox_mAP_iter_60000.pth",
#     # "/mnt/home/jeongjun/neurips2025/mmdet/missing_person/250509_missing_person_faster_rcnn_annotation_crowdhuman_90k/best_coco_bbox_mAP_iter_20000.pth"
#     # "/mnt/home/jeongjun/neurips2025/mmdet/missing_person/250509_missing_person_faster_rcnn_annotation_cocoperson_90k/best_coco_bbox_mAP_iter_70000.pth",
#     "/mnt/home/jeongjun/neurips2025/mmdet/missing_person/250501_missing_person_faster_rcnn_annotation_citypersons_90k/best_coco_bbox_mAP_iter_30000.pth",
#     "/mnt/home/jeongjun/neurips2025/mmdet/missing_person/250501_missing_person_faster_rcnn_annotation_heridal_90k/best_coco_bbox_mAP_iter_10000.pth",
#     "/mnt/home/jeongjun/neurips2025/mmdet/missing_person/250501_missing_person_faster_rcnn_annotation_sard_90k/best_coco_bbox_mAP_iter_50000.pth"
# ]


# config_names = [
#     "exp_config/yolov3.yaml",
#     "exp_config/yolox.yaml",
#     "exp_config/faster_rcnn.yaml",
#     "exp_config/retinanet.yaml",
#     "exp_config/ssd.yaml",
#     "exp_config/detr.yaml",

#     "exp_config/yolov3.yaml",
#     "exp_config/yolox.yaml",
#     "exp_config/faster_rcnn.yaml",
#     "exp_config/retinanet.yaml",
#     "exp_config/ssd.yaml",
#     "exp_config/detr.yaml",
# ]

# checkpoints = [
#     "/mnt/home/jeongjun/neurips2025/mmdet/missing_person/250508_missing_person_summer_yolov3_server4_annotation_v0.3/best_coco_bbox_mAP_epoch_126.pth",
#     "/mnt/home/jeongjun/neurips2025/mmdet/missing_person/250508_missing_person_summer_yolox_annotation_v0.3/best_coco_bbox_mAP_epoch_226.pth",
#     "/mnt/home/jeongjun/neurips2025/mmdet/missing_person/250508_missing_person_summer_faster_rcnn_annotation_v0.3/best_coco_bbox_mAP_epoch_12.pth",
#     "/mnt/home/jeongjun/neurips2025/mmdet/missing_person/250510_missing_person_summer_retinanet_annotation_v0.3/best_coco_bbox_mAP_epoch_9.pth",
#     "/mnt/home/jeongjun/neurips2025/mmdet/missing_person/250508_missing_person_summer_ssd_annotation_v0.3/best_coco_bbox_mAP_epoch_100.pth",
#     "/mnt/home/jeongjun/neurips2025/mmdet/missing_person/250510_missing_person_summer_detr_annotation_v0.3/best_coco_bbox_mAP_epoch_105.pth",


#     "/mnt/home/jeongjun/neurips2025/mmdet/missing_person/250508_missing_person_winter_yolov3_server4_annotation_v0.3/best_coco_bbox_mAP_epoch_84.pth",
#     "/mnt/home/jeongjun/neurips2025/mmdet/missing_person/250508_missing_person_winter_yolox_server8_annotation_v0.3/best_coco_bbox_mAP_epoch_61.pth",
#     "/mnt/home/jeongjun/neurips2025/mmdet/missing_person/250508_missing_person_winter_faster_rcnn_annotation_v0.3/best_coco_bbox_mAP_epoch_9.pth",
#     "/mnt/home/jeongjun/neurips2025/mmdet/missing_person/250508_missing_person_winter_retinanet_server8_annotation_v0.3/best_coco_bbox_mAP_epoch_9.pth",
#     "/mnt/home/jeongjun/neurips2025/mmdet/missing_person/250508_missing_person_winter_ssd_annotation_v0.3/best_coco_bbox_mAP_epoch_105.pth",
#     "/mnt/home/jeongjun/neurips2025/mmdet/missing_person/250510_missing_person_winter_detr_annotation_v0.3/best_coco_bbox_mAP_epoch_75.pth"
# ]

config_names = [
    "exp_config/retinanet.yaml",
    "exp_config/faster_rcnn.yaml",
    "exp_config/ssd.yaml",
    "exp_config/yolov3.yaml",
    "exp_config/yolox.yaml",
]

checkpoints = [
    "missing_person/250423_missing_person_retinanet_annotation_v0.3/best_coco_bbox_mAP_epoch_12.pth",
    "missing_person/250425_missing_person_faster_rcnn_annotation_v0.3/best_coco_bbox_mAP_epoch_9.pth",
    "missing_person/250503_missing_person_ssd_forestperson_v3/epoch_2.pth",
    "missing_person/250425_missing_person_yolov3_annotation_v0.3/epoch_49_best.pth",
    "missing_person/250422_missing_person_yolox_annotation_v0.3/best_coco_bbox_mAP_epoch_74.pth"
]


# def main():
#     args = argparse_setup()
#     with open(args.config, 'r') as file:
#         config = yaml.safe_load(file)
#     print(config)
    
#     # 1. setup the dataloader
#     base_config = Config.fromfile(config["base_config"])
#     base_config["work_dir"] = os.path.join(args.project_dirs, args.exp_name)
    
#     base_config.launcher = args.launcher
#     base_config.load_from = args.checkpoint

#     base_config.test_dataloader.batch_size = 1
    

#     # Get experiment name
#     # base_config["experiment_name"] = config["experiment_name"]
#     base_config["experiment_name"] = args.exp_name

#     base_config["default_hooks"] = dict()

#     base_config.test_evaluator.update(
#         dict(outfile_prefix=args.checkpoint.split("/")[-2])
#     )

#     runner = Runner.from_cfg(base_config)
    
#     # 3. train wiht runner!
#     # runner.val()
#     runner.test()

#     runner.val()

def main():
    args = argparse_setup()
    dicts = {}
    for config_name, checkpoint in zip(config_names, checkpoints):
        args.config = config_name
        args.checkpoint = checkpoint

        print(f"Running with config: {config_name} and checkpoint: {checkpoint}")
        with open(args.config, 'r') as file:
            config = yaml.safe_load(file)
        print(config)
        
        # 1. setup the dataloader
        base_config = Config.fromfile(config["base_config"])
        base_config["work_dir"] = os.path.join(args.project_dirs, args.exp_name)
        
        base_config.launcher = args.launcher
        base_config.load_from = args.checkpoint

        base_config.test_dataloader.batch_size = 1
        

        # Get experiment name
        # base_config["experiment_name"] = config["experiment_name"]
        base_config["experiment_name"] = args.exp_name

        base_config["default_hooks"] = dict()

        season_prediction_result = "season_prediction_results2"
        outfile_prefix = os.path.join(
            season_prediction_result, checkpoint.split("/")[-2]
        )
        base_config.test_evaluator.update(
            dict(outfile_prefix=outfile_prefix)
        )

        runner = Runner.from_cfg(base_config)
        
        # 3. train wiht runner!
        # val = runner.val()
        # test = runner.test()
        # dicts[config_name] = {
        #     "val": val,
        #     "test": test
        # }
        # print(f"Finished running with config: {config_name} and checkpoint: {checkpoint}")
        # print(f"Validation results: {val}")
        # print(f"Test results: {test}")
                # val = runner.val()
        test = runner.test()
        dicts[config_name] = {
            "test": test
        }
        print(f"Finished running with config: {config_name} and checkpoint: {checkpoint}")
        print(f"Test results: {test}")
    print(dicts)
    # Save the results to a file
    with open("overall_results.yaml", 'w') as file:
        yaml.dump(dicts, file)


    
    
if __name__=="__main__":
    
    # This area is for unit test
    
    # 1. load the model config file
    # tmp = Config.fromfile("configs/models/retinanet.py")
    # tmp = Config.fromfile("configs/models/faster_rcnn.py")
    # breakpoint()
    
    # 2. load the concatenated dataset config
    # dummy_dir_path_list = [
    #     "/mnt/home/jeongjun/layout_diffusion/datasets/250130_ms_coco_seed_3/w_acq/guided/cycle_1/3", 
    #     "/mnt/home/jeongjun/layout_diffusion/datasets/250130_ms_coco_seed_3/w_acq/guided/cycle_2/3", 
    # ]
    # dataloader_tmp = create_dataloader_dict_from_data_path_list(dummy_dir_path_list)
    # breakpoint()
    
    # check actual running
    main()