# Isaac Sim Omniverse Replicator 座舱数据合成脚本
# 在 Isaac Sim Python 环境中运行

from omni.isaac.synthetic_data_generation import (
    SyntheticDataGenerator,
    DomainRandomization
)
from omni.isaac.replicator import (
    Camera, Light, Object,
    Scene, Writer
)
import numpy as np

class CabinDataGenerator:
    """
    座舱 DMS 合成数据生成器
    
    基于 NVIDIA Isaac Sim + Omniverse Replicator
    """
    
    def __init__(self, config: dict = None):
        config = config or {}
        self.output_dir = config.get('output_dir', '/tmp/cabin_synthetic')
        self.num_samples = config.get('num_samples', 10000)
        self.image_size = config.get('image_size', (1920, 1080))
        
        # 座舱场景参数
        self.cabin_params = {
            'camera_position': (0.15, -0.35, 0.65),  # 方向盘右上
            'camera_fov': 60,
            'camera_resolution': self.image_size,
        }
        
        # 人员参数范围
        self.person_params = {
            'gender': ['male', 'female'],
            'age_range': (18, 75),
            'height_range': (1.5, 2.0),  # 米
            'weight_range': (50, 120),    # kg
            'ethnicity': ['asian', 'caucasian', 'african', 'latino'],
        }
        
        # 状态参数
        self.states = [
            'normal_driving', 'drowsy', 'distracted_phone',
            'distracted_radio', 'looking_back', 'yawning',
            'eyes_closed', 'talking', 'smoking'
        ]
    
    def generate_batch(self, num: int = None) -> dict:
        """生成一批合成数据"""
        num = num or self.num_samples
        results = []
        
        for i in range(num):
            sample = self._generate_single()
            results.append(sample)
        
        return {'samples': results, 'count': len(results)}
    
    def _generate_single(self) -> dict:
        """生成单个样本"""
        # 随机化参数
        state = np.random.choice(self.states)
        person = self._randomize_person()
        lighting = self._randomize_lighting()
        camera_pose = self._randomize_camera()
        
        return {
            'state': state,
            'person': person,
            'lighting': lighting,
            'camera': camera_pose,
            'annotations': self._generate_annotations(state, person)
        }
    
    def _randomize_person(self) -> dict:
        return {
            'gender': np.random.choice(self.person_params['gender']),
            'age': np.random.randint(*self.person_params['age_range']),
            'height': np.random.uniform(*self.person_params['height_range']),
            'weight': np.random.uniform(*self.person_params['weight_range']),
            'ethnicity': np.random.choice(self.person_params['ethnicity']),
        }
    
    def _randomize_lighting(self) -> dict:
        return {
            'time_of_day': np.random.choice(
                ['dawn', 'morning', 'noon', 'afternoon', 'dusk', 'night']
            ),
            'intensity': np.random.uniform(50, 1000),  # lux
            'color_temp': np.random.uniform(2700, 6500),  # K
            'direction': np.random.uniform(0, 360),  # 度
        }
    
    def _randomize_camera(self) -> dict:
        return {
            'position': (
                np.random.uniform(0.10, 0.20),
                np.random.uniform(-0.40, -0.30),
                np.random.uniform(0.60, 0.70)
            ),
            'rotation': (
                np.random.uniform(-5, 5),
                np.random.uniform(-5, 5),
                np.random.uniform(-2, 2)
            ),
            'fov': np.random.uniform(55, 65)
        }
    
    def _generate_annotations(self, state: str, person: dict) -> dict:
        """生成自动标注"""
        annotations = {
            'classification': state,
            'bounding_box': self._get_face_bbox(),
            'keypoints': self._get_facial_keypoints(state),
            'gaze': self._get_gaze_direction(state),
            'head_pose': self._get_head_pose(state),
            'facial_attributes': {
                'gender': person['gender'],
                'age_group': self._age_group(person['age']),
                'ethnicity': person['ethnicity']
            }
        }
        return annotations
    
    def _get_face_bbox(self) -> list:
        return [0.3, 0.1, 0.7, 0.5]  # 归一化坐标
    
    def _get_facial_keypoints(self, state: str) -> dict:
        keypoints = {}
        if state == 'eyes_closed' or state == 'drowsy':
            keypoints['eye_openness'] = 0.1  # 几乎闭合
        elif state == 'yawning':
            keypoints['mouth_open'] = 0.8
            keypoints['eye_openness'] = 0.3
        else:
            keypoints['eye_openness'] = 0.8
            keypoints['mouth_open'] = 0.1
        
        return keypoints
    
    def _get_gaze_direction(self, state: str) -> list:
        if 'distracted' in state:
            return [0.8, 0.3, 1.0]  # 偏离前方
        elif state == 'looking_back':
            return [-0.5, 0.0, 0.5]
        else:
            return [0.0, 0.0, 1.0]  # 前方
    
    def _get_head_pose(self, state: str) -> dict:
        if state == 'distracted_phone':
            return {'pitch': 30, 'yaw': -15, 'roll': 0}
        elif state == 'looking_back':
            return {'pitch': 0, 'yaw': 60, 'roll': 0}
        elif state == 'drowsy':
            return {'pitch': 15, 'yaw': 0, 'roll': 5}
        else:
            return {'pitch': 0, 'yaw': 0, 'roll': 0}
    
    def _age_group(self, age: int) -> str:
        if age < 25: return 'young'
        elif age < 45: return 'adult'
        elif age < 65: return 'middle'
        else: return 'senior'


# ============ Isaac Sim Replicator 脚本模板 ============

REPLICATOR_SCRIPT = """
# 在 Isaac Sim Python 环境中运行
# ./python.sh cabin_sdg.py

from omni.isaac.replicator import Randomizer, Writer, Annotator
from omni.isaac.core import World

class CabinDMSDataGenerator:
    def __init__(self):
        self.world = World(stage_units_in_meters=True)
        self.scene = self.world.scene
        
    def setup_scene(self):
        # 加载座舱 USD 资产
        cabin_path = "/assets/cabin_interior.usd"
        self.scene.add_reference_to_stage(cabin_path, "/World/Cabin")
        
        # 加载 Metahuman 驾驶员
        driver_path = "/assets/metahuman_driver.usd"
        self.scene.add_reference_to_stage(driver_path, "/World/Driver")
        
        # 设置 DMS 摄像头
        self.dms_camera = Camera(
            prim_path="/World/DMS_Camera",
            position=(0.15, -0.35, 0.65),
            resolution=(1920, 1080),
            fov=60
        )
        
    def randomize(self):
        # 域随机化
        Randomizer.randomize_materials()
        Randomizer.randomize_lighting(
            intensity_range=(50, 1000),
            color_temp_range=(2700, 6500)
        )
        Randomizer.randomize_human_pose(
            state_pool=['normal', 'drowsy', 'distracted', 'yawning']
        )
        
    def write_data(self, output_dir):
        Writer.write(
            output_dir=output_dir,
            annotators=[
                Annotator.RgbImage(),
                Annotator.SemanticSegmentation(),
                Annotator.BoundingBox2D(),
                Annotator.Keypoints2D(),
                Annotator.InstanceSegmentation(),
            ],
            format='coco'
        )

if __name__ == "__main__":
    gen = CabinDMSDataGenerator()
    gen.setup_scene()
    
    for i in range(10000):
        gen.randomize()
        gen.world.render()
        gen.write_data(f"/output/cabin_dms_batch_{i}")
"""

if __name__ == "__main__":
    # 纯 Python 版本测试（不需要 Isaac Sim）
    generator = CabinDataGenerator(config={
        'num_samples': 100,
        'output_dir': '/tmp/cabin_synthetic'
    })
    
    result = generator.generate_batch(num=5)
    
    print("=== NVIDIA Isaac Sim 座舱数据合成 ===\n")
    for i, sample in enumerate(result['samples']):
        print(f"样本 {i+1}:")
        print(f"  状态: {sample['state']}")
        print(f"  人员: {sample['person']['gender']}, "
              f"{sample['person']['age']}岁, "
              f"{sample['person']['ethnicity']}")
        print(f"  光照: {sample['lighting']['time_of_day']}, "
              f"{sample['lighting']['intensity']:.0f} lux")
        print(f"  注视方向: {sample['annotations']['gaze']}")
        print(f"  头部姿态: {sample['annotations']['head_pose']}")
        print()
    
    print(f"总生成数: {result['count']}")
    print("\n=== Isaac Sim 合成数据管线关键指标 ===")
    print("渲染速度: ~1000 fps (RTX 4090)")
    print("标注成本: $0（自动）")
    print("多样性: 无限（参数化生成）")
    print("输出格式: COCO/KITTI/BDD")
    print("标注类型: RGB+语义分割+边界框+关键点+实例分割")