import numpy as np
import torch
import torch.nn as nn
from torch.utils.data import Dataset, DataLoader
from typing import Tuple

class VirtualCO2Dataset(Dataset):
    """
    虚拟CO₂传感器数据集
    
    论文: 445,017样本, 123小时真实驾驶
    输入: CAN总线信号
    输出: CO₂浓度 (ppm)
    """
    
    # CAN总线信号列表
    CAN_FEATURES = [
        'vehicle_speed',          # km/h
        'acceleration',            # m/s²
        'fan_speed_rpm',          # HVAC风扇
        'hvac_mode',              # 0=outside air, 1=recirculation
        'target_temp_c',          # HVAC目标温度
        'ambient_temp_c',         # 外界温度
        'engine_rpm',             # 发动机转速
        'passenger_count',        # 乘员人数
        'window_status',          # 0=closed, 1=open
    ]
    
    def __init__(self, can_data, co2_labels, physics_equilibrium, 
                 sequence_length=60):
        """
        Args:
            can_data: CAN总线信号, shape=(N, 9)
            co2_labels: CO₂浓度标签, shape=(N,)
            physics_equilibrium: 物理模型输出, shape=(N,)
            sequence_length: 时序窗口（60秒@1Hz）
        """
        self.can_data = torch.FloatTensor(can_data)
        self.co2_labels = torch.FloatTensor(co2_labels)
        self.physics_eq = torch.FloatTensor(physics_equilibrium)
        self.seq_len = sequence_length
        
    def __len__(self):
        return len(self.can_data) - self.seq_len
    
    def __getitem__(self, idx):
        x = self.can_data[idx:idx+self.seq_len]  # (60, 9)
        physics = self.physics_eq[idx:idx+self.seq_len]  # (60,)
        x_with_physics = torch.cat([x, physics.unsqueeze(-1)], dim=-1)  # (60, 10)
        y = self.co2_labels[idx+self.seq_len-1]  # 预测最后时刻CO₂
        return x_with_physics, y


class TCN_LSTM(nn.Module):
    """
    Hybrid TCN-LSTM模型（论文最优架构）
    
    TCN: 时序卷积网络，提取局部模式
    LSTM: 长短期记忆，捕获全局依赖
    混合: TCN输出→LSTM→全连接
    
    论文MAPE: 7.93%（含物理特征）
    论文MAPE: 8.35%（纯数据驱动）
    """
    
    class TemporalBlock(nn.Module):
        def __init__(self, in_ch, out_ch, kernel_size, stride, dilation, 
                     dropout=0.2):
            super().__init__()
            padding = (kernel_size - 1) * dilation
            self.conv = nn.Conv1d(in_ch, out_ch, kernel_size,
                                 stride=stride, padding=padding,
                                 dilation=dilation)
            self.chomp = padding  # causal padding
            self.dropout = nn.Dropout(dropout)
            self.net = nn.Sequential(
                self.conv,
                nn.ReLU(),
                self.dropout
            )
            self.downsample = nn.Conv1d(in_ch, out_ch, 1) if in_ch != out_ch else None
            
        def forward(self, x):
            out = self.net(x)
            out = out[:, :, :-self.chomp].contiguous()  # causal
            res = x if self.downsample is None else self.downsample(x)
            return torch.relu(out + res)
    
    def __init__(self, input_dim=10, hidden_dim=64, n_layers=3, 
                 kernel_size=3, dropout=0.2):
        super().__init__()
        
        # TCN层
        tcn_layers = []
        for i in range(n_layers):
            dilation = 2 ** i
            in_ch = input_dim if i == 0 else hidden_dim
            tcn_layers.append(self.TemporalBlock(
                in_ch, hidden_dim, kernel_size, 1, dilation, dropout
            ))
        self.tcn = nn.Sequential(*tcn_layers)
        
        # LSTM层
        self.lstm = nn.LSTM(
            input_size=hidden_dim,
            hidden_size=hidden_dim,
            num_layers=2,
            batch_first=True,
            dropout=dropout
        )
        
        # 输出层
        self.fc = nn.Sequential(
            nn.Linear(hidden_dim, 32),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(32, 1)
        )
        
    def forward(self, x):
        """
        Args:
            x: CAN信号+物理特征, shape=(B, seq_len, input_dim)
        Returns:
            co2_pred: shape=(B, 1) 单位ppm
        """
        # TCN: (B, seq_len, dim) → (B, dim, seq_len) → (B, hidden, seq_len)
        x_t = x.transpose(1, 2)
        tcn_out = self.tcn(x_t)
        
        # LSTM: (B, hidden, seq_len) → (B, seq_len, hidden)
        lstm_in = tcn_out.transpose(1, 2)
        lstm_out, _ = self.lstm(lstm_in)
        
        # 取最后时刻
        last_out = lstm_out[:, -1, :]
        
        # 预测
        co2_pred = self.fc(last_out)
        return co2_pred


class PhysicsCO2Model:
    """
    物理质量平衡模型
    
    论文方法: 简化质量平衡
    dC/dt = (C_out * Q_in + C_prev * Q_recirc - C * Q_out) / V + G * n
    
    其中:
    - C: 车内CO₂浓度
    - C_out: 外界CO₂浓度 (~400ppm)
    - Q_in/Q_out: 通风流量
    - V: 车厢容积
    - G: 每人CO₂产生率 (~0.005 L/s)
    - n: 乘员数量
    """
    
    def __init__(self):
        self.C_OUTDOOR = 400  # ppm
        self.CABIN_VOLUME_L = 3000  # L (典型轿车)
        self.CO2_PROD_PER_PERSON = 0.005  # L/s
        self.AIR_FLOW_RATE = 50  # L/s (空调中速)
        
    def compute_equilibrium(self, hvac_mode, passenger_count, vehicle_speed):
        """
        计算物理平衡CO₂浓度
        
        Args:
            hvac_mode: 0=外循环, 1=内循环
            passenger_count: 乘员数
            vehicle_speed: 车速 km/h
            
        Returns:
            equilibrium_co2: 平衡CO₂浓度 (ppm)
        """
        # 通风流量（随车速增加）
        q_in = self.AIR_FLOW_RATE + vehicle_speed * 0.5  # L/s
        
        if hvac_mode == 1:  # 内循环
            q_in *= 0.1  # 内循环时外气进气量大幅减少
        
        # CO₂产生率
        generation = self.CO2_PROD_PER_PERSON * passenger_count  # L/s
        
        # 平衡浓度
        if q_in > 0:
            c_eq = self.C_OUTDOOR + generation * self.CABIN_VOLUME_L / q_in
        else:
            c_eq = self.C_OUTDOOR + 1000  # 无通风，快速积累
        
        return min(c_eq, 5000)  # 上限


# 实际测试
if __name__ == "__main__":
    print("=== 虚拟CO₂传感器测试 ===\n")
    
    # 1. 物理模型
    physics = PhysicsCO2Model()
    
    scenarios = [
        ("外循环+1人+60km/h", 0, 1, 60),
        ("内循环+2人+30km/h", 1, 2, 30),
        ("内循环+4人+0km/h(堵车)", 1, 4, 0),
        ("外循环+2人+100km/h", 0, 2, 100),
    ]
    
    print("物理模型平衡CO₂预测:")
    for name, mode, n, speed in scenarios:
        co2 = physics.compute_equilibrium(mode, n, speed)
        print(f"  {name}: {co2:.0f} ppm {'⚠️' if co2 > 1000 else '✅'}")
    
    print(f"\n关键阈值: 1000 ppm")
    print(f"  > 1000 ppm: 驾驶员警觉性下降")
    print(f"  > 1500 ppm: 决策能力显著降低")
    print(f"  > 2500 ppm: 相当于3杯啤酒的认知影响")
    
    print(f"\n=== 论文性能 ===")
    print(f"  纯数据驱动 TCN-LSTM: MAPE = 8.35%")
    print(f"  混合方案 TCN-LSTM+物理: MAPE = 7.93%")
    print(f"  物理模型单独: MAPE = 12-15%")
    print(f"  样本数: 445,017 (123小时)")
    print(f"  采样率: 1 Hz")

class MultimodalFatigueModel:
    """
    多模态疲劳模型：CO₂ + PERCLOS + HRV
    
    虚拟CO₂传感器提供环境维度的疲劳因子
    PERCLOS提供直接视觉疲劳指标
    HRV提供生理疲劳指标
    
    三者融合提高疲劳检测准确性
    """
    
    def __init__(self):
        # 各模态权重（初始值，需实际数据校准）
        self.weights = {
            'perclos': 0.45,     # 视觉指标最直接
            'co2_level': 0.20,  # 环境加速因子
            'hrv': 0.25,        # 生理指标
            'driving_time': 0.10  # 时间累积
        }
        
    def predict_fatigue(self, perclos, co2_ppm, hrv_rmssd, driving_hours):
        """
        多模态疲劳预测
        
        Args:
            perclos: 60秒PERCLOS (%)
            co2_ppm: 虚拟传感器CO₂ (ppm)
            hrv_rmssd: HRV RMSSD (ms)
            driving_hours: 连续驾驶小时
            
        Returns:
            fatigue_score: 0-100, >50为疲劳
        """
        # PERCLOS得分
        perclos_score = min(100, perclos * 2)
        
        # CO₂得分（800-2500ppm线性映射）
        co2_score = max(0, min(100, (co2_ppm - 800) / 17))
        
        # HRV得分（RMSSD降低→疲劳）
        hrv_baseline = 40  # ms, 个体基线
        hrv_score = max(0, min(100, (hrv_baseline - hrv_rmssd) / hrv_baseline * 100))
        
        # 驾驶时间得分
        time_score = min(100, driving_hours / 4 * 100)
        
        # 加权融合
        fatigue_score = (
            self.weights['perclos'] * perclos_score +
            self.weights['co2_level'] * co2_score +
            self.weights['hrv'] * hrv_score +
            self.weights['driving_time'] * time_score
        )
        
        return {
            'fatigue_score': fatigue_score,
            'perclos_score': perclos_score,
            'co2_score': co2_score,
            'hrv_score': hrv_score,
            'time_score': time_score,
            'co2_action': 'switch_to_outside_air' if co2_ppm > 1000 else 'normal'
        }