"""
群体管理器
负责创建和管理群体成员
"""

from panda3d.core import Vec3, Point3, LVector3f
import random
import math
import time

class SwarmManager:
    """
    群体管理器
    负责创建和管理群体成员，实现完整的群体智能行为
    """
    
    def __init__(self, world):
        self.world = world
        self.members = []
        self.obstacles = []
        self.predators = []
        self.paused = False
        self.task = None
        self.swarm_type = "鸟类"
        self.target = None
        self.path = []
        self.bounds = {
            'min_x': -50, 'max_x': 50,
            'min_y': -50, 'max_y': 50,
            'min_z': 0, 'max_z': 30
        }
        
        # 多群体支持
        self.swarms = []  # 存储多个群体
        self.swarm_id_counter = 0
        
        # 性能统计
        self.frame_count = 0
        self.last_update_time = time.time()
        self.fps = 0
        
        # 群体参数
        self.cohesion_weight = 1.0
        self.separation_weight = 1.5
        self.alignment_weight = 1.0
        self.max_speed = 5.0
        self.perception_radius = 10.0
        self.member_count = 20
        
        # 初始化高级Boids算法
        from .advanced_boids import AdvancedBoidsAlgorithm
        from .config import Config
        from .interaction_manager import InteractionManager
        from .learning_manager import LearningManager
        from .environment_manager import EnvironmentManager
        self.config = Config()
        self.boids_algorithm = AdvancedBoidsAlgorithm(self.config)
        self.interaction_manager = InteractionManager(self.config)
        self.learning_manager = LearningManager(self.config)
        self.environment_manager = EnvironmentManager(self.config)
        
        # 启动更新任务
        self.task = self.world.taskMgr.add(self._update_task, "swarm_update_task")
        
    def create_example_swarm(self, swarm_type=None, member_count=None, position_offset=Vec3(0, 0, 0)):
        """
        创建示例群体
        根据群体类型创建不同类型的群体成员
        """
        # 如果是创建新的独立群体，不清理现有群体
        if swarm_type is None and member_count is None:
            # 清理现有群体
            self.cleanup()
        
        # 获取群体类型和数量
        swarm_type = swarm_type or self.swarm_type
        count = member_count or self.config.get("member_count", self.member_count)
        
        # 创建指定数量的群体成员
        new_members = []
        for i in range(count):
            # 根据群体类型设置随机位置范围，并应用位置偏移
            if swarm_type == "鸟类":
                pos = Vec3(
                    random.uniform(-20, 20),
                    random.uniform(-20, 20),
                    random.uniform(5, 25)
                ) + position_offset
            elif swarm_type == "鱼类":
                pos = Vec3(
                    random.uniform(-15, 15),
                    random.uniform(-15, 15),
                    random.uniform(2, 15)
                ) + position_offset
            elif swarm_type == "昆虫":
                pos = Vec3(
                    random.uniform(-10, 10),
                    random.uniform(-10, 10),
                    random.uniform(1, 10)
                ) + position_offset
            else:  # 自定义
                pos = Vec3(
                    random.uniform(self.bounds['min_x'], self.bounds['max_x']),
                    random.uniform(self.bounds['min_y'], self.bounds['min_y']),
                    random.uniform(self.bounds['min_z'], self.bounds['max_z'])
                ) + position_offset
            
            # 根据群体类型加载不同的模型
            if swarm_type == "鸟类":
                member_model = "models/misc/sphere"  # 实际应用中应使用鸟类模型
                scale = 0.5
            elif swarm_type == "鱼类":
                member_model = "models/misc/sphere"  # 实际应用中应使用鱼类模型
                scale = 0.3
            elif swarm_type == "昆虫":
                member_model = "models/misc/sphere"  # 实际应用中应使用昆虫模型
                scale = 0.2
            else:
                member_model = "models/misc/sphere"
                scale = 0.5
            
            # 加载模型
            try:
                member_node = self.world.loader.loadModel(member_model)
            except:
                # 如果模型加载失败，使用默认球体
                member_node = self.world.loader.loadModel("models/misc/sphere")
            
            member_node.reparentTo(self.world.render)
            member_node.setPos(pos)
            member_node.setScale(scale)
            
            # 随机初始速度
            velocity = Vec3(
                random.uniform(-1, 1),
                random.uniform(-1, 1),
                random.uniform(-1, 1)
            )
            
            # 添加到群体成员列表
            member = {
                'node': member_node,
                'position': pos,
                'velocity': velocity,
                'index': i,  # 成员索引，用于队形计算
                'type': swarm_type,  # 成员类型
                'swarm_id': self.swarm_id_counter  # 群体ID
            }
            new_members.append(member)
            
            # 为成员添加环境传感器
            if hasattr(self, 'environment_manager'):
                self.environment_manager.add_sensor(id(member))
            
        # 如果是创建新的独立群体
        if swarm_type is not None or member_count is not None:
            # 创建新的群体对象
            swarm = {
                'id': self.swarm_id_counter,
                'type': swarm_type,
                'members': new_members,
                'position_offset': position_offset
            }
            self.swarms.append(swarm)
            self.swarm_id_counter += 1
            
            # 将新成员添加到主成员列表
            self.members.extend(new_members)
            
            print(f"已创建包含{len(new_members)}个成员的{swarm_type}群体 (ID: {swarm['id']})")
            return swarm
        else:
            # 更新主成员列表
            self.members = new_members
            print(f"已创建包含{len(self.members)}个成员的{swarm_type}群体")
            
    def create_multiple_swarms(self, swarm_configs: List[Dict]):
        """
        创建多个群体
        :param swarm_configs: 群体配置列表，每个配置包含type, count, position_offset
        """
        self.cleanup()  # 清理现有群体
        self.swarms.clear()  # 清理群体列表
        self.swarm_id_counter = 0  # 重置ID计数器
        
        for config in swarm_configs:
            swarm_type = config.get('type', '鸟类')
            member_count = config.get('count', 20)
            position_offset = config.get('position_offset', Vec3(0, 0, 0))
            
            self.create_example_swarm(swarm_type, member_count, position_offset)
            
        print(f"已创建{len(self.swarms)}个群体")
        
    def get_swarm_by_id(self, swarm_id: int):
        """
        根据ID获取群体
        """
        for swarm in self.swarms:
            if swarm['id'] == swarm_id:
                return swarm
        return None
        
    def get_swarm_by_type(self, swarm_type: str):
        """
        根据类型获取群体
        """
        return [swarm for swarm in self.swarms if swarm['type'] == swarm_type]
        
    def create_obstacle(self, position, radius):
        """
        创建障碍物
        """
        obstacle = {
            'position': position,
            'radius': radius
        }
        self.obstacles.append(obstacle)
        return obstacle
        
    def create_predator(self, position):
        """
        创建捕食者
        """
        # 创建捕食者模型
        try:
            predator_node = self.world.loader.loadModel("models/misc/sphere")  # 实际应用中应使用捕食者模型
        except:
            predator_node = self.world.loader.loadModel("models/misc/sphere")
            
        predator_node.reparentTo(self.world.render)
        predator_node.setPos(position)
        predator_node.setScale(1.0)
        predator_node.setColor(1, 0, 0, 1)  # 红色表示捕食者
        
        # 创建捕食者对象
        predator = {
            'node': predator_node,
            'position': position,
            'velocity': Vec3(0, 0, 0),
            'target_index': 0  # 目标群体成员索引
        }
        
        self.predators.append(predator)
        return predator
        
    def set_target(self, target):
        """
        设置群体目标位置
        """
        self.target = target
        
    def set_path(self, path):
        """
        设置群体移动路径
        """
        self.path = path
        
    def _update_task(self, task):
        """
        更新任务
        """
        if not self.paused:
            self._update_swarm()
            
            # 更新性能统计
            self.frame_count += 1
            current_time = time.time()
            if current_time - self.last_update_time >= 1.0:
                self.fps = self.frame_count
                self.frame_count = 0
                self.last_update_time = current_time
                
        return task.cont
        
    def _update_swarm(self):
        """
        更新群体行为
        """
        # 更新捕食者行为（如果存在）
        self._update_predators()
        
        # 更新群体间交互
        interaction_forces = {}
        if hasattr(self, 'interaction_manager'):
            # 为每个成员计算交互力
            for member in self.members:
                interaction_forces[id(member)] = {}
                
                # 应用合作力
                if 'cooperative_force' in member:
                    interaction_forces[id(member)]['cooperative_force'] = member['cooperative_force']
                    del member['cooperative_force']
                    
                # 应用竞争力
                if 'competition_force' in member:
                    interaction_forces[id(member)]['competition_force'] = member['competition_force']
                    del member['competition_force']
                    
                # 应用捕食力
                if 'predation_force' in member:
                    interaction_forces[id(member)]['predation_force'] = member['predation_force']
                    del member['predation_force']
                    
                # 应用逃避力
                if 'escape_force' in member:
                    interaction_forces[id(member)]['escape_force'] = member['escape_force']
                    del member['escape_force']
        
        # 更新环境状态
        if hasattr(self, 'environment_manager'):
            self.environment_manager.update_environment_state(
                self.obstacles, self.swarms, self.predators, 
                getattr(self, 'resources', []), getattr(self, 'hazards', [])
            )
            
        # 更新学习和适应
        if hasattr(self, 'learning_manager') and self.config.get("learning_enabled", False):
            # 计算环境状态
            environment_state = self._calculate_environment_state()
            
            # 更新环境模型
            self.learning_manager.update_environment_model(environment_state)
            
            # 适应环境
            self.learning_manager.adapt_to_environment(environment_state)
            
            # 获取适应后的参数
            adapted_params = self.learning_manager.get_adapted_parameters()
            
            # 更新Boids算法参数
            self._update_boids_parameters(adapted_params)
        
        # 检查是否启用并行计算
        if self.config.get("parallel_computation", False) and len(self.members) > 50:
            # 使用并行计算更新群体成员
            self._update_swarm_parallel()
        else:
            # 使用串行计算更新群体成员
            for i, member in enumerate(self.members):
                # 计算邻居
                neighbors = self._find_neighbors(member, i)
                
                # 获取扩展参数
                obstacles = self.obstacles if self.config.get("obstacle_avoidance_enabled", True) else []
                bounds = self.bounds if self.config.get("boundary_enabled", True) else None
                predators = self.predators if self.config.get("predator_avoid_enabled", False) else []
                target = self.target if self.config.get("seek_enabled", False) else None
                path = self.path if self.config.get("path_following_enabled", False) else None
                interactions = interaction_forces.get(id(member), {}) if interaction_forces else {}
                
                # 计算奖励（如果启用学习）
                if hasattr(self, 'learning_manager') and self.config.get("learning_enabled", False):
                    reward = self.learning_manager.calculate_reward(member, neighbors, obstacles, environment_state)
                    
                    # 创建状态表示
                    state = {
                        'neighbor_count': len(neighbors),
                        'obstacle_distance': self._calculate_min_obstacle_distance(member),
                        'boundary_distance': self._calculate_boundary_distance(member),
                        'velocity': member['velocity'].length(),
                        'position_x': member['position'].x,
                        'position_y': member['position'].y,
                        'position_z': member['position'].z
                    }
                    
                    # 添加经验（简化）
                    # 在实际应用中，这里需要更复杂的动作表示
                    action = "maintain"  # 简化动作
                    next_state = state.copy()  # 简化下一状态
                    
                    self.learning_manager.add_experience(state, action, reward, next_state)
                
                # 获取环境响应力
                environment_force = Vec3(0, 0, 0)
                if (hasattr(self, 'environment_manager') and 
                    self.config.get("environment_awareness_enabled", True)):
                    environment_force = self.environment_manager.get_environmental_response(
                        member, member['position'], member['velocity']
                    )
                
                # 使用Boids算法更新成员状态，添加环境响应力
                self.boids_algorithm.update_member(
                    member, neighbors, obstacles, target, bounds, path, predators, interactions
                )
                
                # 应用环境响应力
                if self.config.get("environment_awareness_enabled", True):
                    member['velocity'] += environment_force * 0.1  # 环境力权重
                
                # 更新成员节点位置
                member['node'].setPos(member['position'])
                
                # 更新成员朝向（面向移动方向）
                if member['velocity'].length() > 0:
                    # 创建一个向前的向量
                    forward = member['velocity'].normalized()
                    # 设置节点朝向
                    member['node'].lookAt(member['position'] + forward)
                    
            # 更新行为策略（定期更新）
            if hasattr(self, 'learning_manager') and self.config.get("learning_enabled", False):
                if self.frame_count % 60 == 0:  # 每秒更新一次（假设60FPS）
                    self.learning_manager.update_behavior_policies()
                    
    def _update_swarm_parallel(self):
        """
        使用并行计算更新群体行为
        """
        # 获取扩展参数
        obstacles = self.obstacles if self.config.get("obstacle_avoidance_enabled", True) else []
        bounds = self.bounds if self.config.get("boundary_enabled", True) else None
        predators = self.predators if self.config.get("predator_avoid_enabled", False) else []
        target = self.target if self.config.get("seek_enabled", False) else None
        path = self.path if self.config.get("path_following_enabled", False) else None
        
        # 准备交互力
        interaction_forces = {}
        if hasattr(self, 'interaction_manager'):
            # 为每个成员计算交互力
            for member in self.members:
                interaction_forces[id(member)] = {}
                
                # 应用合作力
                if 'cooperative_force' in member:
                    interaction_forces[id(member)]['cooperative_force'] = member['cooperative_force']
                    del member['cooperative_force']
                    
                # 应用竞争力
                if 'competition_force' in member:
                    interaction_forces[id(member)]['competition_force'] = member['competition_force']
                    del member['competition_force']
                    
                # 应用捕食力
                if 'predation_force' in member:
                    interaction_forces[id(member)]['predation_force'] = member['predation_force']
                    del member['predation_force']
                    
                # 应用逃避力
                if 'escape_force' in member:
                    interaction_forces[id(member)]['escape_force'] = member['escape_force']
                    del member['escape_force']
        
        # 准备环境力参数
        environment_forces = {}
        if (hasattr(self, 'environment_manager') and 
            self.config.get("environment_awareness_enabled", True)):
            for member in self.members:
                environment_forces[id(member)] = self.environment_manager.get_environmental_response(
                    member, member['position'], member['velocity']
                )
        else:
            # 如果环境感知被禁用，设置为零力
            for member in self.members:
                environment_forces[id(member)] = Vec3(0, 0, 0)
        
        # 使用并行算法更新群体成员
        self.boids_algorithm.update_swarm_parallel(
            self.members, self.members, obstacles, target, bounds, path, predators, interaction_forces, environment_forces
        )
        
        # 更新所有成员的节点位置和朝向
        for member in self.members:
            member['node'].setPos(member['position'])
            
            # 更新成员朝向（面向移动方向）
            if member['velocity'].length() > 0:
                # 创建一个向前的向量
                forward = member['velocity'].normalized()
                # 设置节点朝向
                member['node'].lookAt(member['position'] + forward)
                
    def _update_predators(self):
        """
        更新捕食者行为
        """
        for predator in self.predators:
            if self.members:
                # 简单的捕食者行为：寻找最近的群体成员
                nearest_member = None
                nearest_distance = float('inf')
                
                for member in self.members:
                    distance = (predator['position'] - member['position']).length()
                    if distance < nearest_distance:
                        nearest_distance = distance
                        nearest_member = member
                        
                if nearest_member:
                    # 向最近的成员移动
                    direction = (nearest_member['position'] - predator['position']).normalized()
                    predator['velocity'] = direction * 3.0  # 捕食者速度
                    predator['position'] += predator['velocity']
                    predator['node'].setPos(predator['position'])
                    
    def _calculate_environment_state(self):
        """
        计算环境状态
        
        :return: 环境状态字典
        """
        if not self.members:
            return {}
            
        # 计算群体密度
        member_count = len(self.members)
        total_area = (self.bounds['max_x'] - self.bounds['min_x']) * (self.bounds['max_y'] - self.bounds['min_y'])
        density = member_count / total_area if total_area > 0 else 0
        
        # 计算平均速度
        total_speed = sum(member['velocity'].length() for member in self.members)
        avg_speed = total_speed / member_count if member_count > 0 else 0
        
        # 计算与障碍物的平均距离
        total_obstacle_distance = 0
        obstacle_count = 0
        for member in self.members:
            min_distance = float('inf')
            for obstacle in self.obstacles:
                distance = (member['position'] - obstacle['position']).length() - obstacle.get('radius', 1.0)
                min_distance = min(min_distance, distance)
            if min_distance != float('inf'):
                total_obstacle_distance += min_distance
                obstacle_count += 1
                
        avg_obstacle_distance = total_obstacle_distance / obstacle_count if obstacle_count > 0 else 10.0
        
        # 威胁水平（基于与障碍物的距离）
        threat_level = max(0.0, min(1.0, 1.0 - avg_obstacle_distance / 20.0))
        
        # 资源丰富度（简化为群体密度的反比）
        resource_abundance = max(0.0, min(1.0, 1.0 - density * 10.0))
        
        return {
            'density': density,
            'avg_speed': avg_speed,
            'avg_obstacle_distance': avg_obstacle_distance,
            'threat_level': threat_level,
            'resource_abundance': resource_abundance,
            'member_count': member_count
        }
        
    def _calculate_min_obstacle_distance(self, member):
        """
        计算成员与最近障碍物的距离
        
        :param member: 群体成员
        :return: 最小距离
        """
        if not self.obstacles:
            return float('inf')
            
        min_distance = float('inf')
        for obstacle in self.obstacles:
            distance = (member['position'] - obstacle['position']).length() - obstacle.get('radius', 1.0)
            min_distance = min(min_distance, distance)
            
        return min_distance
        
    def _calculate_boundary_distance(self, member):
        """
        计算成员与边界的距离
        
        :param member: 群体成员
        :return: 最小边界距离
        """
        min_distance = float('inf')
        
        # 计算到各边界的距离
        distances = [
            member['position'].x - self.bounds['min_x'],
            self.bounds['max_x'] - member['position'].x,
            member['position'].y - self.bounds['min_y'],
            self.bounds['max_y'] - member['position'].y,
            member['position'].z - self.bounds['min_z'],
            self.bounds['max_z'] - member['position'].z
        ]
        
        return min(distances) if distances else float('inf')
        
    def _update_boids_parameters(self, adapted_params):
        """
        更新Boids算法参数
        
        :param adapted_params: 适应后的参数
        """
        # 更新速度
        current_max_speed = self.config.get("max_speed", 5.0)
        self.config.set("max_speed", current_max_speed * adapted_params['speed'])
        
        # 更新行为权重
        current_cohesion = self.config.get("cohesion_weight", 1.0)
        self.config.set("cohesion_weight", current_cohesion * adapted_params['cohesion'])
        
        current_separation = self.config.get("separation_weight", 1.5)
        self.config.set("separation_weight", current_separation * adapted_params['separation'])
        
        current_alignment = self.config.get("alignment_weight", 1.0)
        self.config.set("alignment_weight", current_alignment * adapted_params['alignment'])
                    
    def _find_neighbors(self, member, index):
        """
        查找邻居
        使用空间分区优化邻居查找（简化实现）
        """
        neighbors = []
        perception_radius = self.config.get("perception_radius", self.perception_radius)
        
        for i, other in enumerate(self.members):
            if i != index:  # 排除自己
                distance = (member['position'] - other['position']).length()
                if distance < perception_radius:
                    neighbors.append(other)
        return neighbors
        
    def add_member(self):
        """
        添加一个新的群体成员
        """
        # 随机位置（靠近现有群体）
        if self.members:
            # 在现有群体附近创建新成员
            existing_member = random.choice(self.members)
            pos = Vec3(
                existing_member['position'].x + random.uniform(-5, 5),
                existing_member['position'].y + random.uniform(-5, 5),
                existing_member['position'].z + random.uniform(-2, 2)
            )
        else:
            # 如果没有现有成员，在随机位置创建
            pos = Vec3(
                random.uniform(-20, 20),
                random.uniform(-20, 20),
                random.uniform(5, 15)
            )
            
        # 加载模型
        try:
            member_node = self.world.loader.loadModel("models/misc/sphere")
        except:
            member_node = self.world.loader.loadModel("models/misc/sphere")
            
        member_node.reparentTo(self.world.render)
        member_node.setPos(pos)
        member_node.setScale(0.5)
        
        # 随机初始速度
        velocity = Vec3(
            random.uniform(-1, 1),
            random.uniform(-1, 1),
            random.uniform(-1, 1)
        )
        
        # 添加到群体成员列表
        member = {
            'node': member_node,
            'position': pos,
            'velocity': velocity,
            'index': len(self.members),  # 成员索引
            'type': self.swarm_type
        }
        
        self.members.append(member)
        self.config.set("member_count", len(self.members))
        return member
        
    def remove_member(self):
        """
        移除一个群体成员
        """
        if self.members:
            # 移除随机成员
            member = self.members.pop()
            member['node'].removeNode()
            self.config.set("member_count", len(self.members))
            return member
        return None
        
    def update_swarm_type(self, swarm_type):
        """
        更新群体类型
        """
        self.swarm_type = swarm_type
        # 重新创建群体以应用新的类型
        self.cleanup()
        self.create_example_swarm()
        
    def get_stats(self):
        """
        获取群体统计信息
        """
        return {
            'member_count': len(self.members),
            'paused': self.paused,
            'fps': self.fps,
            'swarm_type': self.swarm_type
        }
        
    def toggle_pause(self):
        """
        切换暂停状态
        """
        self.paused = not self.paused
        print(f"群体模拟{'已暂停' if self.paused else '已恢复'}")
        
    def cleanup(self):
        """
        清理资源
        """
        # 移除任务
        if self.task:
            self.world.taskMgr.remove(self.task)
            
        # 移除所有群体成员
        for member in self.members:
            member['node'].removeNode()
            
        # 移除所有捕食者
        for predator in self.predators:
            predator['node'].removeNode()
            
        # 清空列表
        self.members.clear()
        self.obstacles.clear()
        self.predators.clear()
        
    def reset(self):
        """
        重置群体
        """
        self.cleanup()
        self.create_example_swarm()