"""
流体域组件
定义流体模拟的边界和域
"""

from panda3d.core import Vec3, Point3, CollisionNode, CollisionSphere, CollisionBox
import math

class FluidDomain:
    """
    流体域类
    定义流体模拟的边界和域
    """
    
    # 边界类型
    BOUNDARY_BOX = "box"
    BOUNDARY_SPHERE = "sphere"
    BOUNDARY_CYLINDER = "cylinder"
    BOUNDARY_CONE = "cone"
    
    # 边界条件
    CONDITION_CLOSED = "closed"  # 封闭边界
    CONDITION_OPEN = "open"      # 开放边界
    CONDITION_PERIODIC = "periodic"  # 周期性边界
    
    def __init__(self, min_bound, max_bound, boundary_type=BOUNDARY_BOX, boundary_condition=CONDITION_CLOSED):
        """
        初始化流体域
        
        Args:
            min_bound (Point3): 最小边界点
            max_bound (Point3): 最大边界点
            boundary_type (str): 边界类型
            boundary_condition (str): 边界条件
        """
        self.min_bound = min_bound
        self.max_bound = max_bound
        self.size = max_bound - min_bound
        self.center = (min_bound + max_bound) * 0.5
        self.boundary_type = boundary_type
        self.boundary_condition = boundary_condition
        
        # 为不同边界类型计算额外参数
        if boundary_type == self.BOUNDARY_SPHERE:
            # 计算球体半径
            self.radius = max(self.size.getX(), self.size.getY(), self.size.getZ()) / 2.0
        elif boundary_type == self.BOUNDARY_CYLINDER:
            # 计算圆柱体半径和高度
            self.radius = max(self.size.getX(), self.size.getY()) / 2.0
            self.height = self.size.getZ()
        elif boundary_type == self.BOUNDARY_CONE:
            # 计算圆锥体半径和高度
            self.radius = max(self.size.getX(), self.size.getY()) / 2.0
            self.height = self.size.getZ()
            
        # 边界属性
        self.restitution = 0.3  # 恢复系数
        self.friction = 0.1     # 摩擦系数
        self.surface_roughness = 0.01  # 表面粗糙度
        
        # 碰撞检测相关
        self.collision_margin = 0.05  # 碰撞边缘
        self.corner_radius = 0.05     # 角点半径
        
        # 动态边界支持
        self.is_dynamic = False
        self.velocity = Vec3(0, 0, 0)
        self.rotation = Vec3(0, 0, 0)
        
        # 边界厚度
        self.boundary_thickness = 0.05
        
        print(f"流体域创建完成: {min_bound} to {max_bound}, 类型: {boundary_type}, 条件: {boundary_condition}")
        
    def update(self, dt):
        """
        更新流体域状态
        
        Args:
            dt (float): 时间步长
        """
        if self.is_dynamic:
            # 更新动态边界的位置和旋转
            self.center += self.velocity * dt
            
            # 在这里可以添加边界旋转的逻辑
            # self.rotation += self.angular_velocity * dt
            
    def handle_particle_collision(self, particle):
        """
        处理粒子与域边界的碰撞
        
        Args:
            particle: 流体粒子对象（字典格式）
        """
        pos = particle['position']
        
        if self.boundary_type == self.BOUNDARY_BOX:
            self._handle_box_collision(particle)
        elif self.boundary_type == self.BOUNDARY_SPHERE:
            self._handle_sphere_collision(particle)
        elif self.boundary_type == self.BOUNDARY_CYLINDER:
            self._handle_cylinder_collision(particle)
        elif self.boundary_type == self.BOUNDARY_CONE:
            self._handle_cone_collision(particle)
            
    def _handle_box_collision(self, particle):
        """
        处理盒状边界碰撞
        
        Args:
            particle: 流体粒子对象
        """
        pos = particle['position']
        vel = particle['velocity']
        
        # 获取边界值
        min_x, max_x = self.min_bound.getX(), self.max_bound.getX()
        min_y, max_y = self.min_bound.getY(), self.max_bound.getY()
        min_z, max_z = self.min_bound.getZ(), self.max_bound.getZ()
        
        # 检查X轴边界
        if pos.getX() <= min_x + self.collision_margin:
            pos.setX(min_x + self.collision_margin)
            vel.setX(-vel.getX() * (1 - self.friction) * self.restitution)
        elif pos.getX() >= max_x - self.collision_margin:
            pos.setX(max_x - self.collision_margin)
            vel.setX(-vel.getX() * (1 - self.friction) * self.restitution)
            
        # 检查Y轴边界
        if pos.getY() <= min_y + self.collision_margin:
            pos.setY(min_y + self.collision_margin)
            vel.setY(-vel.getY() * (1 - self.friction) * self.restitution)
        elif pos.getY() >= max_y - self.collision_margin:
            pos.setY(max_y - self.collision_margin)
            vel.setY(-vel.getY() * (1 - self.friction) * self.restitution)
            
        # 检查Z轴边界
        if pos.getZ() <= min_z + self.collision_margin:
            pos.setZ(min_z + self.collision_margin)
            vel.setZ(-vel.getZ() * (1 - self.friction) * self.restitution)
        elif pos.getZ() >= max_z - self.collision_margin:
            pos.setZ(max_z - self.collision_margin)
            vel.setZ(-vel.getZ() * (1 - self.friction) * self.restitution)
            
        # 处理周期性边界条件
        if self.boundary_condition == self.CONDITION_PERIODIC:
            self._apply_periodic_boundary_conditions(particle)
            
    def _handle_sphere_collision(self, particle):
        """
        处理球状边界碰撞
        
        Args:
            particle: 流体粒子对象
        """
        pos = particle['position']
        vel = particle['velocity']
        
        # 计算粒子到球心的距离
        center_to_particle = pos - self.center
        distance = center_to_particle.length()
        
        # 检查是否超出球体边界
        if distance > self.radius - self.collision_margin:
            # 计算碰撞点
            normal = center_to_particle.normalized()
            collision_point = self.center + normal * (self.radius - self.collision_margin)
            
            # 调整粒子位置
            particle['position'] = collision_point
            
            # 计算碰撞响应
            v_normal = normal.dot(vel) * normal  # 法向速度
            v_tangent = vel - v_normal          # 切向速度
            
            # 应用反弹和摩擦
            new_v_normal = -v_normal * self.restitution
            new_v_tangent = v_tangent * (1 - self.friction)
            
            particle['velocity'] = new_v_normal + new_v_tangent
            
    def _handle_cylinder_collision(self, particle):
        """
        处理圆柱边界碰撞
        
        Args:
            particle: 流体粒子对象
        """
        pos = particle['position']
        vel = particle['velocity']
        
        # 计算在X-Y平面上到圆柱轴的距离
        dx = pos.getX() - self.center.getX()
        dy = pos.getY() - self.center.getY()
        distance_xy = math.sqrt(dx*dx + dy*dy)
        
        # 检查径向边界
        if distance_xy > self.radius - self.collision_margin:
            # 计算径向方向
            radial_dir = Vec3(dx, dy, 0).normalized()
            
            # 调整位置
            new_x = self.center.getX() + radial_dir.getX() * (self.radius - self.collision_margin)
            new_y = self.center.getY() + radial_dir.getY() * (self.radius - self.collision_margin)
            pos.setX(new_x)
            pos.setY(new_y)
            
            # 计算碰撞响应
            v_radial = radial_dir.dot(vel) * radial_dir
            v_tangent = vel - v_radial
            
            # 应用反弹和摩擦
            new_v_radial = -v_radial * self.restitution
            particle['velocity'] = new_v_radial + v_tangent * (1 - self.friction)
            
        # 检查Z轴边界
        min_z = self.center.getZ() - self.height/2.0
        max_z = self.center.getZ() + self.height/2.0
        
        if pos.getZ() <= min_z + self.collision_margin:
            pos.setZ(min_z + self.collision_margin)
            vel.setZ(-vel.getZ() * (1 - self.friction) * self.restitution)
        elif pos.getZ() >= max_z - self.collision_margin:
            pos.setZ(max_z - self.collision_margin)
            vel.setZ(-vel.getZ() * (1 - self.friction) * self.restitution)
            
    def _handle_cone_collision(self, particle):
        """
        处理圆锥边界碰撞
        
        Args:
            particle: 流体粒子对象
        """
        pos = particle['position']
        vel = particle['velocity']
        
        # 计算相对于圆锥顶点的位置
        rel_pos = pos - self.center
        rel_pos.setZ(rel_pos.getZ() + self.height/2)  # 调整Z轴使顶点在底部
        
        # 计算在X-Y平面上到圆锥轴的距离
        distance_xy = math.sqrt(rel_pos.getX()**2 + rel_pos.getY()**2)
        
        # 计算圆锥在当前高度的半径
        cone_radius_at_height = (rel_pos.getZ() / self.height) * self.radius
        
        # 检查是否超出圆锥边界
        if distance_xy > cone_radius_at_height - self.collision_margin and rel_pos.getZ() > 0:
            # 计算圆锥表面的法向量
            radial_dir = Vec3(rel_pos.getX(), rel_pos.getY(), 0).normalized()
            surface_normal = Vec3(
                radial_dir.getX(),
                radial_dir.getY(),
                -self.radius / self.height  # 圆锥的斜率
            ).normalized()
            
            # 调整粒子位置到表面
            adjusted_radius = cone_radius_at_height - self.collision_margin
            particle['position'] = Point3(
                self.center.getX() + radial_dir.getX() * adjusted_radius,
                self.center.getY() + radial_dir.getY() * adjusted_radius,
                self.center.getZ() + rel_pos.getZ() - self.height/2
            )
            
            # 计算碰撞响应
            v_normal = surface_normal.dot(vel) * surface_normal
            v_tangent = vel - v_normal
            
            # 应用反弹和摩擦
            new_v_normal = -v_normal * self.restitution
            particle['velocity'] = new_v_normal + v_tangent * (1 - self.friction)
            
        # 检查Z轴边界
        if rel_pos.getZ() < 0:  # 底部
            pos.setZ(self.center.getZ() - self.height/2 + self.collision_margin)
            vel.setZ(-vel.getZ() * (1 - self.friction) * self.restitution)
        elif rel_pos.getZ() > self.height:  # 顶部（圆锥尖端）
            pos.setZ(self.center.getZ() + self.height/2 - self.collision_margin)
            vel.setZ(-vel.getZ() * (1 - self.friction) * self.restitution)
            
    def _apply_periodic_boundary_conditions(self, particle):
        """
        应用周期性边界条件
        
        Args:
            particle: 流体粒子对象
        """
        pos = particle['position']
        
        # X轴周期性
        if pos.getX() < self.min_bound.getX():
            particle['position'].setX(self.max_bound.getX() - (self.min_bound.getX() - pos.getX()))
        elif pos.getX() > self.max_bound.getX():
            particle['position'].setX(self.min_bound.getX() + (pos.getX() - self.max_bound.getX()))
            
        # Y轴周期性
        if pos.getY() < self.min_bound.getY():
            particle['position'].setY(self.max_bound.getY() - (self.min_bound.getY() - pos.getY()))
        elif pos.getY() > self.max_bound.getY():
            particle['position'].setY(self.min_bound.getY() + (pos.getY() - self.max_bound.getY()))
            
        # Z轴周期性
        if pos.getZ() < self.min_bound.getZ():
            particle['position'].setZ(self.max_bound.getZ() - (self.min_bound.getZ() - pos.getZ()))
        elif pos.getZ() > self.max_bound.getZ():
            particle['position'].setZ(self.min_bound.getZ() + (pos.getZ() - self.max_bound.getZ()))
            
    def is_point_inside(self, point):
        """
        检查点是否在域内
        
        Args:
            point (Point3): 检查点
            
        Returns:
            bool: 是否在域内
        """
        if self.boundary_type == self.BOUNDARY_BOX:
            return (self.min_bound.getX() <= point.getX() <= self.max_bound.getX() and
                    self.min_bound.getY() <= point.getY() <= self.max_bound.getY() and
                    self.min_bound.getZ() <= point.getZ() <= self.max_bound.getZ())
        elif self.boundary_type == self.BOUNDARY_SPHERE:
            distance = (point - self.center).length()
            return distance <= self.radius
        elif self.boundary_type == self.BOUNDARY_CYLINDER:
            dx = point.getX() - self.center.getX()
            dy = point.getY() - self.center.getY()
            distance_xy = math.sqrt(dx*dx + dy*dy)
            z_in_range = self.center.getZ() - self.height/2 <= point.getZ() <= self.center.getZ() + self.height/2
            return distance_xy <= self.radius and z_in_range
        elif self.boundary_type == self.BOUNDARY_CONE:
            rel_point = point - self.center
            rel_point.setZ(rel_point.getZ() + self.height/2)
            dx = rel_point.getX()
            dy = rel_point.getY()
            distance_xy = math.sqrt(dx*dx + dy*dy)
            cone_radius_at_height = (rel_point.getZ() / self.height) * self.radius
            return distance_xy <= cone_radius_at_height and 0 <= rel_point.getZ() <= self.height
            
        return False  # 默认返回False
                
    def get_volume(self):
        """
        计算域体积
        
        Returns:
            float: 域体积
        """
        if self.boundary_type == self.BOUNDARY_BOX:
            return self.size.getX() * self.size.getY() * self.size.getZ()
        elif self.boundary_type == self.BOUNDARY_SPHERE:
            return (4.0/3.0) * math.pi * (self.radius ** 3)
        elif self.boundary_type == self.BOUNDARY_CYLINDER:
            return math.pi * (self.radius ** 2) * self.height
        elif self.boundary_type == self.BOUNDARY_CONE:
            return (1.0/3.0) * math.pi * (self.radius ** 2) * self.height
        else:
            # 默认按盒状计算
            return self.size.getX() * self.size.getY() * self.size.getZ()
            
    def set_boundary_condition(self, condition):
        """
        设置边界条件
        
        Args:
            condition (str): 边界条件类型
        """
        self.boundary_condition = condition
        
    def set_material_properties(self, restitution=None, friction=None, surface_roughness=None):
        """
        设置边界材料属性
        
        Args:
            restitution (float): 恢复系数
            friction (float): 摩擦系数
            surface_roughness (float): 表面粗糙度
        """
        if restitution is not None:
            self.restitution = max(0.0, min(1.0, restitution))
        if friction is not None:
            self.friction = max(0.0, min(1.0, friction))
        if surface_roughness is not None:
            self.surface_roughness = max(0.0, surface_roughness)
            
    def set_dynamic(self, is_dynamic, velocity=None, rotation=None):
        """
        设置边界是否为动态
        
        Args:
            is_dynamic (bool): 是否为动态边界
            velocity (Vec3): 速度
            rotation (Vec3): 旋转
        """
        self.is_dynamic = is_dynamic
        if velocity:
            self.velocity = velocity
        if rotation:
            self.rotation = rotation
            
    def get_surface_area(self):
        """
        计算边界表面积
        
        Returns:
            float: 表面积
        """
        if self.boundary_type == self.BOUNDARY_BOX:
            return 2 * (self.size.getX() * self.size.getY() + 
                       self.size.getX() * self.size.getZ() + 
                       self.size.getY() * self.size.getZ())
        elif self.boundary_type == self.BOUNDARY_SPHERE:
            return 4 * math.pi * (self.radius ** 2)
        elif self.boundary_type == self.BOUNDARY_CYLINDER:
            return 2 * math.pi * self.radius * (self.radius + self.height)
        elif self.boundary_type == self.BOUNDARY_CONE:
            slant_height = math.sqrt(self.radius**2 + self.height**2)
            return math.pi * self.radius * (self.radius + slant_height)
        else:
            # 默认按盒状计算
            return 2 * (self.size.getX() * self.size.getY() + 
                       self.size.getX() * self.size.getZ() + 
                       self.size.getY() * self.size.getZ())
                       
    def get_bounding_box(self):
        """
        获取边界框
        
        Returns:
            tuple: (min_bound, max_bound)
        """
        return (self.min_bound, self.max_bound)
        
    def get_boundary_mesh(self):
        """
        获取边界网格（用于可视化）
        
        Returns:
            dict: 包含网格顶点和面的字典
        """
        if self.boundary_type == self.BOUNDARY_BOX:
            return self._get_box_mesh()
        elif self.boundary_type == self.BOUNDARY_SPHERE:
            return self._get_sphere_mesh()
        elif self.boundary_type == self.BOUNDARY_CYLINDER:
            return self._get_cylinder_mesh()
        elif self.boundary_type == self.BOUNDARY_CONE:
            return self._get_cone_mesh()
        else:
            return self._get_box_mesh()  # 默认返回盒状网格
            
    def _get_box_mesh(self):
        """
        获取盒状边界网格
        
        Returns:
            dict: 网格数据
        """
        vertices = [
            Point3(self.min_bound.getX(), self.min_bound.getY(), self.min_bound.getZ()),
            Point3(self.max_bound.getX(), self.min_bound.getY(), self.min_bound.getZ()),
            Point3(self.max_bound.getX(), self.max_bound.getY(), self.min_bound.getZ()),
            Point3(self.min_bound.getX(), self.max_bound.getY(), self.min_bound.getZ()),
            Point3(self.min_bound.getX(), self.min_bound.getY(), self.max_bound.getZ()),
            Point3(self.max_bound.getX(), self.min_bound.getY(), self.max_bound.getZ()),
            Point3(self.max_bound.getX(), self.max_bound.getY(), self.max_bound.getZ()),
            Point3(self.min_bound.getX(), self.max_bound.getY(), self.max_bound.getZ()),
        ]
        
        faces = [
            [0, 1, 2, 3],  # 底面
            [4, 7, 6, 5],  # 顶面
            [0, 4, 5, 1],  # 前面
            [2, 6, 7, 3],  # 后面
            [0, 3, 7, 4],  # 左面
            [1, 5, 6, 2],  # 右面
        ]
        
        return {"vertices": vertices, "faces": faces}
        
    def _get_sphere_mesh(self, resolution=16):
        """
        获取球状边界网格
        
        Args:
            resolution (int): 网格分辨率
            
        Returns:
            dict: 网格数据
        """
        vertices = []
        faces = []
        
        # 生成球面顶点
        for i in range(resolution + 1):
            phi = math.pi * i / resolution  # 0 到 π
            for j in range(resolution + 1):
                theta = 2 * math.pi * j / resolution  # 0 到 2π
                
                x = self.center.getX() + self.radius * math.sin(phi) * math.cos(theta)
                y = self.center.getY() + self.radius * math.sin(phi) * math.sin(theta)
                z = self.center.getZ() + self.radius * math.cos(phi)
                
                vertices.append(Point3(x, y, z))
                
        # 生成面
        for i in range(resolution):
            for j in range(resolution):
                first = i * (resolution + 1) + j
                second = first + resolution + 1
                
                # 添加两个三角形形成一个四边形面
                faces.append([first, second, first + 1])
                faces.append([second, second + 1, first + 1])
                
        return {"vertices": vertices, "faces": faces}
        
    def _get_cylinder_mesh(self, resolution=16):
        """
        获取圆柱边界网格
        
        Args:
            resolution (int): 网格分辨率
            
        Returns:
            dict: 网格数据
        """
        vertices = []
        faces = []
        
        # 生成底部圆环
        base_center = Point3(self.center.getX(), self.center.getY(), self.center.getZ() - self.height/2)
        for i in range(resolution):
            theta = 2 * math.pi * i / resolution
            x = base_center.getX() + self.radius * math.cos(theta)
            y = base_center.getY() + self.radius * math.sin(theta)
            z = base_center.getZ()
            vertices.append(Point3(x, y, z))
            
        # 添加底部中心点
        vertices.append(base_center)
        
        # 生成顶部圆环
        top_center = Point3(self.center.getX(), self.center.getY(), self.center.getZ() + self.height/2)
        for i in range(resolution):
            theta = 2 * math.pi * i / resolution
            x = top_center.getX() + self.radius * math.cos(theta)
            y = top_center.getY() + self.radius * math.sin(theta)
            z = top_center.getZ()
            vertices.append(Point3(x, y, z))
            
        # 添加顶部中心点
        vertices.append(top_center)
        
        # 生成侧面
        for i in range(resolution):
            next_i = (i + 1) % resolution
            # 底面
            faces.append([i, resolution, (next_i if next_i != 0 else resolution-1)])
            # 顶面
            faces.append([i + resolution + 1, resolution * 2 + 1, (next_i + resolution + 1) if next_i != 0 else resolution])
            # 侧面
            faces.append([i, next_i, i + resolution + 1])
            faces.append([next_i, next_i + resolution + 1, i + resolution + 1])
            
        return {"vertices": vertices, "faces": faces}
        
    def _get_cone_mesh(self, resolution=16):
        """
        获取圆锥边界网格
        
        Args:
            resolution (int): 网格分辨率
            
        Returns:
            dict: 网格数据
        """
        vertices = []
        faces = []
        
        # 生成底部圆环
        base_center = Point3(self.center.getX(), self.center.getY(), self.center.getZ() - self.height/2)
        for i in range(resolution):
            theta = 2 * math.pi * i / resolution
            x = base_center.getX() + self.radius * math.cos(theta)
            y = base_center.getY() + self.radius * math.sin(theta)
            z = base_center.getZ()
            vertices.append(Point3(x, y, z))
            
        # 添加底部中心点
        vertices.append(base_center)
        
        # 圆锥顶点
        apex = Point3(self.center.getX(), self.center.getY(), self.center.getZ() + self.height/2)
        vertices.append(apex)
        
        # 生成侧面
        for i in range(resolution):
            next_i = (i + 1) % resolution
            # 底面
            faces.append([i, resolution, (next_i if next_i != 0 else resolution-1)])
            # 侧面
            faces.append([i, (next_i if next_i != 0 else resolution-1), resolution + 1])
            
        return {"vertices": vertices, "faces": faces}