"""
骨骼动画系统插件 - 高级功能扩展
提供动画状态机、IK系统、动画曲线编辑等高级功能
"""

import math
from panda3d.core import *
from direct.actor.Actor import Actor
from direct.interval.IntervalGlobal import *


class AnimationStateMachine:
    """
    动画状态机
    管理复杂的动画状态转换和混合
    """
    
    def __init__(self, actor):
        self.actor = actor
        self.states = {}  # 状态字典 {state_name: state_data}
        self.current_state = None
        self.transitions = {}  # 状态转换 {from_state: {to_state: condition_func}}
        self.state_callbacks = {}  # 状态回调 {state_name: {event: callback_func}}
        
    def add_state(self, state_name, animations, loop=True, play_rate=1.0, blend_time=0.3):
        """
        添加动画状态
        :param state_name: 状态名称
        :param animations: 动画列表或单个动画名称
        :param loop: 是否循环
        :param play_rate: 播放速度
        :param blend_time: 状态切换混合时间
        """
        if isinstance(animations, str):
            animations = [animations]
            
        self.states[state_name] = {
            'animations': animations,
            'loop': loop,
            'play_rate': play_rate,
            'blend_time': blend_time,
            'active_interval': None
        }
        
        print(f"添加动画状态: {state_name}")
        
    def remove_state(self, state_name):
        """
        移除动画状态
        :param state_name: 状态名称
        """
        if state_name in self.states:
            # 如果是当前状态，先停止
            if self.current_state == state_name:
                self.stop_current_state()
            del self.states[state_name]
            print(f"移除动画状态: {state_name}")
            
    def set_transition(self, from_state, to_state, condition_func=None):
        """
        设置状态转换条件
        :param from_state: 起始状态
        :param to_state: 目标状态
        :param condition_func: 转换条件函数，返回True时执行转换
        """
        if from_state not in self.transitions:
            self.transitions[from_state] = {}
        self.transitions[from_state][to_state] = condition_func
        print(f"设置状态转换: {from_state} -> {to_state}")
        
    def set_state_callback(self, state_name, event, callback_func):
        """
        设置状态回调
        :param state_name: 状态名称
        :param event: 事件类型 ('enter', 'exit', 'update')
        :param callback_func: 回调函数
        """
        if state_name not in self.state_callbacks:
            self.state_callbacks[state_name] = {}
        self.state_callbacks[state_name][event] = callback_func
        
    def enter_state(self, state_name):
        """
        进入指定状态
        :param state_name: 状态名称
        """
        if state_name not in self.states:
            print(f"错误: 状态 {state_name} 不存在")
            return False
            
        # 退出当前状态
        if self.current_state:
            self.exit_current_state()
            
        # 进入新状态
        state_data = self.states[state_name]
        
        # 创建动画序列
        intervals = []
        for anim_name in state_data['animations']:
            if anim_name in self.actor.getAnimNames():
                interval = self.actor.actorInterval(
                    anim_name, 
                    loop=state_data['loop'],
                    playRate=state_data['play_rate']
                )
                intervals.append(interval)
                
        if intervals:
            if len(intervals) == 1:
                state_data['active_interval'] = intervals[0]
            else:
                state_data['active_interval'] = Sequence(*intervals)
                
            # 启动动画
            state_data['active_interval'].start()
            
        self.current_state = state_name
        
        # 调用进入回调
        if state_name in self.state_callbacks and 'enter' in self.state_callbacks[state_name]:
            self.state_callbacks[state_name]['enter']()
            
        print(f"进入动画状态: {state_name}")
        return True
        
    def exit_current_state(self):
        """
        退出当前状态
        """
        if not self.current_state:
            return
            
        state_data = self.states[self.current_state]
        
        # 停止当前动画
        if state_data['active_interval']:
            state_data['active_interval'].finish()
            state_data['active_interval'] = None
            
        # 调用退出回调
        if self.current_state in self.state_callbacks and 'exit' in self.state_callbacks[self.current_state]:
            self.state_callbacks[self.current_state]['exit']()
            
        print(f"退出动画状态: {self.current_state}")
        self.current_state = None
        
    def stop_current_state(self):
        """
        停止当前状态
        """
        if not self.current_state:
            return
            
        state_data = self.states[self.current_state]
        
        # 停止当前动画
        if state_data['active_interval']:
            state_data['active_interval'].finish()
            state_data['active_interval'] = None
            
    def update(self, dt):
        """
        更新状态机
        :param dt: 时间增量
        """
        if not self.current_state:
            return
            
        # 调用更新回调
        if self.current_state in self.state_callbacks and 'update' in self.state_callbacks[self.current_state]:
            self.state_callbacks[self.current_state]['update'](dt)
            
        # 检查状态转换
        if self.current_state in self.transitions:
            for to_state, condition_func in self.transitions[self.current_state].items():
                if condition_func is None or condition_func():
                    self.enter_state(to_state)
                    break
                    
    def get_current_state(self):
        """
        获取当前状态
        :return: 当前状态名称
        """
        return self.current_state
        
    def get_available_states(self):
        """
        获取所有可用状态
        :return: 状态名称列表
        """
        return list(self.states.keys())


class InverseKinematicsSystem:
    """
    逆向运动学系统
    实现骨骼的IK控制
    """
    
    def __init__(self, actor):
        self.actor = actor
        self.ik_chains = {}  # IK链 {chain_name: chain_data}
        self.targets = {}  # IK目标 {target_name: target_node}
        
    def add_ik_chain(self, chain_name, end_effector, pole_vector=None, iterations=10):
        """
        添加IK链
        :param chain_name: 链名称
        :param end_effector: 终端效应器（骨骼名称）
        :param pole_vector: 极向量（用于控制弯曲方向）
        :param iterations: 迭代次数
        """
        # 查找骨骼
        effector_joint = self.actor.exposeJoint(None, "modelRoot", end_effector)
        if not effector_joint:
            print(f"错误: 未找到骨骼 {end_effector}")
            return
            
        self.ik_chains[chain_name] = {
            'end_effector': end_effector,
            'effector_joint': effector_joint,
            'pole_vector': pole_vector,
            'iterations': iterations,
            'enabled': True
        }
        
        print(f"添加IK链: {chain_name} -> {end_effector}")
        
    def remove_ik_chain(self, chain_name):
        """
        移除IK链
        :param chain_name: 链名称
        """
        if chain_name in self.ik_chains:
            del self.ik_chains[chain_name]
            print(f"移除IK链: {chain_name}")
            
    def set_ik_target(self, chain_name, target_pos):
        """
        设置IK目标位置
        :param chain_name: 链名称
        :param target_pos: 目标位置 (x, y, z)
        """
        if chain_name not in self.ik_chains:
            print(f"错误: IK链 {chain_name} 不存在")
            return
            
        chain_data = self.ik_chains[chain_name]
        if not chain_data['enabled']:
            return
            
        # 简化的IK实现（CCD算法）
        effector_name = chain_data['end_effector']
        iterations = chain_data['iterations']
        
        # 获取当前终端效应器位置
        effector_joint = chain_data['effector_joint']
        current_pos = effector_joint.getPos(self.actor)
        target_pos = Point3(*target_pos)
        
        # 计算距离
        distance = (target_pos - current_pos).length()
        if distance < 0.01:  # 阈值
            return
            
        # CCD算法迭代
        for _ in range(iterations):
            # 从终端向根部迭代调整骨骼
            # 这里简化处理，实际应该遍历整条骨骼链
            # 为了示例，我们只调整终端效应器朝向
            direction = target_pos - current_pos
            direction.normalize()
            
            # 简化的朝向调整
            effector_joint.lookAt(target_pos)
            
            # 更新位置
            current_pos = effector_joint.getPos(self.actor)
            distance = (target_pos - current_pos).length()
            if distance < 0.01:
                break
                
    def enable_ik_chain(self, chain_name, enabled=True):
        """
        启用/禁用IK链
        :param chain_name: 链名称
        :param enabled: 是否启用
        """
        if chain_name in self.ik_chains:
            self.ik_chains[chain_name]['enabled'] = enabled
            status = "启用" if enabled else "禁用"
            print(f"{status}IK链: {chain_name}")
            
    def batch_set_targets(self, targets_dict):
        """
        批量设置IK目标
        :param targets_dict: {chain_name: target_pos}
        """
        for chain_name, target_pos in targets_dict.items():
            self.set_ik_target(chain_name, target_pos)


class AnimationCurveEditor:
    """
    动画曲线编辑器
    支持对动画属性进行曲线编辑和自定义
    """
    
    def __init__(self, actor):
        self.actor = actor
        self.curves = {}  # 曲线字典 {curve_name: curve_data}
        self.curve_targets = {}  # 曲线目标 {target_name: target_data}
        
    def create_curve(self, curve_name, keyframes, interpolation='linear'):
        """
        创建动画曲线
        :param curve_name: 曲线名称
        :param keyframes: 关键帧列表 [(time, value), ...]
        :param interpolation: 插值方式 ('linear', 'smooth', 'step')
        """
        self.curves[curve_name] = {
            'keyframes': sorted(keyframes, key=lambda x: x[0]),
            'interpolation': interpolation,
            'current_value': 0.0
        }
        
        print(f"创建动画曲线: {curve_name}")
        
    def remove_curve(self, curve_name):
        """
        移除动画曲线
        :param curve_name: 曲线名称
        """
        if curve_name in self.curves:
            del self.curves[curve_name]
            print(f"移除动画曲线: {curve_name}")
            
    def add_curve_target(self, target_name, curve_name, property_path, transform_type='scale'):
        """
        添加曲线目标
        :param target_name: 目标名称
        :param curve_name: 曲线名称
        :param property_path: 属性路径（如骨骼名称）
        :param transform_type: 变换类型 ('scale', 'position', 'rotation')
        """
        if curve_name not in self.curves:
            print(f"错误: 曲线 {curve_name} 不存在")
            return
            
        self.curve_targets[target_name] = {
            'curve_name': curve_name,
            'property_path': property_path,
            'transform_type': transform_type
        }
        
        print(f"添加曲线目标: {target_name} -> {curve_name}")
        
    def evaluate_curve(self, curve_name, time):
        """
        计算曲线在指定时间的值
        :param curve_name: 曲线名称
        :param time: 时间
        :return: 曲线值
        """
        if curve_name not in self.curves:
            return 0.0
            
        curve_data = self.curves[curve_name]
        keyframes = curve_data['keyframes']
        
        if not keyframes:
            return 0.0
            
        # 边界情况
        if time <= keyframes[0][0]:
            return keyframes[0][1]
        if time >= keyframes[-1][0]:
            return keyframes[-1][1]
            
        # 查找相邻关键帧
        for i in range(len(keyframes) - 1):
            if keyframes[i][0] <= time <= keyframes[i + 1][0]:
                t1, v1 = keyframes[i]
                t2, v2 = keyframes[i + 1]
                
                # 根据插值方式计算值
                if curve_data['interpolation'] == 'step':
                    return v1
                elif curve_data['interpolation'] == 'smooth':
                    # 平滑插值（三次 Hermite）
                    t = (time - t1) / (t2 - t1)
                    t2_sq = t * t
                    t3_sq = t2_sq * t
                    h00 = 2 * t3_sq - 3 * t2_sq + 1
                    h10 = t3_sq - 2 * t2_sq + t
                    h01 = -2 * t3_sq + 3 * t2_sq
                    h11 = t3_sq - t2_sq
                    return h00 * v1 + h10 * (t2 - t1) + h01 * v2 + h11 * (t2 - t1)
                else:  # linear
                    t = (time - t1) / (t2 - t1)
                    return v1 + (v2 - v1) * t
                    
        return 0.0
        
    def update_curves(self, time):
        """
        更新所有曲线
        :param time: 当前时间
        """
        for target_name, target_data in self.curve_targets.items():
            curve_name = target_data['curve_name']
            property_path = target_data['property_path']
            transform_type = target_data['transform_type']
            
            # 计算曲线值
            value = self.evaluate_curve(curve_name, time)
            
            # 应用到目标
            self.apply_curve_value(property_path, transform_type, value)
            
            # 更新曲线数据
            if curve_name in self.curves:
                self.curves[curve_name]['current_value'] = value
                
    def apply_curve_value(self, property_path, transform_type, value):
        """
        应用曲线值到目标属性
        :param property_path: 属性路径
        :param transform_type: 变换类型
        :param value: 值
        """
        # 查找目标骨骼
        joint = self.actor.exposeJoint(None, "modelRoot", property_path)
        if not joint:
            return
            
        # 根据变换类型应用值
        if transform_type == 'scale':
            joint.setScale(value, value, value)
        elif transform_type == 'position':
            current_pos = joint.getPos()
            if property_path.endswith('.x'):
                joint.setPos(value, current_pos.y, current_pos.z)
            elif property_path.endswith('.y'):
                joint.setPos(current_pos.x, value, current_pos.z)
            elif property_path.endswith('.z'):
                joint.setPos(current_pos.x, current_pos.y, value)
        elif transform_type == 'rotation':
            # 简化处理，实际应该分解为具体的旋转轴
            joint.setHpr(value, value, value)


class AdvancedAnimationManager:
    """
    高级动画管理器
    集成状态机、IK系统和曲线编辑器
    """
    
    def __init__(self, actor):
        self.actor = actor
        self.state_machine = AnimationStateMachine(actor)
        self.ik_system = InverseKinematicsSystem(actor)
        self.curve_editor = AnimationCurveEditor(actor)
        self.enabled_systems = {
            'state_machine': True,
            'ik_system': True,
            'curve_editor': True
        }
        
    def update(self, dt, time):
        """
        更新所有高级动画系统
        :param dt: 时间增量
        :param time: 当前时间
        """
        if self.enabled_systems['state_machine']:
            self.state_machine.update(dt)
            
        if self.enabled_systems['curve_editor']:
            self.curve_editor.update_curves(time)
            
    def enable_system(self, system_name, enabled=True):
        """
        启用/禁用子系统
        :param system_name: 系统名称 ('state_machine', 'ik_system', 'curve_editor')
        :param enabled: 是否启用
        """
        if system_name in self.enabled_systems:
            self.enabled_systems[system_name] = enabled
            status = "启用" if enabled else "禁用"
            print(f"{status}动画子系统: {system_name}")


# 使用示例和测试代码
def example_state_machine_usage(actor):
    """
    状态机使用示例
    """
    print("=== 动画状态机使用示例 ===")
    
    # 创建状态机
    state_machine = AnimationStateMachine(actor)
    
    # 添加状态
    state_machine.add_state("idle", "idle_anim", loop=True)
    state_machine.add_state("walk", "walk_anim", loop=True)
    state_machine.add_state("run", "run_anim", loop=True)
    state_machine.add_state("jump", ["jump_start", "jump_loop", "jump_end"], loop=False)
    
    # 设置状态转换
    state_machine.set_transition("idle", "walk")
    state_machine.set_transition("walk", "run")
    state_machine.set_transition("run", "idle")
    
    # 进入初始状态
    state_machine.enter_state("idle")
    
    print("状态机示例设置完成")
    return state_machine


def example_ik_system_usage(actor):
    """
    IK系统使用示例
    """
    print("=== 逆向运动学系统使用示例 ===")
    
    # 创建IK系统
    ik_system = InverseKinematicsSystem(actor)
    
    # 添加IK链
    ik_system.add_ik_chain("right_arm", "right_hand")
    ik_system.add_ik_chain("left_arm", "left_hand")
    ik_system.add_ik_chain("right_leg", "right_foot")
    ik_system.add_ik_chain("left_leg", "left_foot")
    
    # 设置IK目标
    ik_system.set_ik_target("right_arm", (1, 0, 1))
    ik_system.set_ik_target("left_arm", (-1, 0, 1))
    
    print("IK系统示例设置完成")
    return ik_system


def example_curve_editor_usage(actor):
    """
    曲线编辑器使用示例
    """
    print("=== 动画曲线编辑器使用示例 ===")
    
    # 创建曲线编辑器
    curve_editor = AnimationCurveEditor(actor)
    
    # 创建曲线
    scale_curve_frames = [(0, 1.0), (1, 1.5), (2, 1.0), (3, 0.8), (4, 1.0)]
    curve_editor.create_curve("bounce_scale", scale_curve_frames, 'smooth')
    
    position_curve_frames = [(0, 0), (1, 0.5), (2, 0), (3, -0.3), (4, 0)]
    curve_editor.create_curve("bounce_position", position_curve_frames, 'smooth')
    
    # 添加曲线目标
    curve_editor.add_curve_target("character_scale", "bounce_scale", "character", "scale")
    curve_editor.add_curve_target("character_height", "bounce_position", "character", "position")
    
    print("曲线编辑器示例设置完成")
    return curve_editor


def run_advanced_examples(actor):
    """
    运行所有高级功能示例
    """
    print("开始运行高级动画功能示例...")
    
    state_machine = example_state_machine_usage(actor)
    ik_system = example_ik_system_usage(actor)
    curve_editor = example_curve_editor_usage(actor)
    
    print("所有高级功能示例运行完成!")
    
    return state_machine, ik_system, curve_editor


if __name__ == "__main__":
    # 这里可以添加更多测试代码
    print("高级动画功能模块加载完成")