from itertools import product
import numpy as np

class DMSScenarioGenerator:
    """
    DMS 安全关键场景生成器
    
    基于 OmniSCS 的可编辑世界思路
    系统性组合各维度生成测试场景
    """
    
    DIMENSIONS = {
        'driver_state': ['normal', 'light_fatigue', 'moderate_fatigue', 
                         'severe_fatigue', 'phone_distraction', 
                         'cognitive_distraction', 'anger', 'anxiety'],
        'lighting': ['day', 'dusk', 'night', 'backlit', 'tunnel'],
        'weather': ['clear', 'rain', 'fog', 'snow'],
        'road': ['highway', 'urban', 'rural', 'tunnel'],
        'speed': ['stationary', 'low_30', 'mid_60', 'high_120'],
    }
    
    def generate_matrix(self):
        """生成完整场景矩阵"""
        combos = list(product(*self.DIMENSIONS.values()))
        return combos
    
    def priority_filter(self, combos):
        """
        按安全关键性优先级过滤
        
        不是所有组合都需要测试，按风险排序
        """
        # 高风险组合（必须测试）
        high_risk = []
        for c in combos:
            driver, light, weather, road, speed = c
            
            # 疲劳+夜间+高速 = 极高风险
            if 'fatigue' in driver and light == 'night' and speed == 'high_120':
                high_risk.append((c, 'CRITICAL'))
            # 分心+城市+低速 = 行人风险
            elif 'distraction' in driver and road == 'urban' and speed == 'low_30':
                high_risk.append((c, 'HIGH'))
            # 愤怒+高速 = 攻击性驾驶
            elif driver == 'anger' and speed == 'high_120':
                high_risk.append((c, 'HIGH'))
            # 逆光+任何分心
            elif light == 'backlit' and 'distraction' in driver:
                high_risk.append((c, 'HIGH'))
        
        return high_risk
    
    def generate_test_plan(self):
        """生成测试计划"""
        combos = self.generate_matrix()
        high_risk = self.priority_filter(combos)
        
        print(f"总组合数: {len(combos)}")
        print(f"高优先级场景: {len(high_risk)}")
        
        # 按优先级排序
        critical = [c for c, p in high_risk if p == 'CRITICAL']
        high = [c for c, p in high_risk if p == 'HIGH']
        
        print(f"\nCRITICAL ({len(critical)} 场景):")
        for c in critical[:5]:  # 显示前5个
            print(f"  {c}")
        
        print(f"\nHIGH ({len(high)} 场景):")
        for c in high[:5]:
            print(f"  {c}")
        
        return critical + high


# 测试
if __name__ == "__main__":
    gen = DMSScenarioGenerator()
    test_plan = gen.generate_test_plan()