# -*- coding: utf-8 -*-
"""
统一颜色管理器
- 维护映射表（2D坐标、3D坐标、强度、距离、颜色）
- 管理配色方案
- 提供可扩展的过滤接口
"""

import numpy as np
from typing import Dict, List, Tuple, Optional, Callable

from lidar_manager.color_utils.color_mapper import ColorMapper
from lidar_manager.logger import get_module_logger

logger = get_module_logger(__name__)


MAPPING_DTYPE = [
    ('pixel_2d', 'i4', 2),
    ('xyz_3d', 'f4', 3),
    ('intensity', 'f4'),
    ('distance', 'f4'),
    ('color_bgr', 'u1', 3),
    ('normalized', 'f4'),
    ('cam_xyz', 'f4', 3),
]


class ColorManager:
    """统一颜色管理器（核心模块）"""

    def __init__(self, initial_distance_range: Optional[Tuple[float, float]] = None,
                 initial_intensity_range: Optional[Tuple[float, float]] = None):
        self._mapping_table: Optional[np.ndarray] = None
        self._filter_mask: Optional[np.ndarray] = None
        self._color_scheme: Optional[List] = None
        self._color_map_func: Optional[Callable] = None
        self._global_attr_min: float = 0.0
        self._global_attr_max: float = 1.0
        self._initial_distance_range: Optional[Tuple[float, float]] = initial_distance_range
        self._initial_intensity_range: Optional[Tuple[float, float]] = initial_intensity_range
        self._has_used_initial_range: bool = False
        self._color_type: str = 'distance'
        self._filter_registry: Dict[str, Callable] = {
            'attribute': self._filter_by_attribute_impl,
            'color': self._filter_by_color_impl,
            'range': self._filter_by_range_impl,
        }

    def build_mapping_table(self,
                            points_2d: np.ndarray,
                            points_xyz: np.ndarray,
                            intensities: np.ndarray,
                            distances: Optional[np.ndarray],
                            cam_xyz: np.ndarray,
                            color_type: str = 'distance') -> None:
        n_points = len(points_2d)
        if n_points == 0:
            logger.warning("[ColorManager] 输入点数为0，无法构建映射表")
            self._mapping_table = None
            return

        is_first_build = (self._mapping_table is None or len(self._mapping_table) == 0)
        if is_first_build:
            self._has_used_initial_range = False

        self._mapping_table = np.zeros(n_points, dtype=MAPPING_DTYPE)
        self._mapping_table['pixel_2d'] = points_2d
        self._mapping_table['xyz_3d'] = points_xyz
        self._mapping_table['intensity'] = intensities

        if distances is None or len(distances) == 0:
            distances = np.linalg.norm(points_xyz, axis=1)
        else:
            computed_distances = np.linalg.norm(points_xyz, axis=1)
            if not np.allclose(distances, computed_distances, rtol=1e-3):
                logger.warning("[ColorManager] 传入的距离与LiDAR系计算不一致，使用计算值")
                distances = computed_distances

        self._mapping_table['distance'] = distances
        self._mapping_table['cam_xyz'] = cam_xyz
        self._color_type = color_type
        self._compute_normalized_and_colors()
        self._filter_mask = None

    def get_mapping_table(self, filtered: bool = False) -> Optional[np.ndarray]:
        if self._mapping_table is None:
            return None
        if filtered and self._filter_mask is not None:
            return self._mapping_table[self._filter_mask]
        return self._mapping_table

    def update_color_scheme(self, color_list: List, recompute_normalized: bool = False,
                           lazy_compute: bool = False) -> None:
        self._color_scheme = color_list
        self._color_map_func = ColorMapper.create_color_map_function(color_list)
        if lazy_compute:
            return
        if self._mapping_table is not None:
            if recompute_normalized:
                self._compute_normalized_and_colors(use_filtered=True)
            else:
                self._compute_colors()
        logger.info(f"[ColorManager] 配色方案已更新，颜色数量: {len(color_list)}")

    def set_normalization_range(self, min_val: float, max_val: float, color_type: str = None) -> None:
        if color_type is not None:
            self._color_type = color_type
        if max_val <= min_val:
            logger.warning(f"[ColorManager] 无效的范围: [{min_val}, {max_val}]")
            return
        self._global_attr_min = float(min_val)
        self._global_attr_max = float(max_val)
        if self._mapping_table is not None:
            self._compute_normalized_and_colors(use_manual_range=True)

    def _compute_normalized_and_colors(self, use_filtered: bool = False, use_manual_range: bool = False) -> None:
        if self._mapping_table is None:
            return
        if self._color_type == 'distance':
            attr_values = self._mapping_table['distance']
        elif self._color_type == 'intensity':
            attr_values = self._mapping_table['intensity']
        else:
            attr_values = self._mapping_table['distance']

        if use_manual_range:
            pass
        elif use_filtered and self._filter_mask is not None:
            filtered_attr_values = attr_values[self._filter_mask]
            if len(filtered_attr_values) > 0:
                self._global_attr_min = float(np.min(filtered_attr_values))
                self._global_attr_max = float(np.max(filtered_attr_values))
            else:
                self._global_attr_min, self._global_attr_max = 0.0, 1.0
        else:
            if not self._has_used_initial_range:
                if self._color_type == 'distance' and self._initial_distance_range is not None:
                    self._global_attr_min = float(self._initial_distance_range[0])
                    self._global_attr_max = float(self._initial_distance_range[1])
                    self._has_used_initial_range = True
                elif self._color_type == 'intensity' and self._initial_intensity_range is not None:
                    self._global_attr_min = float(self._initial_intensity_range[0])
                    self._global_attr_max = float(self._initial_intensity_range[1])
                    self._has_used_initial_range = True
                else:
                    if len(attr_values) > 0:
                        self._global_attr_min = float(np.min(attr_values))
                        self._global_attr_max = float(np.max(attr_values))
                    else:
                        self._global_attr_min, self._global_attr_max = 0.0, 1.0
                    self._has_used_initial_range = True
            else:
                if len(attr_values) > 0:
                    self._global_attr_min = float(np.min(attr_values))
                    self._global_attr_max = float(np.max(attr_values))
                else:
                    self._global_attr_min, self._global_attr_max = 0.0, 1.0

        if self._global_attr_max > self._global_attr_min:
            normalized = (attr_values - self._global_attr_min) / (self._global_attr_max - self._global_attr_min)
            normalized = np.clip(normalized, 0.0, 1.0)
        else:
            normalized = np.ones_like(attr_values) * 0.5
        self._mapping_table['normalized'] = normalized
        self._compute_colors()

    def _compute_colors(self) -> None:
        if self._color_map_func is None or self._mapping_table is None:
            return
        normalized = self._mapping_table['normalized']
        colors = np.array([self._color_map_func(n) for n in normalized])
        self._mapping_table['color_bgr'] = colors

    def apply_filter(self, filter_type: str, threshold_params: Dict) -> np.ndarray:
        if self._mapping_table is None:
            return np.array([], dtype=bool)
        if filter_type not in self._filter_registry:
            logger.error(f"[ColorManager] 未知的过滤类型: {filter_type}")
            return np.ones(len(self._mapping_table), dtype=bool)
        mask = self._filter_registry[filter_type](threshold_params)
        if self._filter_mask is None:
            self._filter_mask = mask
        else:
            self._filter_mask = self._filter_mask & mask
        return self._filter_mask

    def filter_by_attribute(self, reference: float, threshold: float, attr_type: str = 'distance') -> np.ndarray:
        return self.apply_filter('attribute', {'reference': reference, 'threshold': threshold, 'attr_type': attr_type})

    def filter_by_color(self, reference_color: np.ndarray, threshold: float) -> np.ndarray:
        return self.apply_filter('color', {'reference_color': reference_color, 'threshold': threshold})

    def filter_by_range(self, min_val: float, max_val: float, attr_type: str = 'distance') -> np.ndarray:
        return self.apply_filter('range', {'min_val': min_val, 'max_val': max_val, 'attr_type': attr_type})

    def clear_filter(self) -> None:
        self._filter_mask = None

    def _filter_by_attribute_impl(self, params: Dict) -> np.ndarray:
        reference, threshold = params['reference'], params['threshold']
        attr_type = params.get('attr_type', 'distance')
        attr_values = self._mapping_table['distance'] if attr_type == 'distance' else self._mapping_table['intensity']
        if attr_type not in ('distance', 'intensity'):
            attr_values = self._mapping_table['distance']
        return np.abs(attr_values - reference) <= threshold

    def _filter_by_color_impl(self, params: Dict) -> np.ndarray:
        reference_color = np.array(params['reference_color'])
        colors = self._mapping_table['color_bgr']
        return np.linalg.norm(colors - reference_color, axis=1) <= params['threshold']

    def _filter_by_range_impl(self, params: Dict) -> np.ndarray:
        min_val, max_val = params['min_val'], params['max_val']
        attr_type = params.get('attr_type', 'distance')
        attr_values = self._mapping_table['distance'] if attr_type == 'distance' else self._mapping_table['intensity']
        if attr_type not in ('distance', 'intensity'):
            attr_values = self._mapping_table['distance']
        return (attr_values >= min_val) & (attr_values <= max_val)

    def register_filter(self, filter_type: str, filter_func: Callable) -> None:
        self._filter_registry[filter_type] = filter_func

    def get_points_by_pixel(self, px: int, py: int) -> np.ndarray:
        if self._mapping_table is None:
            return np.array([], dtype=MAPPING_DTYPE)
        mask = (self._mapping_table['pixel_2d'][:, 0] == px) & (self._mapping_table['pixel_2d'][:, 1] == py)
        if self._filter_mask is not None:
            mask = mask & self._filter_mask
        return self._mapping_table[mask]

    def get_points_in_radius(self, px: int, py: int, radius: int) -> np.ndarray:
        if self._mapping_table is None:
            return np.array([], dtype=MAPPING_DTYPE)
        pixels = self._mapping_table['pixel_2d']
        mask = np.sqrt((pixels[:, 0] - px)**2 + (pixels[:, 1] - py)**2) <= radius
        if self._filter_mask is not None:
            mask = mask & self._filter_mask
        result = self._mapping_table[mask]
        if len(result) > 0:
            result = result[np.argsort(result['distance'])]
        return result

    def get_global_attr_range(self) -> Tuple[float, float]:
        return (self._global_attr_min, self._global_attr_max)

    def is_valid(self) -> bool:
        return self._mapping_table is not None and len(self._mapping_table) > 0

    def get_point_count(self, filtered: bool = False) -> int:
        if self._mapping_table is None:
            return 0
        if filtered and self._filter_mask is not None:
            return int(np.sum(self._filter_mask))
        return len(self._mapping_table)
