# from turtle import right
import cv2
import threading
import sys
import os
import argparse
import numpy as np
from surround_view import FisheyeCameraModel, BirdView
import surround_view.param_settings as settings

right_frame = None

class MultiCameraBirdView:
    def __init__(self):
        self.running = True
        self.names = settings.camera_names  # e.g., ['front', 'back', 'left', 'right']
        self.yamls = [os.path.join(os.getcwd(), "yaml", name + ".yaml") for name in self.names]
        self.camera_models = [
            FisheyeCameraModel(camera_file, camera_name)
            for camera_file, camera_name in zip(self.yamls, self.names)
        ]
        
        # 摄像头设备映射
        self.which_cameras = {"front": 0, "back": 2, "left": 1, "right": 3}
        self.caps = []

        print("[INFO] 初始化摄像头...")
        for name in self.names:
            cap = cv2.VideoCapture(self.which_cameras[name], 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_BUFFERSIZE, 1)  # 最小缓冲，降低延迟
            if not cap.isOpened():
                print(f"[ERROR] 无法打开 {name} 摄像头 (设备 {self.which_cameras[name]})", file=sys.stderr)
                self.running = False
                return
            self.caps.append(cap)
        
        self.birdview = BirdView()
        self._initialize_weights()

    def _initialize_weights(self):
        try:
            images = [os.path.join(os.getcwd(), "images", name + ".png") for name in self.names]
            static_frames = []
            for img_path, cam_model in zip(images, self.camera_models):
                img = cv2.imread(img_path)
                if img is None:
                    print(f"[WARN] 静态图缺失: {img_path}，用黑图代替")
                    img = np.zeros((1080, 1920, 3), dtype=np.uint8)
                img = cam_model.undistort(img)
                img = cam_model.project(img)
                img = cam_model.flip(img)
                static_frames.append(img)
            if len(static_frames) == 4:
                self.birdview.get_weights_and_masks(static_frames)
                print("[INFO] 权重矩阵初始化成功")
        except Exception as e:
            print(f"[ERROR] 权重初始化失败: {e}", file=sys.stderr)

    def process_frame_once(self, frame, model):
        """只处理一次：去畸变 + 投影 + 翻转"""
        frame = model.undistort(frame)
        frame = model.project(frame)
        frame = model.flip(frame)
        return frame

    def process_frame_undistort(self, frame, model):
        """只处理一次：去畸变 + 投影 + 翻转"""
        frame = model.undistort(frame)
        # frame = model.project(frame)
        # frame = model.flip(frame)
        return frame

    def run(self):
        current_view = "front"  # 默认显示前视图

        while self.running:
            raw_frames = {}
            processed_frames = []
            valid = True

            # 读取并处理所有摄像头帧
            for i, (cap, model, name) in enumerate(zip(self.caps, self.camera_models, self.names)):
                ret, frame = cap.read()
                if not ret or frame is None:
                    print(f"[WARN] 跳过 {name} 帧")
                    valid = False
                    break
                # 保存原始帧（用于右侧显示）
                raw_frames[name] = frame.copy()
                # 处理用于鸟瞰图的帧
                p_frame = self.process_frame_once(frame, model)
                processed_frames.append(p_frame)

            if not valid or len(processed_frames) != 4:
                continue

            # 更新鸟瞰图
            self.birdview.update_frames(processed_frames)
            self.birdview.stitch_all_parts()
            self.birdview.make_white_balance()
            self.birdview.copy_car_image()

            # 获取当前选中的单路图像（已校正）
            single_img = self.process_frame_undistort(raw_frames[current_view], self.camera_models[self.names.index(current_view)])
            birdview_img = self.birdview.image

            # 拼接显示：左侧鸟瞰图（1/3），右侧单路（2/3）
            h_display, w_display = 720, 1280
            w_bird = w_display // 3
            w_single = w_display - w_bird

            bird_resized = cv2.resize(birdview_img, (w_bird, h_display))
            single_resized = cv2.resize(single_img, (w_single, h_display))
            display = np.hstack((bird_resized, single_resized))



            cv2.namedWindow('Video', cv2.WND_PROP_FULLSCREEN)
            cv2.setWindowProperty('Video', cv2.WND_PROP_FULLSCREEN, cv2.WINDOW_FULLSCREEN)
            cv2.imshow("Video", display)

            key = cv2.waitKey(1) & 0xFF
            if key == ord('q'):
                self.running = False
                break
            elif key == ord('1'):
                current_view = "front"
            elif key == ord('2'):
                current_view = "back"
            elif key == ord('3'):
                current_view = "left"
            elif key == ord('4'):
                current_view = "right"

        # 释放资源
        for cap in self.caps:
            cap.release()
        cv2.destroyAllWindows()


def main():
    print("🚀 启动实时四路环视系统...")
    print("操作说明:")
    print("  1 - 前视图 | 2 - 后视图 | 3 - 左视图 | 4 - 右视图")
    print("  q - 退出程序")
    multi_cam = MultiCameraBirdView()
    if multi_cam.running:
        multi_cam.run()
    else:
        print("[ERROR] 摄像头初始化失败")

if __name__ == "__main__":
    main()