import cv2
import threading
import sys
import os
import argparse
import numpy as np
import time  


running = True
capture_flag = False

# 内参 (K) 与 畸变系数 (D)
K = np.array([[ 3.9875655282072148e+02, 0., 6.3788553583130647e+02],
              [0.,3.7214138046449295e+02, 3.4915903287538981e+02],
              [  0., 0., 1. ]])
D = np.array([ 8.6662165805452579e-02, -2.8155159006680735e-02,
               -3.5111914522638211e-02, 1.4621518110074484e-02 ])

W, H = 1280, 720


def input_thread():

    global running, capture_flag
    print("\n终端输入's'截图,'q'退出程序。")
    while running:
        try:
            cmd = input().strip().lower()
            if cmd == 's':
                capture_flag = True
            elif cmd == 'q':
                print("退出命令触发")
                running = False
                break
        except EOFError:
            break
        except Exception:
            pass


def main():
    global running, capture_flag


    parser = argparse.ArgumentParser(description="Multi-Camera Hardware Reference Tool")
    parser.add_argument("--i", type=str, required=True, 
                        choices=["front", "back", "left", "right"], 
                        help="摄像头方位 (front, back, left, right)")
    args = parser.parse_args()
    
    camera_mapping = {"front": 0, "left": 1, "back": 2, "right": 3}
    which_camera = camera_mapping[args.i]
    cam_name = args.i.upper()
    print(f"无法打开摄像头节点: {cam_name} (Node: /dev/video{which_camera})")


    save_dir = os.path.join(os.getcwd(), "images")
    if not os.path.exists(save_dir):
        os.makedirs(save_dir)

  
    cmd_listener = threading.Thread(target=input_thread, daemon=True)
    cmd_listener.start()


    cap = None
    try:
        cap = cv2.VideoCapture(which_camera, cv2.CAP_V4L2)
        cap.set(cv2.CAP_PROP_FOURCC, cv2.VideoWriter_fourcc(*"YUYV"))
        cap.set(cv2.CAP_PROP_FRAME_WIDTH, W)
        cap.set(cv2.CAP_PROP_FRAME_HEIGHT, H)
        cap.set(cv2.CAP_PROP_BUFFERSIZE, 3) 

        if not cap.isOpened():
            print(f"[ERROR] Cannot open camera /dev/video{which_camera}", file=sys.stderr)
            running = False
            return

     
        map1, map2 = cv2.fisheye.initUndistortRectifyMap(K, D, np.eye(3), K, (W, H), cv2.CV_16SC2)   # 预计算鱼眼镜头的抗畸变映射矩阵

        cv2.namedWindow('PowLu camera tools', cv2.WND_PROP_FULLSCREEN)
        cv2.setWindowProperty('Video old vs new', cv2.WND_PROP_FULLSCREEN, cv2.WINDOW_FULLSCREEN)

        while running:
            ret, f = cap.read()
            if not ret or f is None:
                time.sleep(0.01)
                continue

            raw_frame = f.copy()

            if capture_flag:
                filename = os.path.join(save_dir, f"{args.i.lower()}.png")
                cv2.imwrite(filename, raw_frame)
                print(f"[SUCCESS] Image saved accurately: {filename}")
                capture_flag = False

            undistorted = cv2.remap(f, map1, map2, interpolation=cv2.INTER_LINEAR, borderMode=cv2.BORDER_CONSTANT)
            comparison = np.hstack((f, undistorted))

            scale = min(1920 / comparison.shape[1], 1080 / comparison.shape[0])
            if scale < 1:
                comparison = cv2.resize(comparison, (int(comparison.shape[1] * scale), int(comparison.shape[0] * scale)))
            
            comparison = np.ascontiguousarray(comparison, dtype=np.uint8)

            text_info = f"Camera: {cam_name}  |  Press 'q'='Quit', 's'='Screenshot'"
            cv2.putText(comparison, text_info, (20, 50),
                        cv2.FONT_HERSHEY_SIMPLEX, 1.2, (0, 255, 255), 3, cv2.LINE_AA)
            
            w_half = comparison.shape[1] // 2
            h_bottom = comparison.shape[0] - 30
            cv2.putText(comparison, "FISHEYE", (20, h_bottom), 
                        cv2.FONT_HERSHEY_SIMPLEX, 1.5, (0, 255, 0), 3)
            cv2.putText(comparison, "UNDISTORTED", (w_half + 20, h_bottom), 
                        cv2.FONT_HERSHEY_SIMPLEX, 1.5, (0, 255, 0), 3)

            cv2.imshow('PowLu camera tools', comparison)

            key = cv2.waitKey(1) & 0xFF
            if key == ord('q'):
                running = False
                break
            elif key == ord('s'):
                capture_flag = True

    finally:
        running = False
        if cap is not None and cap.isOpened():
            cap.release()
        cv2.destroyAllWindows()

if __name__ == "__main__":
    try:
        main()
    except KeyboardInterrupt:
        running = False