import omni.isaac.core.utils.prims as prim_utils
from omni.isaac.core.simulation_context import SimulationContext
from omni.isaac.core.utils.stage import get_current_stage
from pxr import UsdGeom, UsdPhysics, PhysxSchema, Gf, Sdf

# 引入 PhysX 粒子工具 (Omniverse 内置库)
import omni.physx.scripts.utils as physxUtils
import omni.physx.scripts.particleUtils as particleUtils

def create_fluid_scene():
    stage = get_current_stage()
    
    # 1. 初始化仿真环境
    sim = SimulationContext()
    # 确保物理频率适合流体 (流体通常需要更小的步长或更多的子步)
    sim.set_simulation_dt(physics_dt=1.0 / 60.0, rendering_dt=1.0 / 60.0)

    # 2. 创建粒子系统 (Particle System)
    # 这是管理所有流体粒子的容器
    particle_system_path = "/World/FluidSystem"
    particleUtils.add_physx_particle_system(
        stage=stage,
        particle_system_path=particle_system_path,
        simulation_owner=None,
        radius=0.02,  # 粒子半径（决定流体分辨率）
        contact_offset=0.03,
        solid_rest_offset=0.02,
        fluid_rest_offset=0.02,
        mode="fluid" # 设置为流体模式
    )

    # 3. 创建流体材质 (定义粘度)
    # 对于像蜂蜜或洗发水这样的"带流性"液体，需要高粘度
    material_path = "/World/FluidMaterial"
    particleUtils.add_pbd_particle_material(
        stage=stage,
        path=material_path,
        cohesion=0.02,    # 内聚力 (让液体聚在一起)
        viscosity=0.5,     # 粘度 (关键参数：数值越大越粘稠，如蜂蜜)
        surface_tension=0.01, # 表面张力
        friction=0.1,      # 粒子间摩擦
        damping=0.1       # 阻尼
    )

    # 4. 创建流体几何体 (在试管初始位置生成一堆粒子)
    # 假设我们在 (0, 0, 0.5) 的位置生成一团液体
    fluid_path = "/World/Liquid"
    # 使用网格采样创建粒子点集
    particleUtils.add_physx_particleset_pointinstancer(
        stage=stage,
        path=fluid_path,
        particle_system_path=particle_system_path,
        self_collision=True,
        fluid=True,
        particle_group=0,
        particle_mass=0.001,
        density=1000.0 # 水的密度
    )
    
    # 这里的关键是将材质绑定到几何体
    physxUtils.add_physics_material_to_prim(stage, stage.GetPrimAtPath(fluid_path), material_path)

    # 这里的逻辑通常需要编写一个简单的生成器，在空间中填满粒子
    # 为了演示简单，我们假设通过 create_grid_particles 生成位置
    # 在实际 Isaac Lab 中，你通常会加载一个预先保存好的带粒子的 USD，或者使用 SamplingAPI
    points = []
    for x in range(5):
        for y in range(5):
             for z in range(10):
                 points.append(Gf.Vec3f(x*0.04, y*0.04, 0.5 + z*0.04))
    
    # 将坐标赋予粒子系统
    points_attr = stage.GetPrimAtPath(fluid_path).GetAttribute("positions")
    points_attr.Set(points)

    # 5. 创建容器 (烧杯和试管)
    # 这里使用简单的圆柱管演示，实际应用需加载 .usd 模型
    
    # 烧杯 (底部接收)
    beaker_path = "/World/Beaker"
    prim_utils.create_cyl(
        prim_path=beaker_path,
        radius=0.15,
        height=0.2,
        position=Gf.Vec3d(0, 0, 0.1)
    )
    # 注意：PrimUtils创建的是实心凸包，流体进不去。
    # 关键步骤：必须将容器碰撞改为 Mesh (SDF) 或手动用多个面拼成空心杯子。
    # 在代码中通常加载自定义 USD 资产，这里仅做说明：
    # setup_concave_collision(beaker_path) 

    # 试管 (上方倒水)
    tube_path = "/World/TestTube"
    tube = prim_utils.create_cyl(
        prim_path=tube_path,
        radius=0.05,
        height=0.2,
        position=Gf.Vec3d(0.2, 0, 0.6) # 位于烧杯上方侧面
    )
    # 为试管添加刚体属性以便旋转
    physxUtils.set_rigid_body_properties(stage, tube_path)

    return sim, tube_path

def run_simulation(sim, tube_path):
    # 6. 开始循环并执行“倒水”动作
    sim.initialize_physics()
    sim.play()
    
    stage = get_current_stage()
    tube_prim = stage.GetPrimAtPath(tube_path)
    xform = UsdGeom.Xformable(tube_prim)

    # 模拟 500 帧
    for i in range(500):
        # 简单的运动学控制：慢慢旋转试管
        if i < 200: 
            # 旋转以倒出液体 (欧拉角或四元数)
            rotation = Gf.Rotation(Gf.Vec3d(0, 1, 0), i * 0.5) # 绕Y轴旋转
            # 注意：实际代码需处理完整的 Transform 设置
            current_xform = xform.GetLocalTransformation()
            # 这里简化处理，直接更新姿态逻辑需根据具体场景编写
            
        sim.step()

# 执行
if __name__ == "__main__":
    sim, tube = create_fluid_scene()
    run_simulation(sim, tube)