import argparse
import os
import cv2
from tqdm import tqdm
from scripts.data_converter.visual_utils import *

import warnings
warnings.filterwarnings("ignore")

def kitti_visual_tool(data_root, pred_root, demo_dir):
    if not os.path.exists(data_root):
        raise ValueError("data_root Not Found")
    image_path = os.path.join(data_root, "training/image_2")
    calib_path = os.path.join(data_root, "training/calib")
    gt_path = os.path.join(data_root, "training/label_2")
    label_path = pred_root
    image_ids = []
    label_files = os.listdir(label_path)
    # for image_file in os.listdir(image_path):
    #     image_ids.append(image_file.split(".")[0])
    # for i in range(len(image_ids)):
    #     if os.path.exists(os.path.join(image_path, str(image_ids[i]) + ".png")):
    #         image_2_file = os.path.join(image_path, str(image_ids[i]) + ".png")
    #     elif os.path.exists(os.path.join(image_path, str(image_ids[i]) + ".jpg")):
    #         image_2_file = os.path.join(image_path, str(image_ids[i]) + ".jpg")
    #     else:
    #         print("Error: image file not found.")
    #     calib_file = os.path.join(calib_path, str(image_ids[i]) + ".txt")
    #     label_2_file = os.path.join(label_path, str(image_ids[i]) + ".txt")
    #     image = cv2.imread(image_2_file)
    #     _, P2, denorm = load_calib(calib_file)
    #     image = draw_3d_box_on_image(image, label_2_file, P2, denorm)
    #     cv2.imwrite(os.path.join(demo_dir, str(image_ids[i]) + ".jpg"), image)
    # import pdb;pdb.set_trace()
    for index, label_file in enumerate(tqdm(label_files)):
        
        image_2_file = os.path.join(image_path, label_file[:-4] + ".jpg")
        calib_file = os.path.join(calib_path, label_file[:-4] + ".txt")
        gt = os.path.join(gt_path, label_file[:-4] + ".txt")
        label = os.path.join(label_path, label_file)
        image = cv2.imread(image_2_file)
        _, P2, denorm = load_calib(calib_file)
        image = draw_3d_box_on_image(image, gt, P2, denorm, (0,255,0))      # green
        image = draw_3d_box_on_image(image, label, P2, denorm, (255,0,0))   # blue
        
        cv2.imwrite(os.path.join(demo_dir, label_file[:-4] + ".jpg"), image)

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Dataset in KITTI format Checking ...")
    parser.add_argument("--data_root", type=str,
                        default="",
                        help="Path to Dataset root in KITTI format")
    parser.add_argument("--pred_root", type=str,
                        default="/home/ps/zz2/BEVHeight/outputs_dair/data",
                        help="Path to Dataset root in KITTI format")
    parser.add_argument("--demo_dir", type=str,
                        default="",
                        help="Path to demo directions")
    args = parser.parse_args()
    os.makedirs(args.demo_dir, exist_ok=True)
    kitti_visual_tool(args.data_root, args.pred_root, args.demo_dir)