"""
跨源视觉负荷指数(VLI)预测框架
论文：Sensors 2026, 26(17), 5521
核心思路：训练用眼动数据，推理用车辆运动学
"""

import numpy as np
from sklearn.ensemble import GradientBoostingRegressor
from sklearn.model_selection import GroupKFold
from sklearn.metrics import mean_squared_error, r2_score
from scipy.signal import savgol_filter

class VisualLoadIndexPredictor:
    """
    跨源VLI预测器
    
    训练阶段：使用眼动特征计算VLI代理标签
    推理阶段：仅使用车辆运动学特征预测VLI
    """
    
    def __init__(self, config: dict):
        super().__init__()
        # 车辆运动学特征
        self.kinematic_features = config.get('kinematic_features', [
            'steering_angle',       # 方向盘转角
            'steering_velocity',    # 方向盘角速度
            'vehicle_speed',        # 车速
            'lateral_accel',        # 横向加速度
            'longitudinal_accel',   # 纵向加速度
            'yaw_rate',             # 横摆角速度
            'lane_distance',        # 车道偏离距离
            'time_since_departure'  # 出发后时长
        ])
        
        # VLI计算参数
        self.gaze_threshold = config.get('gaze_threshold', 0.3)  # 视线偏离阈值
        self.window_sec = config.get('window_sec', 10)            # 滑动窗口
        self.fps = config.get('fps', 30)                          # 采样率
        
        # 模型
        self.model = None
        self.architectures = ['GBDT', 'RandomForest', 'XGBoost', 'Linear', 'MLP']
        
    def compute_vli_proxy(self, gaze_data: np.ndarray) -> np.ndarray:
        """
        从眼动数据计算VLI代理标签
        
        Args:
            gaze_data: 眼动数据, shape=(N, 4)
                       columns: [gaze_x, gaze_y, pupil_diameter, blink_rate]
        
        Returns:
            vli: 视觉负荷指数, shape=(N-window,)
                 0.0 = 低负荷（视线在路上）
                 1.0 = 高负荷（视线偏离道路）
        
        公式（论文Section 2.2）：
            VLI(t) = mean(|gaze_offset(t)| > threshold) over window
            + α * (1 - gaze_stability(t))
            + β * blink_rate_anomaly(t)
        """
        window = int(self.window_sec * self.fps)
        vli_values = []
        
        for i in range(len(gaze_data) - window):
            w = gaze_data[i:i+window]
            
            # 视线偏离度
            gaze_offset = np.sqrt(w[:, 0]**2 + w[:, 1]**2)
            off_road_ratio = np.mean(gaze_offset > self.gaze_threshold)
            
            # 注视稳定性（分散度）
            gaze_stability = 1.0 - np.std(gaze_offset) / (np.mean(gaze_offset) + 1e-6)
            
            # 眨眼异常
            blink_anomaly = np.abs(w[:, 3] - np.median(w[:, 3])) / (np.median(w[:, 3]) + 1e-6)
            blink_score = np.mean(blink_anomaly)
            
            # 加权融合
            alpha = 0.3
            beta = 0.2
            vli = off_road_ratio + alpha * (1 - gaze_stability) + beta * blink_score
            vli = np.clip(vli, 0.0, 1.0)
            
            vli_values.append(vli)
        
        return np.array(vli_values)
    
    def extract_kinematic_features(self, can_data: np.ndarray) -> np.ndarray:
        """
        从CAN总线数据提取运动学特征
        
        Args:
            can_data: CAN总线数据, shape=(N, 8)
                      columns: [steering_angle, speed, accel_x, accel_y,
                                yaw_rate, lane_dist, brake, throttle]
        
        Returns:
            features: 运动学特征矩阵, shape=(N, M)
        """
        features = np.zeros((len(can_data), 12))
        
        # 基础特征
        features[:, 0] = can_data[:, 0]  # 方向盘转角
        features[:, 1] = can_data[:, 1]  # 车速
        features[:, 2] = can_data[:, 2]  # 横向加速度
        features[:, 3] = can_data[:, 3]  # 纵向加速度
        features[:, 4] = can_data[:, 4]  # 横摆角速度
        features[:, 5] = can_data[:, 5]  # 车道偏离
        
        # 派生特征
        features[:, 6] = np.gradient(can_data[:, 0])  # 方向盘角速度
        features[:, 7] = np.gradient(can_data[:, 1])  # 加速度变化率
        features[:, 8] = savgol_filter(can_data[:, 2], 11, 2)  # 平滑横向加速
        features[:, 9] = can_data[:, 0] * can_data[:, 1]  # 转向×速度交互
        features[:, 10] = np.abs(can_data[:, 4]) * can_data[:, 1]  # 横摆×速度
        features[:, 11] = np.gradient(can_data[:, 5])  # 车道偏离速率
        
        return features
    
    def train_cross_source(self, 
                            can_data: np.ndarray, 
                            gaze_data: np.ndarray,
                            driver_ids: np.ndarray,
                            architecture: str = 'GBDT'):
        """
        跨源训练：用眼动数据生成标签，用车载数据训练模型
        
        关键：GroupKFold按驾驶员分组，确保驾驶员独立性
        
        Args:
            can_data: CAN总线数据
            gaze_data: 同步眼动数据
            driver_ids: 驾驶员ID（用于GroupKFold）
            architecture: 模型架构
        """
        # Step 1: 计算VLI代理标签
        vli_labels = self.compute_vli_proxy(gaze_data)
        
        # Step 2: 对齐时间窗口
        window = int(self.window_sec * self.fps)
        kinematic_features = self.extract_kinematic_features(can_data[:-window])
        
        # Step 3: 按驾驶员分组交叉验证
        gkf = GroupKFold(n_splits=5)
        
        results = []
        for train_idx, test_idx in gkf.split(kinematic_features, vli_labels, 
                                              groups=driver_ids[:-window]):
            X_train, X_test = kinematic_features[train_idx], kinematic_features[test_idx]
            y_train, y_test = vli_labels[train_idx], vli_labels[test_idx]
            
            # 选择架构
            if architecture == 'GBDT':
                model = GradientBoostingRegressor(
                    n_estimators=200, max_depth=5, learning_rate=0.05,
                    subsample=0.8, random_state=42
                )
            elif architecture == 'XGBoost':
                import xgboost as xgb
                model = xgb.XGBRegressor(
                    n_estimators=200, max_depth=5, learning_rate=0.05,
                    subsample=0.8, random_state=42
                )
            elif architecture == 'RandomForest':
                from sklearn.ensemble import RandomForestRegressor
                model = RandomForestRegressor(
                    n_estimators=200, max_depth=8, random_state=42
                )
            elif architecture == 'Linear':
                from sklearn.linear_model import Ridge
                model = Ridge(alpha=1.0)
            elif architecture == 'MLP':
                from sklearn.neural_network import MLPRegressor
                model = MLPRegressor(
                    hidden_layer_sizes=(64, 32), max_iter=500, random_state=42
                )
            
            model.fit(X_train, y_train)
            y_pred = model.predict(X_test)
            
            rmse = np.sqrt(mean_squared_error(y_test, y_pred))
            r2 = r2_score(y_test, y_pred)
            
            results.append({
                'architecture': architecture,
                'rmse': rmse,
                'r2': r2,
                'n_train': len(train_idx),
                'n_test': len(test_idx)
            })
        
        self.model = model
        return results
    
    def predict_realtime(self, can_data: np.ndarray) -> float:
        """
        实时推理：仅需CAN总线数据，无需眼动仪
        
        Args:
            can_data: 最近10秒CAN数据, shape=(300, 8) (30fps)
        
        Returns:
            vli_pred: 预测视觉负荷指数 [0, 1]
        """
        features = self.extract_kinematic_features(can_data)
        latest_feature = features[-1:].reshape(1, -1)
        vli_pred = self.model.predict(latest_feature)[0]
        return float(np.clip(vli_pred, 0.0, 1.0))


# ============ 实际测试 ============
if __name__ == "__main__":
    import time
    
    np.random.seed(42)
    
    # 模拟数据：10分钟驾驶，30fps
    n_samples = 10 * 60 * 30  # 18000
    
    # CAN总线数据模拟
    can_data = np.zeros((n_samples, 8))
    can_data[:, 0] = np.cumsum(np.random.normal(0, 0.5, n_samples))  # 方向盘
    can_data[:, 1] = 60 + 20 * np.sin(np.linspace(0, 10, n_samples))  # 车速
    can_data[:, 2] = np.random.normal(0, 0.1, n_samples)  # 横向加速度
    can_data[:, 3] = np.random.normal(0, 0.05, n_samples)  # 纵向加速度
    can_data[:, 4] = np.random.normal(0, 0.02, n_samples)  # 横摆角速度
    can_data[:, 5] = np.random.normal(0, 0.3, n_samples)  # 车道偏离
    
    # 眼动数据模拟（训练时使用）
    gaze_data = np.zeros((n_samples, 4))
    gaze_data[:, 0] = np.random.normal(0, 0.2, n_samples)  # gaze_x
    gaze_data[:, 1] = np.random.normal(0, 0.15, n_samples)  # gaze_y
    gaze_data[:, 2] = 4.0 + np.random.normal(0, 0.3, n_samples)  # 瞳孔直径
    gaze_data[:, 3] = 15 + np.random.poisson(1, n_samples)  # 眨眼率
    
    # 驾驶员ID（10名驾驶员）
    driver_ids = np.repeat(np.arange(10), n_samples // 10)
    
    # 训练
    predictor = VisualLoadIndexPredictor({'fps': 30, 'window_sec': 10})
    
    print("=" * 60)
    print("跨源VLI预测实验")
    print("=" * 60)
    
    for arch in ['GBDT', 'RandomForest', 'Linear']:
        results = predictor.train_cross_source(can_data, gaze_data, 
                                                driver_ids, architecture=arch)
        avg_rmse = np.mean([r['rmse'] for r in results])
        avg_r2 = np.mean([r['r2'] for r in results])
        print(f"\n{arch}:")
        print(f"  RMSE: {avg_rmse:.4f} (±{np.std([r['rmse'] for r in results]):.4f})")
        print(f"  R²: {avg_r2:.4f} (±{np.std([r['r2'] for r in results]):.4f})")
    
    # 实时预测测试
    test_can = can_data[:300]  # 10秒数据
    vli = predictor.predict_realtime(test_can)
    print(f"\n实时VLI预测: {vli:.3f}")
    print(f"  VLI > 0.5: {'⚠️ 高视觉负荷' if vli > 0.5 else '✅ 正常'}")