"""
Quantum-Inspired Driver Behavior Representation
基于 arXiv:2608.25907

核心思路：
1. 将驾驶行为特征（速度/加速度/转向/刹车等）映射为量子态
2. 构建密度矩阵表示驾驶员综合状态
3. 通过谱分解提取隐藏的驾驶模式
"""

import numpy as np
from scipy import linalg
from typing import Tuple, List
from dataclasses import dataclass

@dataclass
class DrivingFeatures:
    """驾驶行为特征"""
    speed: float              # km/h
    acceleration: float       # m/s²
    steering_angle: float     # 度
    steering_velocity: float  # 度/秒
    brake_pressure: float     # 0-1
    throttle_position: float  # 0-1
    lane_deviation: float     # m
    headway_distance: float   # m
    time_to_collision: float  # s

class QuantumDriverRepresentation:
    """
    量子启发的驾驶员行为表示
    
    将驾驶特征映射到量子态空间，
    构建密度矩阵，通过谱分解发现隐藏模式
    """
    
    def __init__(self, n_features: int = 9, n_modes: int = 3):
        self.n_features = n_features
        self.n_modes = n_modes  # 预期驾驶模式数
        self.feature_basis = None  # 特征基向量
        self.density_matrix = None
        
    def features_to_quantum_state(self, features: DrivingFeatures) -> np.ndarray:
        """
        将驾驶特征映射为量子态向量
        
        |ψ⟩ = Σᵢ αᵢ |fᵢ⟩
        
        其中 |fᵢ⟩ 是特征基，αᵢ 是振幅
        """
        # 提取特征向量
        feat_vec = np.array([
            features.speed / 120,           # 归一化
            np.tanh(features.acceleration / 3),
            np.tanh(features.steering_angle / 180),
            np.tanh(features.steering_velocity / 60),
            features.brake_pressure,
            features.throttle_position,
            np.tanh(features.lane_deviation / 2),
            np.tanh(features.headway_distance / 50),
            np.tanh(features.time_to_collision / 5),
        ])
        
        # 归一化为量子态（|ψ|²=1）
        norm = np.linalg.norm(feat_vec)
        if norm > 0:
            quantum_state = feat_vec / norm
        else:
            quantum_state = np.zeros_like(feat_vec)
            quantum_state[0] = 1.0
            
        return quantum_state
    
    def build_density_matrix(self, states: List[np.ndarray]) -> np.ndarray:
        """
        构建密度矩阵
        
        ρ = (1/N) Σₙ |ψₙ⟩⟨ψₙ|
        
        密度矩阵捕捉多个时间步的驾驶行为统计特性
        """
        n = self.n_features
        rho = np.zeros((n, n))
        
        for state in states:
            # |ψ⟩⟨ψ| 外积
            rho += np.outer(state, state)
        
        rho /= len(states)  # 归一化
        
        # 确保迹为1（量子力学约束）
        trace = np.trace(rho)
        if trace > 0:
            rho /= trace
            
        self.density_matrix = rho
        return rho
    
    def spectral_decomposition(self, rho: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
        """
        密度矩阵谱分解
        
        ρ = Σᵢ λᵢ |vᵢ⟩⟨vᵢ|
        
        特征值 λᵢ 对应各驾驶模式的概率
        特征向量 |vᵢ⟩ 对应各驾驶模式的特征
        """
        eigenvalues, eigenvectors = np.linalg.eigh(rho)
        
        # 降序排列
        idx = np.argsort(eigenvalues)[::-1]
        eigenvalues = eigenvalues[idx]
        eigenvectors = eigenvectors[:, idx]
        
        # 只保留前 n_modes 个
        eigenvalues = eigenvalues[:self.n_modes]
        eigenvectors = eigenvectors[:, :self.n_modes]
        
        # 归一化特征值
        total = np.sum(eigenvalues)
        if total > 0:
            eigenvalues /= total
            
        return eigenvalues, eigenvectors
    
    def classify_mode(self, features: DrivingFeatures) -> dict:
        """分类当前驾驶模式"""
        # 转换为量子态
        state = self.features_to_quantum_state(features)
        
        # 计算与各模式的内积（投影概率）
        if self.feature_basis is None:
            return {'mode': 'unknown', 'confidence': 0}
        
        projections = []
        for i in range(self.n_modes):
            proj = np.abs(np.dot(state, self.feature_basis[:, i]))**2
            projections.append(proj)
        
        projections = np.array(projections)
        projections /= projections.sum()  # 归一化
        
        mode_idx = np.argmax(projections)
        mode_names = ['free_flow', 'transition', 'congestion']
        
        return {
            'mode': mode_names[mode_idx] if mode_idx < 3 else f'mode_{mode_idx}',
            'confidence': projections[mode_idx],
            'all_projections': projections.tolist(),
        }
    
    def fit(self, feature_sequences: List[DrivingFeatures]) -> dict:
        """
        无监督训练：从驾驶数据中学习驾驶模式
        
        不需要标注——通过密度矩阵谱分解自动发现模式
        """
        # 转换为量子态序列
        quantum_states = [self.features_to_quantum_state(f) for f in feature_sequences]
        
        # 构建密度矩阵
        rho = self.build_density_matrix(quantum_states)
        
        # 谱分解
        eigenvalues, eigenvectors = self.spectral_decomposition(rho)
        
        # 保存模式基
        self.feature_basis = eigenvectors
        
        return {
            'eigenvalues': eigenvalues.tolist(),
            'mode_names': ['free_flow', 'transition', 'congestion'],
            'mode_weights': eigenvalues.tolist(),
            'density_matrix_trace': np.trace(rho),
            'density_matrix_rank': np.linalg.matrix_rank(rho),
        }


# 测试
if __name__ == "__main__":
    model = QuantumDriverRepresentation(n_features=9, n_modes=3)
    
    # 生成模拟驾驶数据
    np.random.seed(42)
    
    # 自由流模式：高速、小转向、稳定
    free_flow = [
        DrivingFeatures(
            speed=np.random.normal(90, 5),
            acceleration=np.random.normal(0, 0.3),
            steering_angle=np.random.normal(0, 3),
            steering_velocity=np.random.normal(0, 5),
            brake_pressure=np.random.uniform(0, 0.05),
            throttle_position=np.random.uniform(0.4, 0.6),
            lane_deviation=np.random.normal(0, 0.2),
            headway_distance=np.random.normal(45, 5),
            time_to_collision=np.random.normal(5, 0.5),
        )
        for _ in range(100)
    ]
    
    # 拥堵模式：低速、频繁刹车
    congestion = [
        DrivingFeatures(
            speed=np.random.normal(20, 5),
            acceleration=np.random.normal(-0.5, 1),
            steering_angle=np.random.normal(0, 10),
            steering_velocity=np.random.normal(0, 15),
            brake_pressure=np.random.uniform(0.1, 0.5),
            throttle_position=np.random.uniform(0.1, 0.3),
            lane_deviation=np.random.normal(0, 0.4),
            headway_distance=np.random.normal(8, 2),
            time_to_collision=np.random.normal(1.5, 0.3),
        )
        for _ in range(100)
    ]
    
    # 过渡模式：中速、变化大
    transition = [
        DrivingFeatures(
            speed=np.random.normal(50, 15),
            acceleration=np.random.normal(0, 1.5),
            steering_angle=np.random.normal(0, 8),
            steering_velocity=np.random.normal(0, 12),
            brake_pressure=np.random.uniform(0.05, 0.2),
            throttle_position=np.random.uniform(0.2, 0.5),
            lane_deviation=np.random.normal(0, 0.3),
            headway_distance=np.random.normal(20, 5),
            time_to_collision=np.random.normal(3, 0.8),
        )
        for _ in range(100)
    ]
    
    # 合并所有数据（无标签）
    all_data = free_flow + congestion + transition
    np.random.shuffle(all_data)
    
    # 无监督训练
    result = model.fit(all_data)
    print("=== 无监督模式发现结果 ===")
    print(f"特征值（模式权重）: {result['eigenvalues']}")
    print(f"模式名称: {result['mode_names']}")
    print(f"密度矩阵迹: {result['density_matrix_trace']:.4f}")
    
    # 分类测试
    print("\n=== 分类测试 ===")
    test_free = DrivingFeatures(
        speed=95, acceleration=0.1, steering_angle=1,
        steering_velocity=2, brake_pressure=0.02, throttle_position=0.5,
        lane_deviation=0.1, headway_distance=48, time_to_collision=5.5
    )
    r1 = model.classify_mode(test_free)
    print(f"自由流场景: mode={r1['mode']}, confidence={r1['confidence']:.2%}")
    
    test_congestion = DrivingFeatures(
        speed=15, acceleration=-0.8, steering_angle=8,
        steering_velocity=18, brake_pressure=0.35, throttle_position=0.15,
        lane_deviation=0.35, headway_distance=7, time_to_collision=1.2
    )
    r2 = model.classify_mode(test_congestion)
    print(f"拥堵场景: mode={r2['mode']}, confidence={r2['confidence']:.2%}")