# -*- coding: utf-8 -*-
"""
投影管理器 - 缓存管理模块
负责去畸变缓存、点云投影缓存、背景缓存的管理
"""

from typing import Dict, List, Optional, Tuple

import numpy as np

from lidar_manager.logger import get_module_logger

logger = get_module_logger(__name__)


class ProjectionCacheManager:
    """投影缓存管理器"""

    def __init__(self, max_cache_size_mb: float = 500.0):
        self._max_cache_size_mb = max_cache_size_mb
        self._cached_undistorted_image: Optional[np.ndarray] = None
        self._cached_original_image_hash: Optional[int] = None
        self._cached_points_hash: Optional[int] = None
        self._cached_K_hash: Optional[int] = None
        self._cached_T_hash: Optional[int] = None
        self._cached_front_points_xyz: Optional[np.ndarray] = None
        self._cached_front_points_cam: Optional[np.ndarray] = None
        self._cached_front_points_2d: Optional[np.ndarray] = None
        self._cached_front_valid_mask: Optional[np.ndarray] = None
        self._cached_front_intensities: Optional[np.ndarray] = None
        self._cached_front_distances: Optional[np.ndarray] = None
        self._angle_update_threshold: float = 5.0
        self._baseline_extrinsic: Optional[np.ndarray] = None
        self._baseline_rotation_matrix: Optional[np.ndarray] = None
        self._cached_background_color: Optional[List[int]] = None
        self._cached_background_image: Optional[np.ndarray] = None
        self._cached_background_size: Optional[Tuple[int, int]] = None
        self._cache_stats = {
            'undistort_hits': 0, 'undistort_misses': 0,
            'pointcloud_hits': 0, 'pointcloud_misses': 0,
        }

    def clear_undistort_cache(self) -> None:
        self._cached_undistorted_image = None
        self._cached_original_image_hash = None

    def clear_pointcloud_cache(self) -> None:
        self._cached_points_hash = None
        self._cached_K_hash = None
        self._cached_T_hash = None
        self._cached_front_points_xyz = None
        self._cached_front_points_cam = None
        self._cached_front_points_2d = None
        self._cached_front_valid_mask = None
        self._cached_front_intensities = None
        self._cached_front_distances = None

    def clear_background_cache(self) -> None:
        self._cached_background_image = None
        self._cached_background_size = None

    def clear_all_cache(self) -> None:
        self.clear_undistort_cache()
        self.clear_pointcloud_cache()
        self.clear_background_cache()

    def get_undistorted_image(self, image_hash: int) -> Optional[np.ndarray]:
        if (self._cached_undistorted_image is not None and
                self._cached_original_image_hash == image_hash):
            self._cache_stats['undistort_hits'] += 1
            return self._cached_undistorted_image.copy()
        return None

    def set_undistorted_image(self, image: np.ndarray, image_hash: int) -> None:
        self._cached_undistorted_image = image.copy()
        self._cached_original_image_hash = image_hash
        self._cache_stats['undistort_misses'] += 1

    def get_pointcloud_cache(self, points_hash: int, K_hash: Optional[int]) -> Optional[Dict]:
        if (self._cached_points_hash == points_hash and
                self._cached_K_hash == K_hash and
                self._cached_front_points_xyz is not None):
            self._cache_stats['pointcloud_hits'] += 1
            return {
                'points_xyz': self._cached_front_points_xyz,
                'points_cam': self._cached_front_points_cam,
                'points_2d': self._cached_front_points_2d,
                'valid_mask': self._cached_front_valid_mask,
                'intensities': self._cached_front_intensities,
                'distances': self._cached_front_distances,
            }
        return None

    def set_pointcloud_cache(self, points_hash: int, K_hash: Optional[int],
                              T_hash: Optional[int], cache_data: Dict) -> None:
        self._cached_points_hash = points_hash
        self._cached_K_hash = K_hash
        self._cached_T_hash = T_hash
        self._cached_front_points_xyz = cache_data.get('points_xyz')
        self._cached_front_points_cam = cache_data.get('points_cam')
        self._cached_front_points_2d = cache_data.get('points_2d')
        self._cached_front_valid_mask = cache_data.get('valid_mask')
        self._cached_front_intensities = cache_data.get('intensities')
        self._cached_front_distances = cache_data.get('distances')
        self._cache_stats['pointcloud_misses'] += 1

    def get_background_image(self, size: Tuple[int, int],
                             background_color: Optional[List[int]]) -> Optional[np.ndarray]:
        if (self._cached_background_image is not None and
                self._cached_background_size == size and
                self._cached_background_color == background_color):
            return self._cached_background_image.copy()
        return None

    def set_background_image(self, image: np.ndarray, size: Tuple[int, int],
                             background_color: Optional[List[int]]) -> None:
        self._cached_background_image = image.copy()
        self._cached_background_size = size
        self._cached_background_color = background_color

    def set_background_color(self, color: Optional[List[int]]) -> None:
        self._cached_background_color = color
        self._cached_background_image = None

    def get_background_color(self) -> Optional[List[int]]:
        return self._cached_background_color

    def reset_baseline_extrinsic(self, T: Optional[np.ndarray] = None) -> None:
        if T is None:
            self._baseline_extrinsic = None
            self._baseline_rotation_matrix = None
        else:
            self._baseline_extrinsic = T.copy()
            self._baseline_rotation_matrix = T[:3, :3].copy()

    def get_baseline_rotation_matrix(self) -> Optional[np.ndarray]:
        return self._baseline_rotation_matrix

    def should_update_cache_by_angle(self, current_T: np.ndarray,
                                      compute_angle_diff) -> bool:
        if self._baseline_rotation_matrix is None:
            return True
        current_R = current_T[:3, :3]
        angle_diff = compute_angle_diff(self._baseline_rotation_matrix, current_R)
        return angle_diff >= self._angle_update_threshold

    def estimate_cache_size(self) -> float:
        size_mb = 0.0
        if self._cached_undistorted_image is not None:
            size_mb += self._cached_undistorted_image.nbytes / (1024 * 1024)
        if self._cached_front_points_xyz is not None:
            size_mb += self._cached_front_points_xyz.nbytes / (1024 * 1024)
        if self._cached_front_points_cam is not None:
            size_mb += self._cached_front_points_cam.nbytes / (1024 * 1024)
        if self._cached_front_points_2d is not None:
            size_mb += self._cached_front_points_2d.nbytes / (1024 * 1024)
        if self._cached_front_intensities is not None:
            size_mb += self._cached_front_intensities.nbytes / (1024 * 1024)
        if self._cached_front_distances is not None:
            size_mb += self._cached_front_distances.nbytes / (1024 * 1024)
        if self._cached_background_image is not None:
            size_mb += self._cached_background_image.nbytes / (1024 * 1024)
        return size_mb

    def get_max_cache_size_mb(self) -> float:
        return self._max_cache_size_mb

    def get_cache_stats(self) -> Dict:
        undistort_total = self._cache_stats['undistort_hits'] + self._cache_stats['undistort_misses']
        pointcloud_total = self._cache_stats['pointcloud_hits'] + self._cache_stats['pointcloud_misses']
        undistort_hit_rate = (self._cache_stats['undistort_hits'] / undistort_total * 100
                             if undistort_total > 0 else 0.0)
        pointcloud_hit_rate = (self._cache_stats['pointcloud_hits'] / pointcloud_total * 100
                              if pointcloud_total > 0 else 0.0)
        return {
            'undistort': {
                'hits': self._cache_stats['undistort_hits'],
                'misses': self._cache_stats['undistort_misses'],
                'total': undistort_total,
                'hit_rate': f"{undistort_hit_rate:.2f}%",
            },
            'pointcloud': {
                'hits': self._cache_stats['pointcloud_hits'],
                'misses': self._cache_stats['pointcloud_misses'],
                'total': pointcloud_total,
                'hit_rate': f"{pointcloud_hit_rate:.2f}%",
            },
            'cache_size_mb': self.estimate_cache_size(),
        }
