"""
流体求解器模块
负责流体动力学方程的求解，实现基于SPH的流体模拟
"""

from panda3d.core import Vec3, Point3
import math
import numpy as np
from scipy.spatial import cKDTree

class FluidSolver:
    """
    流体求解器类
    实现基于SPH（光滑粒子流体动力学）的流体模拟
    """
    
    def __init__(self, fluid_world):
        """
        初始化流体求解器
        
        Args:
            fluid_world: 流体世界对象的引用
        """
        self.fluid_world = fluid_world
        
        # SPH参数
        self.smoothing_radius = 1.0  # 光滑半径
        self.rest_density = 1000.0   # 静止密度
        self.stiffness = 1000.0      # 刚度系数
        self.viscosity = 0.001       # 粘度
        self.surface_tension = 0.01  # 表面张力
        self.gravity = Vec3(0, 0, -9.81)  # 重力
        
        # 数值参数
        self.time_step = 0.016       # 时间步长
        self.speed_limit = 100.0     # 速度限制
        
        # 粒子数据结构
        self.particle_positions = []  # 粒子位置数组
        self.particle_velocities = []  # 粒子速度数组
        self.particle_densities = []   # 粒子密度数组
        self.particle_pressures = []   # 粒子压力数组
        
        # KD树用于近邻搜索
        self.kd_tree = None
        self.particle_indices = []
        
        print("SPH流体求解器初始化完成")
        
    def set_smoothing_radius(self, radius):
        """
        设置光滑半径
        
        Args:
            radius (float): 光滑半径
        """
        self.smoothing_radius = radius
        
    def set_rest_density(self, density):
        """
        设置静止密度
        
        Args:
            density (float): 静止密度
        """
        self.rest_density = density
        
    def set_stiffness(self, stiffness):
        """
        设置刚度系数
        
        Args:
            stiffness (float): 刚度系数
        """
        self.stiffness = stiffness
        
    def set_viscosity(self, viscosity):
        """
        设置流体粘度
        
        Args:
            viscosity (float): 粘度值
        """
        self.viscosity = viscosity
        
    def set_surface_tension(self, tension):
        """
        设置表面张力
        
        Args:
            tension (float): 表面张力
        """
        self.surface_tension = tension
        
    def set_gravity(self, gravity):
        """
        设置重力
        
        Args:
            gravity (Vec3): 重力向量
        """
        self.gravity = gravity
        
    def solve(self, dt):
        """
        执行流体求解
        
        Args:
            dt (float): 时间步长
        """
        # 更新粒子数据结构
        self._update_particle_data()
        
        # 构建KD树用于近邻搜索
        self._build_kd_tree()
        
        # 计算粒子密度
        self._calculate_densities()
        
        # 计算粒子压力
        self._calculate_pressures()
        
        # 计算作用力
        self._calculate_forces()
        
        # 更新速度
        self._update_velocities(dt)
        
        # 处理碰撞
        self._handle_collisions()
        
        # 更新位置
        self._update_positions(dt)
        
        # 处理边界条件
        self._handle_boundary_conditions()
        
    def _update_particle_data(self):
        """
        更新粒子数据结构
        """
        # 清空旧数据
        self.particle_positions.clear()
        self.particle_velocities.clear()
        self.particle_densities.clear()
        self.particle_pressures.clear()
        self.particle_indices.clear()
        
        # 填充新数据
        for i, particle in enumerate(self.fluid_world.fluid_particles):
            self.particle_positions.append(particle['position'])
            self.particle_velocities.append(particle['velocity'])
            self.particle_densities.append(0.0)
            self.particle_pressures.append(0.0)
            self.particle_indices.append(i)
            
    def _build_kd_tree(self):
        """
        构建KD树用于近邻搜索
        """
        if len(self.particle_positions) > 0:
            # 将位置转换为numpy数组
            positions = np.array([[p.getX(), p.getY(), p.getZ()] for p in self.particle_positions])
            self.kd_tree = cKDTree(positions)
        else:
            self.kd_tree = None
            
    def _calculate_densities(self):
        """
        计算每个粒子的密度
        """
        if not self.kd_tree or len(self.particle_positions) == 0:
            return
            
        # 对每个粒子计算密度
        for i, position in enumerate(self.particle_positions):
            # 使用KD树查找近邻粒子
            neighbors = self.kd_tree.query_ball_point(
                [position.getX(), position.getY(), position.getZ()], 
                self.smoothing_radius
            )
            
            # 计算密度
            density = 0.0
            for j in neighbors:
                # 计算距离
                r = (position - self.particle_positions[j]).length()
                if r < self.smoothing_radius:
                    # 使用SPH核函数计算贡献
                    density += self._sph_kernel(r)
                    
            # 乘以粒子质量（假设所有粒子质量相等）
            self.particle_densities[i] = density * 1.0  # 假设粒子质量为1.0
            
    def _calculate_pressures(self):
        """
        计算每个粒子的压力
        """
        # 使用状态方程计算压力
        for i, density in enumerate(self.particle_densities):
            # 状态方程: P = k * (ρ - ρ₀)
            pressure = self.stiffness * (density - self.rest_density)
            # 确保压力非负
            self.particle_pressures[i] = max(0.0, pressure)
            
    def _calculate_forces(self):
        """
        计算作用在每个粒子上的力
        """
        if not self.kd_tree or len(self.particle_positions) == 0:
            return
            
        # 对每个粒子计算力
        for i, position in enumerate(self.particle_positions):
            # 初始化力
            pressure_force = Vec3(0, 0, 0)
            viscosity_force = Vec3(0, 0, 0)
            surface_tension_force = Vec3(0, 0, 0)
            gravity_force = self.gravity * self.particle_densities[i]  # F = m * g
            
            # 查找近邻粒子
            neighbors = self.kd_tree.query_ball_point(
                [position.getX(), position.getY(), position.getZ()], 
                self.smoothing_radius
            )
            
            # 计算与近邻粒子的相互作用力
            for j in neighbors:
                if i == j:  # 跳过自己
                    continue
                    
                # 计算相对位置和距离
                r_vec = position - self.particle_positions[j]
                r = r_vec.length()
                
                if 0 < r < self.smoothing_radius:
                    # 归一化方向向量
                    r_norm = r_vec / r if r > 0 else Vec3(0, 0, 0)
                    
                    # 计算压力力
                    pressure_force += self._pressure_force(i, j, r, r_norm)
                    
                    # 计算粘性力
                    viscosity_force += self._viscosity_force(i, j, r)
                    
                    # 计算表面张力
                    surface_tension_force += self._surface_tension_force(i, j, r, r_norm)
                    
            # 合成总力
            total_force = pressure_force + viscosity_force + surface_tension_force + gravity_force
            
            # 应用力到粒子（F = m * a => a = F / m）
            if self.particle_densities[i] > 0:
                acceleration = total_force / self.particle_densities[i]
                # 限制加速度
                if acceleration.length() > 1000.0:
                    acceleration.normalize()
                    acceleration *= 1000.0
                    
                # 更新粒子速度（临时存储）
                self.particle_velocities[i] += acceleration * self.time_step
                
    def _pressure_force(self, i, j, r, r_norm):
        """
        计算压力力
        
        Args:
            i (int): 粒子i索引
            j (int): 粒子j索引
            r (float): 距离
            r_norm (Vec3): 归一化方向向量
            
        Returns:
            Vec3: 压力力
        """
        # SPH压力梯度核函数
        pressure_grad = self._sph_pressure_kernel_gradient(r)
        
        # 压力力计算
        force = -r_norm * pressure_grad * (
            self.particle_pressures[i] + self.particle_pressures[j]
        ) / (2.0 * self.particle_densities[j])
        
        return force
        
    def _viscosity_force(self, i, j, r):
        """
        计算粘性力（使用Monaghan粘性模型）
        
        Args:
            i (int): 粒子i索引
            j (int): 粒子j索引
            r (float): 距离
            
        Returns:
            Vec3: 粘性力
        """
        # 确保不计算自己与自己的相互作用
        if i == j:
            return Vec3(0, 0, 0)
            
        # 获取粒子速度
        velocity_i = self.particle_velocities[i]
        velocity_j = self.particle_velocities[j]
        
        # 计算速度差
        velocity_diff = velocity_j - velocity_i
        
        # 使用Monaghan粘性模型
        # 核函数拉普拉斯近似
        viscosity_laplacian = self._sph_viscosity_kernel_laplacian(r)
        
        # 粘性力计算
        # F_viscosity = μ * (v_j - v_i) * ∇²W_ij / ρ_j
        force_magnitude = self.viscosity * viscosity_laplacian / self.particle_densities[j]
        force = velocity_diff * force_magnitude
        
        return force
        
    def _surface_tension_force(self, i, j, r, r_norm):
        """
        计算表面张力（使用CSF模型 - Continuum Surface Force）
        
        Args:
            i (int): 粒子i索引
            j (int): 粒子j索引
            r (float): 距离
            r_norm (Vec3): 归一化方向向量
            
        Returns:
            Vec3: 表面张力
        """
        # 为避免除零错误，确保距离不为零
        if r <= 0:
            return Vec3(0, 0, 0)
            
        # 计算表面张力需要更复杂的算法，这里使用简化的模型
        # 首先计算密度梯度来估计表面
        density_gradient = self._calculate_density_gradient(i, j, r, r_norm)
        
        # 使用连续表面力模型 (CSF)
        # 表面张力力 = σ * κ * n * δ(Γ)
        # 其中 σ 是表面张力系数，κ 是曲率，n 是表面法向量，δ(Γ) 是界面delta函数
        
        # 简化的表面张力计算
        # 基于密度变化来估算是否在表面
        density_diff = abs(self.particle_densities[i] - self.particle_densities[j])
        
        # 只有在密度变化显著时才应用表面张力（即在表面附近）
        if density_diff > self.rest_density * 0.1:  # 表面判断阈值
            # 计算表面曲率（简化计算）
            curvature = self._estimate_curvature(i, j, r, r_norm)
            
            # 表面张力力计算
            force = r_norm * self.surface_tension * curvature
        else:
            # 不在表面附近，表面张力影响较小
            force = Vec3(0, 0, 0)
        
        return force
        
    def _calculate_density_gradient(self, i, j, r, r_norm):
        """
        计算密度梯度
        
        Args:
            i (int): 粒子i索引
            j (int): 粒子j索引
            r (float): 距离
            r_norm (Vec3): 归一化方向向量
            
        Returns:
            Vec3: 密度梯度
        """
        # 密度梯度 = Σ m_j * (ρ_j - ρ_i) / ρ_j * ∇W_ij
        if r <= 0:
            return Vec3(0, 0, 0)
            
        # 获取粒子i和j的密度
        rho_i = self.particle_densities[i]
        rho_j = self.particle_densities[j]
        
        # 计算SPH核函数梯度
        grad_w = r_norm * self._sph_kernel_gradient(r)
        
        # 密度梯度的近似计算
        if rho_j > 0:
            density_grad = (rho_j - rho_i) / rho_j * grad_w
        else:
            density_grad = Vec3(0, 0, 0)
            
        return density_grad
        
    def _estimate_curvature(self, i, j, r, r_norm):
        """
        估算表面曲率
        
        Args:
            i (int): 粒子i索引
            j (int): 粒子j索引
            r (float): 距离
            r_norm (Vec3): 归一化方向向量
            
        Returns:
            float: 曲率值
        """
        # 简化的曲率估算方法
        # 基于密度梯度的幅度
        density_grad = self._calculate_density_gradient(i, j, r, r_norm)
        grad_magnitude = density_grad.length()
        
        # 基于密度梯度估算曲率
        # 这只是一个近似，实际的曲率计算更复杂
        if grad_magnitude > 0:
            # 使用密度梯度幅度作为曲率的粗略估计
            # 在更精确的实现中，这将涉及计算拉普拉斯算子
            curvature = grad_magnitude / r if r > 0 else 0
        else:
            curvature = 0
            
        return curvature
        
    def _sph_kernel(self, r):
        """
        SPH核函数 (Poly6 kernel)
        
        Args:
            r (float): 距离
            
        Returns:
            float: 核函数值
        """
        if r >= self.smoothing_radius or r < 0:
            return 0.0
            
        h = self.smoothing_radius
        coefficient = 315.0 / (64.0 * math.pi * h**9)
        return coefficient * (h*h - r*r)**3
        
    def _sph_kernel_gradient(self, r):
        """
        SPH核函数梯度 (Poly6 kernel gradient)
        
        Args:
            r (float): 距离
            
        Returns:
            float: 核函数梯度值
        """
        if r >= self.smoothing_radius or r <= 0:
            return 0.0
            
        h = self.smoothing_radius
        coefficient = -945.0 / (32.0 * math.pi * h**9)
        return coefficient * (h*h - r*r)**2 * r
        
    def _sph_pressure_kernel_gradient(self, r):
        """
        SPH压力核函数梯度 (Spiky kernel gradient)
        这是用于压力力计算的核函数
        
        Args:
            r (float): 距离
            
        Returns:
            float: 核函数梯度值
        """
        if r >= self.smoothing_radius or r <= 0:
            return 0.0
            
        h = self.smoothing_radius
        coefficient = -45.0 / (math.pi * h**6)
        
        # 确保核函数值在定义域内
        if r <= h:
            return coefficient * (h - r)**2
        else:
            return 0.0
        
    def _sph_viscosity_kernel_laplacian(self, r):
        """
        SPH粘性核函数拉普拉斯 (Viscosity kernel laplacian)
        用于粘性力计算的核函数
        
        Args:
            r (float): 距离
            
        Returns:
            float: 核函数拉普拉斯值
        """
        if r >= self.smoothing_radius or r < 0:
            return 0.0
            
        h = self.smoothing_radius
        coefficient = 45.0 / (math.pi * h**6)
        
        # 确保核函数值在定义域内
        if r <= h:
            return coefficient * (h - r)
        else:
            return 0.0
        
    def _sph_surface_tension_kernel(self, r):
        """
        SPH表面张力核函数 (Poly6 kernel)
        用于表面张力计算的核函数
        
        Args:
            r (float): 距离
            
        Returns:
            float: 核函数值
        """
        if r >= self.smoothing_radius or r < 0:
            return 0.0
            
        h = self.smoothing_radius
        coefficient = 315.0 / (64.0 * math.pi * h**9)
        
        # 确保核函数值在定义域内
        if r <= h:
            return coefficient * (h*h - r*r)**3
        else:
            return 0.0
        
    def _update_velocities(self, dt):
        """
        更新粒子速度
        
        Args:
            dt (float): 时间步长
        """
        # 将计算得到的速度应用到粒子
        for i, particle in enumerate(self.fluid_world.fluid_particles):
            # 限制速度
            if self.particle_velocities[i].length() > self.speed_limit:
                self.particle_velocities[i].normalize()
                self.particle_velocities[i] *= self.speed_limit
                
            particle['velocity'] = self.particle_velocities[i]
            
    def _handle_collisions(self):
        """
        处理粒子与边界和其他粒子的碰撞
        """
        for particle in self.fluid_world.fluid_particles:
            # 与边界碰撞检测
            for domain in self.fluid_world.fluid_domains:
                domain.handle_particle_collision(particle)
                
            # 与其他粒子碰撞检测（使用KD树加速）
            if self.kd_tree and len(self.particle_positions) > 0:
                pos = particle['position']
                neighbors = self.kd_tree.query_ball_point(
                    [pos.getX(), pos.getY(), pos.getZ()], 
                    0.5  # 碰撞检测半径
                )
                
                # 处理与近邻粒子的碰撞
                for j in neighbors:
                    other_particle = self.fluid_world.fluid_particles[j]
                    if particle != other_particle:
                        distance = (particle['position'] - other_particle['position']).length()
                        if distance < 0.2:  # 碰撞距离
                            # 碰撞响应
                            direction = particle['position'] - other_particle['position']
                            if direction.length() > 0:
                                direction.normalize()
                                particle['velocity'] += direction * 0.5
                                
    def _update_positions(self, dt):
        """
        更新粒子位置
        
        Args:
            dt (float): 时间步长
        """
        for particle in self.fluid_world.fluid_particles:
            # 更新位置
            particle['position'] += particle['velocity'] * dt
            
            # 简单的速度衰减
            particle['velocity'] *= 0.99
            
    def _handle_boundary_conditions(self):
        """
        处理边界条件
        """
        # 确保粒子不会获得过大的速度
        for particle in self.fluid_world.fluid_particles:
            if particle['velocity'].length() > self.speed_limit:
                particle['velocity'].normalize()
                particle['velocity'] *= self.speed_limit