# 相机参考工具
# @路怀帅 2025.12.19 10:32am 
# PS : 新增 加载相机畸变KD 视频对比显示功能

import cv2
import threading
import sys
from datetime import datetime
import os
import argparse
import numpy as np

# 全局变量
frame = None
running = True
which_camera = 0


# 加载鱼眼相机参数
K = np.array([[5.3402108990030604e+02, 0., 9.2598444295282172e+02], 
              [0., 5.3455325152709827e+02, 5.7771767919091610e+02],
              [0., 0., 1.]])
D = np.array([-1.8724887233075402e-02, 6.4408558584901701e-03,
              -5.2069636709412993e-03, 8.4815411645490968e-04])
W, H = 1920, 1080



def video_thread():
    global frame, running
    cap = cv2.VideoCapture(which_camera, cv2.CAP_ANY)
    cap.set(cv2.CAP_PROP_FOURCC,cv2.VideoWriter_fourcc(*"YUYV"))
    cap.set(cv2.CAP_PROP_FRAME_WIDTH, 1920)
    cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 1080)
    # cap.set(cv2.CAP_PROP_FPS, 30)

    if not cap.isOpened():
        print("[ERROR] Cannot open camera", file=sys.stderr)
        running = False
        return
    # 鱼眼相机去畸变
    map1, map2 = cv2.fisheye.initUndistortRectifyMap(K, D, np.eye(3), K, (W, H), cv2.CV_16SC2)

    while running:
        ret, f = cap.read()
        if not ret:
            break

        # 图像去畸变
        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 = comparison.astype(np.uint8) 

        # 设置视频流全屏显示
        text_info = f"Camera: {args.i.upper()}  |  Press 'q' to quit, 's' to screenshot"
        cv2.putText(comparison, text_info, (10, 30),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 255, 0), 2, cv2.LINE_AA)

        cv2.namedWindow('Video old vs new', cv2.WND_PROP_FULLSCREEN)
        cv2.setWindowProperty('Video old vs new', cv2.WND_PROP_FULLSCREEN, cv2.WINDOW_FULLSCREEN)
        cv2.imshow('Video old vs new', comparison)

        if cv2.waitKey(1) & 0xFF == ord('q'):
            running = False
            break

    cap.release()
    cv2.destroyAllWindows()

def input_thread():
    global running
    print("Commands: 's' = screenshot, 'q' = quit")
    while running:
        try:
            cmd = input().strip().lower()
            if cmd == 's':
                if frame is not None:
                    filename = f"./save_img/shot_{datetime.now().strftime('%Y%m%d_%H%M%S')}.jpg"
                    cv2.imwrite(filename, frame)
                    print(f"[SSH] Saved: {os.path.abspath(filename)}")
                else:
                    print("[SSH] No frame available yet.")
            elif cmd == 'q':
                running = False
                break
            else:
                print("[SSH] Unknown command. Use 's' or 'q'.")
        except EOFError:
            break

if __name__ == "__main__":

    # 获取用户参数启动
    parser = argparse.ArgumentParser(description="Example script with argparse")

    # 获取用户输入的摄像头方位 （front back left right）
    parser.add_argument("--i", type=str, required=True, help="摄像头方位")

    args = parser.parse_args()

    print("相机方位:", args.i)

    if args.i == "front":
        which_camera = 0
    elif args.i == "back":
        which_camera = 2
    elif args.i == "left":
        which_camera = 1
    elif args.i == "right":
        which_camera = 3
    else:
        print("[ERROR] Invalid camera direction. Use 'front', 'back', 'left', or 'right'.", file=sys.stderr)
        running = False
        exit(1)

    # 启动视频线程
    vt = threading.Thread(target=video_thread, daemon=True)
    vt.start()

    # 主线程监听 SSH 输入
    input_thread()

    print("[INFO] Exiting...")