import numpy as np
from typing import Tuple

class GazeErrorMeasurement:
    """
    道路注视误差测量框架
    
    论文 Section 3.2 的方法实现
    
    测量流程：
    1. 驾驶员声明注视目标（语音）
    2. YOLO检测车辆+车牌
    3. 车牌中心作为参考点
    4. 计算注视光线与参考光线的角度差
    """
    
    def __init__(self, camera_intrinsics: dict):
        """
        Args:
            camera_intrinsics: 相机内参
                fx, fy: 焦距
                cx, cy: 主点
                dist: 畸变系数
        """
        self.fx = camera_intrinsics['fx']
        self.fy = camera_intrinsics['fy']
        self.cx = camera_intrinsics['cx']
        self.cy = camera_intrinsics['cy']
        self.dist = camera_intrinsics.get('dist', None)
    
    def pixel_to_ray(self, pixel: np.ndarray) -> np.ndarray:
        """
        像素坐标转单位光线
        
        Args:
            pixel: (2,) 像素坐标 [u, v]
        
        Returns:
            ray: (3,) 单位光线方向
        """
        # 去畸变
        u, v = pixel
        x = (u - self.cx) / self.fx
        y = (v - self.cy) / self.fy
        
        # 归一化
        norm = np.sqrt(x**2 + y**2 + 1)
        ray = np.array([x, y, 1]) / norm
        return ray
    
    def angular_gaze_error(self, gaze_pixel: np.ndarray,
                           target_pixel: np.ndarray) -> float:
        """
        计算角度注视误差
        
        Args:
            gaze_pixel: (2,) 注视点像素坐标
            target_pixel: (2,) 目标点像素坐标
        
        Returns:
            error_deg: 角度误差（度）
        """
        gaze_ray = self.pixel_to_ray(gaze_pixel)
        target_ray = self.pixel_to_ray(target_pixel)
        
        # 两光线间夹角
        cos_angle = np.dot(gaze_ray, target_ray)
        cos_angle = np.clip(cos_angle, -1, 1)
        angle_rad = np.arccos(cos_angle)
        angle_deg = np.degrees(angle_rad)
        
        return angle_deg
    
    def signed_bias(self, gaze_pixel: np.ndarray,
                    target_pixel: np.ndarray) -> Tuple[float, float]:
        """
        计算有符号偏差（方位角/仰角）
        
        Returns:
            azimuth_bias: 方位角偏差（度）
            elevation_bias: 仰角偏差（度）
        """
        gaze_ray = self.pixel_to_ray(gaze_pixel)
        target_ray = self.pixel_to_ray(target_pixel)
        
        # 分解为方位角和仰角
        gaze_az = np.arctan2(gaze_ray[0], gaze_ray[2])
        gaze_el = np.arcsin(gaze_ray[1])
        target_az = np.arctan2(target_ray[0], target_ray[2])
        target_el = np.arcsin(target_ray[1])
        
        az_bias = np.degrees(gaze_az - target_az)
        el_bias = np.degrees(gaze_el - target_el)
        
        return az_bias, el_bias
    
    def fixation_precision(self, gaze_samples: np.ndarray,
                           target_pixel: np.ndarray) -> float:
        """
        注视内精度（一个注视事件内的分散度）
        
        Args:
            gaze_samples: (N, 2) 一次注视的N帧注视点
            target_pixel: (2,) 目标点
        
        Returns:
            precision_deg: 注视内精度（度）
        """
        errors = []
        for g in gaze_samples:
            errors.append(self.angular_gaze_error(g, target_pixel))
        
        # 中位精度
        return np.median(errors)


class OffsetCalibration:
    """
    偏移校正框架
    
    论文 Section 3.4 的方法：
    1. 从室内校准场景估计偏移量
    2. 将偏移转移到道路场景
    3. 偏移 = 校准场景中注视到目标的角度差的中位数
    """
    
    def __init__(self):
        self.offset = None  # (azimuth, elevation) in degrees
        self.scene_offsets = {}  # 每个校准场景的偏移
    
    def estimate_offset(self, gaze_pixels: np.ndarray,
                        target_pixels: np.ndarray,
                        camera_intrinsics: dict) -> Tuple[float, float]:
        """
        从校准数据估计偏移
        
        Args:
            gaze_pixels: (N, 2) 校准时的注视点
            target_pixels: (N, 2) 校准目标点
            camera_intrinsics: 相机内参
        
        Returns:
            offset: (az_offset, el_offset) 偏移量（度）
        """
        measurer = GazeErrorMeasurement(camera_intrinsics)
        
        az_biases = []
        el_biases = []
        
        for g, t in zip(gaze_pixels, target_pixels):
            az_b, el_b = measurer.signed_bias(g, t)
            az_biases.append(az_b)
            el_biases.append(el_b)
        
        # 中位数偏移
        az_offset = np.median(az_biases)
        el_offset = np.median(el_biases)
        
        self.offset = (az_offset, el_offset)
        return self.offset
    
    def apply_offset(self, gaze_pixels: np.ndarray,
                     camera_intrinsics: dict) -> np.ndarray:
        """
        将偏移应用到道路注视数据
        
        Args:
            gaze_pixels: (N, 2) 原始注视点
        
        Returns:
            corrected_pixels: (N, 2) 校正后注视点
        """
        if self.offset is None:
            return gaze_pixels
        
        # 将偏移从角度转回像素偏移
        az_off, el_off = self.offset
        fx = camera_intrinsics['fx']
        fy = camera_intrinsics['fy']
        
        # 角度偏移→像素偏移（小角度近似）
        du = fx * np.radians(az_off)
        dv = fy * np.radians(el_off)
        
        corrected = gaze_pixels.copy()
        corrected[:, 0] += du
        corrected[:, 1] += dv
        
        return corrected


# 测试
if __name__ == "__main__":
    print("=" * 60)
    print("可穿戴注视追踪精度评估 - 测试")
    print("=" * 60)
    
    # 模拟相机内参
    intrinsics = {'fx': 600, 'fy': 600, 'cx': 320, 'cy': 240}
    
    measurer = GazeErrorMeasurement(intrinsics)
    
    # 模拟注视点和目标
    gaze = np.array([350, 200])  # 注视点
    target = np.array([340, 210])  # 目标车牌中心
    
    error = measurer.angular_gaze_error(gaze, target)
    az_b, el_b = measurer.signed_bias(gaze, target)
    
    print(f"\n注视点: {gaze}")
    print(f"目标点: {target}")
    print(f"角度误差: {error:.2f}°")
    print(f"方位角偏差: {az_b:.2f}°")
    print(f"仰角偏差: {el_b:.2f}°")
    
    # 偏移校正
    calibrator = OffsetCalibration()
    
    # 模拟校准数据（有系统性偏移）
    np.random.seed(42)
    n_cal = 100
    targets_cal = np.random.randint(200, 440, (n_cal, 2)).astype(float)
    gazes_cal = targets_cal + np.array([8, -5]) + np.random.randn(n_cal, 2) * 3
    
    offset = calibrator.estimate_offset(gazes_cal, targets_cal, intrinsics)
    print(f"\n估计偏移: 方位角={offset[0]:.2f}°, 仰角={offset[1]:.2f}°")
    
    # 模拟道路数据
    road_targets = np.random.randint(200, 440, (41, 2)).astype(float)
    road_gazes = road_targets + np.array([8, -5]) + np.random.randn(41, 2) * 4
    
    # 校正前误差
    errors_before = [measurer.angular_gaze_error(g, t) 
                     for g, t in zip(road_gazes, road_targets)]
    print(f"\n校正前: 平均误差={np.mean(errors_before):.2f}°")
    
    # 校正后误差
    corrected_gazes = calibrator.apply_offset(road_gazes, intrinsics)
    errors_after = [measurer.angular_gaze_error(g, t)
                    for g, t in zip(corrected_gazes, road_targets)]
    print(f"校正后: 平均误差={np.mean(errors_after):.2f}°")
    print(f"改善: {(1-np.mean(errors_after)/np.mean(errors_before))*100:.1f}%")
    print(f"改善场景数: {sum(1 for a,b in zip(errors_after, errors_before) if a < b)}/41")

class OnlineRecalibration:
    """
    在线重校准策略
    
    论文发现：偏移在不同session间变化
    → 可靠BEV监督需要在线重校准和条件相关的偏移估计
    """
    
    def __init__(self, initial_offset=(0, 0), 
                 adaptation_rate=0.01,
                 min_confidence=0.5):
        """
        Args:
            initial_offset: 初始偏移
            adaptation_rate: 自适应学习率
            min_confidence: 最小置信度才更新偏移
        """
        self.offset = np.array(initial_offset)
        self.eta = adaptation_rate
        self.min_conf = min_confidence
        self.confidence = 0.0
        self.frame_count = 0
        self.condition_offset = {}  # 条件相关偏移
        
    def update(self, detected_target: np.ndarray,
               gaze_pixel: np.ndarray,
               condition: str = 'default'):
        """
        在线更新偏移（当有可靠检测目标时）
        
        Args:
            detected_target: 检测到的目标位置（如交通标志）
            gaze_pixel: 当前注视点
            condition: 当前驾驶条件（如highway/urban/night）
        """
        # 估计瞬时偏移
        instant_offset = detected_target - gaze_pixel
        
        # 置信度基于检测置信度
        if self.confidence > self.min_conf:
            # 指数移动平均
            self.offset = (1 - self.eta) * self.offset + self.eta * instant_offset
            
            # 条件相关偏移
            if condition not in self.condition_offset:
                self.condition_offset[condition] = self.offset.copy()
            else:
                co = self.condition_offset[condition]
                self.condition_offset[condition] = (1 - self.eta) * co + self.eta * instant_offset
        
        self.frame_count += 1
        return self.offset
    
    def get_condition_offset(self, condition: str) -> np.ndarray:
        """获取条件相关偏移"""
        if condition in self.condition_offset:
            return self.condition_offset[condition]
        return self.offset

class UncertaintyAwareGaze:
    """
    不确定性感知的注视表示
    
    论文建议：在BEV中表示为不确定区域
    而非精确点，用于安全关键决策
    """
    
    def __init__(self, gaze_point, angular_error_std):
        """
        Args:
            gaze_point: (2,) 注视点像素坐标
            angular_error_std: 角度误差标准差（度）
        """
        self.gaze = np.array(gaze_point)
        self.sigma = np.radians(angular_error_std)
        
    def to_ellipse(self, fx, fy, cx, cy, distance=10.0):
        """
        转换为3D空间中的不确定性椭圆
        
        Args:
            distance: 到目标的距离（米）
        
        Returns:
            ellipse_params: (center, axes, angle)
        """
        # 角度误差→空间误差
        spatial_sigma = distance * self.sigma  # 米
        
        # 椭圆参数
        center = self.gaze
        axes = (fx * self.sigma, fy * self.sigma)  # 像素
        
        return {
            'center': center,
            'semi_axes': axes,
            'angle': 0,  # 简化为圆形
            'spatial_radius_m': spatial_sigma
        }
    
    def covers_target(self, target_point, confidence=0.95):
        """
        判断目标是否在注视不确定性区域内
        
        Args:
            target_point: (2,) 目标像素坐标
            confidence: 置信度
        
        Returns:
            is_covered: bool
            probability: 目标在区域内的概率
        """
        from scipy.stats import chi2
        
        dist = np.linalg.norm(self.gaze - target_point)
        threshold = np.sqrt(chi2.ppf(confidence, df=2)) * self.sigma
        
        return dist < threshold * 1000  # 像素尺度转换