import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np

class DriverDensityMatrix(nn.Module):
    """
    驾驶员行为密度矩阵
    
    将驾驶行为编码为量子密度矩阵，统一表征：
    - 行为不确定性（对角元素 = 概率）
    - 行为相干性（非对角元素 = 状态间干涉）
    - 时序演化（密度矩阵的幺正/近似幺正变换）
    
    论文Section 3: Quantum-Inspired Representation
    
    Args:
        n_basis: 行为基矢维度（默认3: 自由流/过渡/拥堵）
        n_features: 输入特征维度（速度/加速度/间距等）
    """
    
    def __init__(self, n_basis: int = 3, n_features: int = 6):
        super().__init__()
        self.n_basis = n_basis
        
        # 特征提取网络
        self.encoder = nn.Sequential(
            nn.Linear(n_features, 64),
            nn.ReLU(),
            nn.Linear(64, 128),
            nn.ReLU(),
            nn.Linear(128, n_basis * n_basis)
        )
        
        # 演化算子网络（学习近似幺正变换）
        self.evolution = nn.Sequential(
            nn.Linear(n_basis * n_basis, 128),
            nn.ReLU(),
            nn.Linear(128, n_basis * n_basis)
        )
        
    def to_density_matrix(self, features: torch.Tensor) -> torch.Tensor:
        """
        将驾驶特征转换为密度矩阵
        
        Args:
            features: (B, n_features) 包含速度/加速度/间距/相对速度等
            
        Returns:
            rho: (B, n_basis, n_basis) 密度矩阵，满足Tr(ρ)=1, ρ≥0
        """
        # 编码为矩阵元素
        raw = self.encoder(features)  # (B, n_basis²)
        mat = raw.view(-1, self.n_basis, self.n_basis)
        
        # 厄米化：ρ = (M + M†) / 2
        mat = (mat + mat.transpose(-1, -2)) / 2
        
        # 正定化：ρ → ρ² + λI，再归一化
        rho = F.softmax(mat.flatten(1), dim=1).view_as(mat)
        
        # 迹归一化：Tr(ρ) = 1
        trace = torch.diagonal(rho, dim1=-1, dim2=-2).sum(-1, keepdim=True)
        rho = rho / (trace.unsqueeze(-1) + 1e-8)
        
        return rho
    
    def evolve(self, rho: torch.Tensor) -> torch.Tensor:
        """
        密度矩阵演化（近似幺正变换）
        
        Args:
            rho: (B, n_basis, n_basis) 当前密度矩阵
            
        Returns:
            rho_next: (B, n_basis, n_basis) 演化后密度矩阵
        """
        # 学习演化算子
        flat = rho.flatten(1)
        evolved = self.evolution(flat)
        U = evolved.view(-1, self.n_basis, self.n_basis)
        
        # 近似幺正：U → U + U† 的极分解近似
        rho_next = U @ rho @ U.transpose(-1, -2)
        
        # 厄米化和归一化
        rho_next = (rho_next + rho_next.transpose(-1, -2)) / 2
        trace = torch.diagonal(rho_next, dim1=-1, dim2=-2).sum(-1, keepdim=True)
        rho_next = rho_next / (trace.unsqueeze(-1) + 1e-8)
        
        return rho_next
    
    def measure(self, rho: torch.Tensor, observable: torch.Tensor) -> torch.Tensor:
        """
        测量：计算可观测量的期望值
        
        ⟨O⟩ = Tr(Oρ)
        
        Args:
            rho: (B, n_basis, n_basis) 密度矩阵
            observable: (n_basis, n_basis) 可观测算子（如速度、加速度）
            
        Returns:
            expectation: (B,) 期望值
        """
        # ⟨O⟩ = Tr(O @ ρ) = Σ_i Σ_j O_ij * ρ_ji
        product = torch.matmul(observable.unsqueeze(0), rho)
        expectation = torch.diagonal(product, dim1=-1, dim2=-2).sum(-1)
        return expectation
    
    def forward(self, features: torch.Tensor) -> dict:
        """
        完整前向传播
        
        Args:
            features: (B, n_features) 驾驶特征序列
            
        Returns:
            dict: {
                'rho': 当前密度矩阵,
                'rho_next': 演化后密度矩阵,
                'profile': 行为剖面概率
            }
        """
        rho = self.to_density_matrix(features)
        rho_next = self.evolve(rho)
        
        # 对角元素 = 各行为基模式的概率
        profile = torch.diagonal(rho, dim1=-1, dim2=-2)
        
        return {
            'rho': rho,
            'rho_next': rho_next,
            'profile': profile  # (B, n_basis) 自由流/过渡/拥堵概率
        }


class QuantumDriverModel(nn.Module):
    """
    量子启发的驾驶行为完整模型
    
    架构:
      [驾驶特征] → [密度矩阵编码] → [时序演化] → [行为预测]
    
    论文Section 4: Framework Architecture
    
    Args:
        n_basis: 行为基矢数
        n_features: 输入特征维度
        n_steps: 历史步数
        predict_horizon: 预测步数
    """
    
    def __init__(self, n_basis: int = 3, n_features: int = 6,
                 n_steps: int = 10, predict_horizon: int = 5):
        super().__init__()
        self.n_basis = n_basis
        self.n_steps = n_steps
        self.predict_horizon = predict_horizon
        
        # 密度矩阵编码器
        self.encoder = DriverDensityMatrix(n_basis, n_features)
        
        # 时序演化（多步）
        self.evolution_layers = nn.ModuleList([
            DriverDensityMatrix(n_basis, n_features)
            for _ in range(predict_horizon)
        ])
        
        # 可观测量（速度、加速度）
        self.observables = nn.Parameter(
            torch.randn(predict_horizon, n_basis, n_basis) * 0.1
        )
        
        # 行为剖面分类
        self.classifier = nn.Sequential(
            nn.Linear(n_basis, 32),
            nn.ReLU(),
            nn.Linear(32, n_basis)
        )
        
    def forward(self, x: torch.Tensor) -> dict:
        """
        Args:
            x: (B, n_steps, n_features) 驾驶特征历史序列
            
        Returns:
            dict: {
                'predictions': (B, predict_horizon) 预测值,
                'profiles': (B, n_basis) 行为剖面
            }
        """
        B = x.size(0)
        
        # 编码当前状态
        current_features = x[:, -1, :]  # (B, n_features)
        result = self.encoder(current_features)
        rho = result['rho']
        
        # 多步演化
        predictions = []
        profiles = []
        
        for t, layer in enumerate(self.evolution_layers):
            # 演化
            rho = layer.evolve(rho)
            
            # 测量预测值
            obs = self.observables[t]
            pred = self.encoder.measure(rho, obs)
            predictions.append(pred)
            
            # 行为剖面
            profile = torch.diagonal(rho, dim1=-1, dim2=-2)
            profiles.append(profile)
        
        predictions = torch.stack(predictions, dim=1)  # (B, predict_horizon)
        profiles = torch.stack(profiles, dim=1).mean(dim=1)  # (B, n_basis)
        
        # 分类
        logits = self.classifier(profiles)
        
        return {
            'predictions': predictions,
            'profiles': profiles,
            'logits': logits
        }


# ============ 完整测试代码 ============
if __name__ == "__main__":
    import time
    from sklearn.metrics import accuracy_score
    
    # 模型参数
    N_BASIS = 3          # 自由流/过渡/拥堵
    N_FEATURES = 6       # 速度/加速度/间距/相对速度/车道偏移/方向盘角
    N_STEPS = 10         # 历史10步
    PREDICT_HORIZON = 5  # 预测5步
    
    # 行为模式名称
    PROFILES = ['Free Flow', 'Transition', 'Congestion']
    
    # 初始化模型
    model = QuantumDriverModel(
        n_basis=N_BASIS,
        n_features=N_FEATURES,
        n_steps=N_STEPS,
        predict_horizon=PREDICT_HORIZON
    )
    
    total_params = sum(p.numel() for p in model.parameters())
    print(f"QuantumDriverModel参数量: {total_params:,}")
    
    # 模拟I-24 MOTION数据
    np.random.seed(42)
    BATCH_SIZE = 32
    
    def generate_driving_sequence(profile: str, batch=1):
        """生成模拟驾驶序列"""
        t = np.linspace(0, 5, N_STEPS)
        
        if profile == 'Free Flow':
            # 自由流：高速、小加减速度、大间距
            speed = 28 + np.random.randn(batch, N_STEPS) * 1.5
            accel = np.random.randn(batch, N_STEPS) * 0.3
            gap = 40 + np.random.randn(batch, N_STEPS) * 5
            rel_speed = np.random.randn(batch, N_STEPS) * 0.5
            lane_offset = np.random.randn(batch, N_STEPS) * 0.1
            steering = np.random.randn(batch, N_STEPS) * 0.02
        elif profile == 'Transition':
            # 过渡：中速、变化大
            speed = 18 + np.sin(t * 2) * 3 + np.random.randn(batch, N_STEPS) * 1.5
            accel = np.cos(t * 2) * 0.8 + np.random.randn(batch, N_STEPS) * 0.3
            gap = 25 + np.sin(t * 2) * 5 + np.random.randn(batch, N_STEPS) * 3
            rel_speed = np.sin(t * 2) * 1 + np.random.randn(batch, N_STEPS) * 0.5
            lane_offset = np.random.randn(batch, N_STEPS) * 0.2
            steering = np.sin(t * 2) * 0.05 + np.random.randn(batch, N_STEPS) * 0.02
        else:  # Congestion
            # 拥堵：低速、大加减速度、小间距
            speed = 5 + np.random.randn(batch, N_STEPS) * 2
            accel = np.random.randn(batch, N_STEPS) * 1.5
            gap = 8 + np.random.randn(batch, N_STEPS) * 2
            rel_speed = np.random.randn(batch, N_STEPS) * 1.5
            lane_offset = np.random.randn(batch, N_STEPS) * 0.3
            steering = np.random.randn(batch, N_STEPS) * 0.05
        
        features = np.stack([speed, accel, gap, rel_speed, lane_offset, steering], 
                           axis=-1)  # (batch, N_STEPS, N_FEATURES)
        return torch.tensor(features, dtype=torch.float32)
    
    # 生成测试数据
    print("\n--- 生成模拟数据 ---")
    x_data = []
    y_data = []
    for i, profile in enumerate(PROFILES):
        seq = generate_driving_sequence(profile, batch=20)
        x_data.append(seq)
        y_data.append(torch.full((20,), i, dtype=torch.long))
    
    x_all = torch.cat(x_data)  # (60, N_STEPS, N_FEATURES)
    y_all = torch.cat(y_data)  # (60,)
    
    # 前向传播
    print("\n--- 前向传播 ---")
    model.eval()
    with torch.no_grad():
        out = model(x_all)
        
        print(f"输入: {x_all.shape}")
        print(f"预测值: {out['predictions'].shape}")
        print(f"行为剖面: {out['profiles'].shape}")
        
        # 行为剖面分析
        for i, profile in enumerate(PROFILES):
            idx = (y_all == i)
            mean_profile = out['profiles'][idx].mean(0)
            print(f"  {profile:12s}: Free={mean_profile[0]:.3f}, "
                  f"Trans={mean_profile[1]:.3f}, "
                  f"Cong={mean_profile[2]:.3f}")
    
    # 训练流程
    print("\n--- 训练流程 ---")
    criterion = nn.CrossEntropyLoss()
    optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3, weight_decay=1e-4)
    
    model.train()
    for epoch in range(5):
        optimizer.zero_grad()
        out = model(x_all)
        
        # 多任务损失：分类 + 预测
        cls_loss = criterion(out['logits'], y_all)
        # 预测损失（下一帧速度）
        pred_loss = F.mse_loss(
            out['predictions'][:, 0], 
            x_all[:, -1, 0]  # 目标速度
        )
        loss = cls_loss + 0.1 * pred_loss
        
        loss.backward()
        optimizer.step()
        
        acc = accuracy_score(y_all.numpy(), 
                            torch.argmax(out['logits'], dim=1).detach().numpy())
        print(f"Epoch {epoch+1}: Total={loss.item():.4f}, "
              f"Cls={cls_loss.item():.4f}, Pred={pred_loss.item():.4f}, "
              f"Acc={acc:.2%}")
    
    # 推理速度
    model.eval()
    with torch.no_grad():
        start = time.time()
        for _ in range(100):
            _ = model(x_all[:1])
        elapsed = (time.time() - start) / 100
        print(f"\n单样本推理延迟: {elapsed*1000:.1f}ms")
        print(f"理论帧率: {1/elapsed:.0f} FPS")
    
    # 密度矩阵可视化
    print("\n--- 密度矩阵分析 ---")
    with torch.no_grad():
        for i, profile in enumerate(PROFILES):
            rho = model.encoder.to_density_matrix(
                x_all[y_all == i, -1, :].mean(0, keepdim=True)
            )
            print(f"\n{profile} 密度矩阵:")
            print(f"  对角线（概率）: {torch.diagonal(rho[0]).numpy().round(3)}")
            print(f"  非对角线（相干）: {rho[0, 0, 1].item():.3f}, "
                  f"{rho[0, 0, 2].item():.3f}")
            print(f"  迹: {torch.trace(rho[0]).item():.3f}")

class DriverBehaviorBaseline:
    """
    利用量子行为模型构建驾驶员个性化基线
    
    应用场景：
    1. 驾驶员状态异常检测（行为偏离基线 → 疲劳/分心/损伤）
    2. 个性化DMS校准（不同驾驶员行为习惯不同）
    3. 行为预测（短期预测用于ADAS协同）
    """
    
    def __init__(self, driver_id: str):
        self.driver_id = driver_id
        self.model = QuantumDriverModel(n_basis=3, n_features=6)
        self.baseline_rho = None  # 驾驶员基线密度矩阵
        self.deviation_threshold = 0.15  # 行为偏离阈值
        
    def calibrate(self, driving_data: torch.Tensor):
        """校准驾驶员基线"""
        with torch.no_grad():
            result = self.model(driving_data)
            self.baseline_rho = result['profiles'].mean(0)
            
    def detect_anomaly(self, current_features: torch.Tensor) -> dict:
        """
        异常检测：当前行为是否偏离基线
        
        Returns:
            {
                'is_anomalous': bool,
                'deviation': float,  # 偏离程度
                'profile': str,     # 当前行为模式
                'risk': str          # 风险等级
            }
        """
        with torch.no_grad():
            result = self.model(current_features.unsqueeze(0))
            current_rho = result['profiles'][0]
            
            if self.baseline_rho is not None:
                deviation = torch.norm(current_rho - self.baseline_rho).item()
            else:
                deviation = 0.0
                
        # 判定
        is_anomalous = deviation > self.deviation_threshold
        profile_idx = torch.argmax(current_rho).item()
        profiles = ['Free Flow', 'Transition', 'Congestion']
        
        risk = 'NORMAL'
        if is_anomalous and deviation > 0.3:
            risk = 'HIGH'  # 可能疲劳/分心/损伤
        elif is_anomalous:
            risk = 'MEDIUM'
            
        return {
            'is_anomalous': is_anomalous,
            'deviation': deviation,
            'profile': profiles[profile_idx],
            'risk': risk
        }