# lib/base_img_process.py
# -*- coding: utf-8 -*-
"""
图像处理基础模块
- 图像去畸变
- 相机内参调整
- 通用图像变换
"""

import cv2
import numpy as np
from typing import Optional, Tuple, Dict

from lidar_manager.logger import get_module_logger

logger = get_module_logger(__name__)
_warned_legacy = False


class ImageProcessor:
    """图像处理器基类（支持针孔和鱼眼相机）"""

    def __init__(self, camera_intrinsic: np.ndarray,
                 distortion_coeffs: Optional[np.ndarray] = None,
                 image_size: Optional[Tuple[int, int]] = None,
                 undistort_scale: float = 1.0,
                 camera_model: str = 'pinhole',
                 camera_id: Optional[str] = None):
        """
        初始化图像处理器

        参数:
            camera_intrinsic: 3x3相机内参矩阵 K
            distortion_coeffs: 畸变系数 (可选)
            image_size: 图像尺寸 (width, height) (可选)
            undistort_scale: 去畸变scale参数，统一使用此值 (默认1.0，保留完整图像)
            camera_model: 相机模型类型 ('pinhole' 或 'fisheye')
                - 'pinhole': 针孔相机模型（使用Brown-Conrady畸变模型）
                - 'fisheye': 鱼眼相机模型（使用Kannala-Brandt畸变模型）
            camera_id: 相机标识（如 frame_id），用于日志定位，可选
        """
        self.K = camera_intrinsic.copy()
        self.D = distortion_coeffs.copy() if distortion_coeffs is not None else None
        self.image_size = image_size
        self.camera_id = camera_id or ""

        # 🔧 统一去畸变scale参数
        self._undistort_scale = undistort_scale

        # 🔧 相机模型类型
        if camera_model not in ['pinhole', 'fisheye']:
            logger.warning(f"[ImageProcessor] 未知的相机模型 '{camera_model}'，使用默认值 'pinhole'")
            camera_model = 'pinhole'
        self.camera_model = camera_model
        logger.info(f"[ImageProcessor] 相机模型: {camera_model}" + (f", 相机={camera_id}" if camera_id else ""))

        # 缓存去畸变映射
        self._undistort_map1 = None
        self._undistort_map2 = None
        self._new_K = None
        
        # 有效ROI区域（去畸变后的有效区域）
        self._valid_roi = None  # (x, y, width, height)

    def undistort_image(self, image: np.ndarray,
                        scale: Optional[float] = None,
                        use_cached_map: bool = True) -> Tuple[np.ndarray, np.ndarray]:
        """
        去畸变图像（统一接口）
        
        参数:
            image: 输入图像
            scale: 去畸变后保留的有效区域比例 (0-1)
                  如果为None，使用初始化时设置的统一scale参数
                  如果指定，会覆盖统一参数（不推荐，建议通过配置统一设置）
            use_cached_map: 是否使用缓存的映射表（加速）
        
        返回:
            (undistorted_image, new_K): 去畸变后的图像和新的内参矩阵
        
        注意：
            - 🔧 统一接口：默认使用初始化时设置的scale参数，确保所有去畸变使用相同方法
            - scale=1.0: 保留完整的去畸变图像（不裁剪，推荐）
            - scale=0.0: 裁剪成方形，只保留完全有效的区域
        """
        # 🔧 统一使用初始化时设置的scale参数
        if scale is None:
            scale = self._undistort_scale
        else:
            if abs(scale - self._undistort_scale) > 1e-6:
                logger.warning(f"[图像去畸变] ⚠️ 传入的scale({scale})与统一参数({self._undistort_scale})不一致，"
                              f"建议统一使用配置中的scale参数")
        
        if self.D is None or len(self.D) == 0:
            # 无畸变系数，直接返回
            return image.copy(), self.K.copy()
        
        h, w = image.shape[:2]
        
        # 🔧 根据相机模型选择去畸变方法
        if self.camera_model == 'fisheye':
            return self._undistort_fisheye(image, scale, use_cached_map)
        else:
            return self._undistort_pinhole(image, scale, use_cached_map)
    
    def _undistort_pinhole(self, image: np.ndarray, scale: float, use_cached_map: bool) -> Tuple[np.ndarray, np.ndarray]:
        """
        针孔相机去畸变（Brown-Conrady模型）
        
        使用 cv2.getOptimalNewCameraMatrix 和 cv2.initUndistortRectifyMap
        """
        h, w = image.shape[:2]
        
        # 如果没有缓存或图像尺寸变化，重新计算
        if (not use_cached_map or
                self._undistort_map1 is None or
                self._undistort_map2 is None or
                self.image_size != (w, h)):
            
            # 仅首次或参数变化时输出详细日志，避免每次渲染刷屏
            log_params = not getattr(self, '_undistort_pinhole_params_logged', False)
            if log_params:
                logger.debug(f"[针孔相机去畸变] 原始内参矩阵K:\n{self.K}")
                logger.debug(f"[针孔相机去畸变] 畸变系数D: {self.D} (长度: {len(self.D)})")
            
            # 🔧 详细输出参数含义
            if len(self.D) >= 5:
                if log_params:
                    logger.debug(f"[针孔相机去畸变] 基础5参数: k1={self.D[0]:.6f}, k2={self.D[1]:.6f}, "
                               f"p1={self.D[2]:.6f}, p2={self.D[3]:.6f}, k3={self.D[4]:.6f}")
            if len(self.D) >= 8:
                if log_params:
                    logger.debug(f"[针孔相机去畸变] 扩展8参数: k4={self.D[5]:.6f}, "
                               f"k5={self.D[6]:.6f}, k6={self.D[7]:.6f}")
                logger.info(f"[针孔相机去畸变] ✅ 使用8参数Brown-Conrady模型（全部参数）" + (f", 相机={self.camera_id}" if self.camera_id else ""))
            elif len(self.D) == 5:
                logger.info(f"[针孔相机去畸变] ✅ 使用5参数Brown-Conrady模型" + (f", 相机={self.camera_id}" if self.camera_id else ""))
            else:
                camera_desc = f", 相机={self.camera_id}" if self.camera_id else ""
                logger.warning(
                    f"[针孔相机去畸变] ⚠️ 畸变参数数量异常: {len(self.D)} (期望5或8){camera_desc}。"
                    f" 若该相机为鱼眼，请在标定中设置 camera_model='fisheye'。"
                    f" 参见文档: camera_model 与 distortion 参数说明。"
                )
            
            # 计算最优新内参矩阵
            self._new_K, roi = cv2.getOptimalNewCameraMatrix(
                self.K, self.D, (w, h), scale
            )
            
            x, y, roi_w, roi_h = roi
            if log_params:
                logger.debug(f"[针孔相机去畸变] 新内参矩阵new_K:\n{self._new_K}")
                logger.debug(f"[针孔相机去畸变] 有效ROI区域: x={x}, y={y}, w={roi_w}, h={roi_h}")
                logger.debug(f"[针孔相机去畸变] 注意: ROI尺寸({roi_w}x{roi_h})可能小于原始尺寸({w}x{h})")
            self._undistort_pinhole_params_logged = True
            
            # 保存有效ROI区域
            self._valid_roi = (x, y, roi_w, roi_h)
            
            # 生成去畸变映射
            self._undistort_map1, self._undistort_map2 = cv2.initUndistortRectifyMap(
                self.K, self.D, None, self._new_K, (w, h), cv2.CV_32FC1
            )
            
            self.image_size = (w, h)
        
        # 应用去畸变
        undistorted = cv2.remap(
            image,
            self._undistort_map1,
            self._undistort_map2,
            cv2.INTER_LINEAR
        )
        
        out_h, out_w = undistorted.shape[:2]
        
        return undistorted, self._new_K.copy()
    
    def _undistort_fisheye(self, image: np.ndarray, scale: float, use_cached_map: bool) -> Tuple[np.ndarray, np.ndarray]:
        """
        鱼眼相机去畸变（Kannala-Brandt模型）
        
        使用 cv2.fisheye.getOptimalNewCameraMatrix 和 cv2.fisheye.initUndistortRectifyMap
        """
        h, w = image.shape[:2]
        
        # 检查畸变系数格式（鱼眼相机需要4个参数：k1, k2, k3, k4）
        if len(self.D) != 4:
            logger.error(f"[鱼眼相机去畸变] ⚠️ 鱼眼相机畸变系数应为4个参数，当前为{len(self.D)}个")
            logger.error(f"[鱼眼相机去畸变] 畸变系数: {self.D}")
            logger.warning("[鱼眼相机去畸变] 尝试使用前4个参数")
            D_fisheye = self.D[:4] if len(self.D) >= 4 else np.pad(self.D, (0, 4 - len(self.D)), 'constant')
        else:
            D_fisheye = self.D.copy()
        
        # 如果没有缓存或图像尺寸变化，重新计算
        if (not use_cached_map or
                self._undistort_map1 is None or
                self._undistort_map2 is None or
                self.image_size != (w, h)):
            
            log_params = not getattr(self, '_undistort_fisheye_params_logged', False)
            if log_params:
                logger.debug(f"[鱼眼相机去畸变] 原始内参矩阵K:\n{self.K}")
                logger.debug(f"[鱼眼相机去畸变] 畸变系数D: {D_fisheye} (k1, k2, k3, k4)")
            
            # 计算最优新内参矩阵（鱼眼相机）
            self._new_K = cv2.fisheye.estimateNewCameraMatrixForUndistortRectify(
                self.K, D_fisheye, (w, h), np.eye(3), balance=scale
            )
            
            if log_params:
                logger.debug(f"[鱼眼相机去畸变] 新内参矩阵new_K:\n{self._new_K}")
            self._undistort_fisheye_params_logged = True
            
            # 生成去畸变映射（鱼眼相机）
            self._undistort_map1, self._undistort_map2 = cv2.fisheye.initUndistortRectifyMap(
                self.K, D_fisheye, np.eye(3), self._new_K, (w, h), cv2.CV_32FC1
            )
            
            # 鱼眼相机没有ROI概念，整个图像都是有效的
            self._valid_roi = (0, 0, w, h)
            self.image_size = (w, h)
        
        # 应用去畸变
        undistorted = cv2.remap(
            image,
            self._undistort_map1,
            self._undistort_map2,
            cv2.INTER_LINEAR,
            cv2.BORDER_CONSTANT
        )
        
        out_h, out_w = undistorted.shape[:2]
        
        return undistorted, self._new_K.copy()

    def adjust_focal_length(self, focal_length_delta: float) -> np.ndarray:
        """
        调整焦距

        参数:
            focal_length_delta: 焦距增量（像素）

        返回:
            调整后的内参矩阵
        """
        K_adjusted = self.K.copy()
        K_adjusted[0, 0] += focal_length_delta  # fx
        K_adjusted[1, 1] += focal_length_delta  # fy
        return K_adjusted

    def get_intrinsic_matrix(self) -> np.ndarray:
        """获取当前内参矩阵"""
        return self.K.copy()

    def get_distortion_coeffs(self) -> Optional[np.ndarray]:
        """获取畸变系数"""
        return self.D.copy() if self.D is not None else None

    def get_valid_roi(self) -> Optional[Tuple[int, int, int, int]]:
        """
        获取去畸变后的有效ROI区域
        
        返回:
            (x, y, width, height) 或 None（如果没有去畸变）
        """
        return self._valid_roi

    def update_intrinsic(self, new_K: np.ndarray):
        """
        更新内参矩阵

        参数:
            new_K: 新的3x3内参矩阵
        """
        self.K = new_K.copy()
        # 清除缓存的去畸变映射
        self._undistort_map1 = None
        self._undistort_map2 = None
        self._undistort_pinhole_params_logged = False
        self._undistort_fisheye_params_logged = False

    def update_distortion(self, new_D: np.ndarray):
        """
        更新畸变系数

        参数:
            new_D: 新的畸变系数
        """
        self.D = new_D.copy()
        # 清除缓存的去畸变映射
        self._undistort_map1 = None
        self._undistort_map2 = None
        self._undistort_pinhole_params_logged = False
        self._undistort_fisheye_params_logged = False

    @staticmethod
    def resize_image(image: np.ndarray,
                     scale: float = 1.0,
                     interpolation: int = cv2.INTER_LINEAR) -> np.ndarray:
        """
        缩放图像

        参数:
            image: 输入图像
            scale: 缩放比例
            interpolation: 插值方法

        返回:
            缩放后的图像
        """
        if abs(scale - 1.0) < 1e-6:
            return image.copy()

        h, w = image.shape[:2]
        new_w = int(w * scale)
        new_h = int(h * scale)

        return cv2.resize(image, (new_w, new_h), interpolation=interpolation)

    @staticmethod
    def crop_image(image: np.ndarray, roi: Tuple[int, int, int, int]) -> np.ndarray:
        """
        裁剪图像

        参数:
            image: 输入图像
            roi: (x, y, width, height) 感兴趣区域

        返回:
            裁剪后的图像
        """
        x, y, w, h = roi
        return image[y:y + h, x:x + w].copy()


class ImageProcessorFactory:
    """图像处理器工厂类"""

    @staticmethod
    def create_from_config(calib_data: Dict, undistort_scale: float = 1.0) -> ImageProcessor:
        """
        从标定数据创建图像处理器（统一接口）

        参数:
            calib_data: 标定数据字典，包含:
                - intrinsic: 内参矩阵 (9元素列表)
                - distortion: 畸变系数 (可选)
                - width: 图像宽度 (可选)
                - height: 图像高度 (可选)
                - camera_model: 相机模型类型 ('pinhole' 或 'fisheye') (可选，默认'pinhole')
            undistort_scale: 去畸变scale参数，统一使用此值 (默认1.0，保留完整图像)

        返回:
            ImageProcessor实例
        
        注意：
            - 🔧 统一接口：所有ImageProcessor实例使用相同的undistort_scale参数
            - 🔧 支持针孔和鱼眼相机模型：根据camera_model自动选择去畸变方法
            - 建议从配置文件中读取此参数，确保全局一致
        """
        global _warned_legacy
        if not _warned_legacy:
            logger.warning("[Legacy] image.processor.ImageProcessorFactory 仍在使用，计划迁移到 ProjectionManager + base_lib 统一接口后删除")
            _warned_legacy = True
        K = np.array(calib_data['intrinsic']).reshape(3, 3)

        D = None
        if 'distortion' in calib_data and calib_data['distortion'] is not None:
            D = np.array(calib_data['distortion'])

        image_size = None
        if 'width' in calib_data and 'height' in calib_data:
            image_size = (calib_data['width'], calib_data['height'])

        # 🔧 读取相机模型类型（支持 'pinhole' 和 'fisheye'）
        camera_model = calib_data.get('camera_model', 'pinhole')
        if camera_model not in ['pinhole', 'fisheye']:
            logger.warning(f"[ImageProcessorFactory] 未知的相机模型 '{camera_model}'，使用默认值 'pinhole'")
            camera_model = 'pinhole'
        camera_id = calib_data.get('frame_id') or calib_data.get('camera_id') or None

        logger.info(f"[ImageProcessorFactory] 创建ImageProcessor: 相机模型={camera_model}, scale={undistort_scale}" + (f", 相机={camera_id}" if camera_id else ""))

        return ImageProcessor(K, D, image_size, undistort_scale=undistort_scale, camera_model=camera_model, camera_id=camera_id)


# ========== 测试代码 ==========

if __name__ == "__main__":
    # 模拟标定数据
    calib_data = {
        'intrinsic': [
            1000, 0, 960,
            0, 1000, 540,
            0, 0, 1
        ],
        'distortion': [-0.1, 0.05, 0, 0, 0],
        'width': 1920,
        'height': 1080
    }

    # 创建处理器
    processor = ImageProcessorFactory.create_from_config(calib_data)

    # 模拟图像
    test_image = np.random.randint(0, 255, (1080, 1920, 3), dtype=np.uint8)

    # 测试去畸变
    undistorted, new_K = processor.undistort_image(test_image, scale=0.9)
    logger.info("去畸变完成")
    logger.info(f"  原始尺寸: {test_image.shape}")
    logger.info(f"  去畸变后尺寸: {undistorted.shape}")
    logger.info(f"  新内参矩阵:\n{new_K}")

    # 测试焦距调整
    K_adjusted = processor.adjust_focal_length(50.0)
    logger.info("焦距调整 +50像素")
    logger.info(f"  新内参矩阵:\n{K_adjusted}")