#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
自定义步态测试脚本
测试俯身步态、爬坡步态、石板路步态等特殊步态
"""

import time
import sys
import os
import toml
import copy
import math

# 添加父目录到路径，以便能够导入utils
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
# 添加task_3和task_4目录
sys.path.append(os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), 'task_3'))
sys.path.append(os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), 'task_4'))

from utils.log_helper import LogHelper, get_logger, section, info, debug, warning, error, success, timing
from base_move.go_straight import go_straight
from file_send_lcmt import file_send_lcmt

# 创建日志记录器
logger = get_logger("自定义步态测试")

# 机器人命令模板
robot_cmd = {
    'mode':0, 'gait_id':0, 'contact':0, 'life_count':0,
    'vel_des':[0.0, 0.0, 0.0],
    'rpy_des':[0.0, 0.0, 0.0],
    'pos_des':[0.0, 0.0, 0.0],
    'acc_des':[0.0, 0.0, 0.0, 0.0, 0.0, 0.0],
    'ctrl_point':[0.0, 0.0, 0.0],
    'foot_pose':[0.0, 0.0, 0.0, 0.0, 0.0, 0.0],
    'step_height':[0.0, 0.0],
    'value':0,  'duration':0
}

def test_stoop_gait(ctrl, msg, observe=True):
    """测试俯身步态"""
    section('测试：俯身步态', "步态测试")
    info('开始测试俯身步态...', "测试")
    
    try:
        # 加载俯身步态配置
        gait_config_path = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), 
                                       'task_4', 'Gait_Params_stoop.toml')
        gait_def_path = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), 
                                    'task_4', 'Gait_Def_stoop.toml')
        
        if not os.path.exists(gait_config_path):
            error(f"俯身步态配置文件不存在: {gait_config_path}", "错误")
            return
        
        steps = toml.load(gait_config_path)
        full_steps = {'step':[robot_cmd]}
        
        k = 0
        for i in steps['step']:
            cmd = copy.deepcopy(robot_cmd)
            cmd['duration'] = i['duration']
            if i['type'] == 'usergait':                
                cmd['mode'] = 11  # LOCOMOTION
                cmd['gait_id'] = 110  # USERGAIT
                cmd['vel_des'] = i['body_vel_des']
                cmd['rpy_des'] = i['body_pos_des'][0:3]
                cmd['pos_des'] = i['body_pos_des'][3:6]
                cmd['foot_pose'][0:2] = i['landing_pos_des'][0:2]
                cmd['foot_pose'][2:4] = i['landing_pos_des'][3:5]
                cmd['foot_pose'][4:6] = i['landing_pos_des'][6:8]
                cmd['ctrl_point'][0:2] = i['landing_pos_des'][9:11]
                cmd['step_height'][0] = math.ceil(i['step_height'][0] * 1e3) + math.ceil(i['step_height'][1] * 1e3) * 1e3
                cmd['step_height'][1] = math.ceil(i['step_height'][2] * 1e3) + math.ceil(i['step_height'][3] * 1e3) * 1e3
                cmd['acc_des'] = i['weight']
                cmd['value'] = i['use_mpc_traj']
                cmd['contact'] = math.floor(i['landing_gain'] * 1e1)
                cmd['ctrl_point'][2] = i['mu']
            if k == 0:
                full_steps['step'] = [cmd]
            else:
                full_steps['step'].append(cmd)
            k = k + 1
        
        # 保存完整步态文件
        full_gait_path = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), 
                                     'task_4', 'Gait_Params_stoop_full.toml')
        with open(full_gait_path, 'w') as f:
            f.write("# Gait Params\n")
            f.writelines(toml.dumps(full_steps))
        
        # 发送步态文件
        usergait_msg = file_send_lcmt()
        
        # 发送步态定义
        with open(gait_def_path, 'r') as f:
            usergait_msg.data = f.read()
        ctrl.lc_s.publish("user_gait_file", usergait_msg.encode())
        time.sleep(0.5)
        
        # 发送步态参数
        with open(full_gait_path, 'r') as f:
            usergait_msg.data = f.read()
        ctrl.lc_s.publish("user_gait_file", usergait_msg.encode())
        time.sleep(0.1)
        
        # 记录初始位置
        initial_pos = ctrl.odo_msg.xyz
        info(f"初始位置: x={initial_pos[0]:.3f}, y={initial_pos[1]:.3f}", "位置")
        
        # 执行俯身步态
        msg.mode = 62
        msg.value = 0
        msg.contact = 15
        msg.gait_id = 110
        msg.duration = 1000
        msg.life_count += 1
        
        info("开始执行俯身步态，持续5秒", "执行")
        start_time = time.time()
        
        while time.time() - start_time < 5.0:  # 执行5秒
            ctrl.Send_cmd(msg)
            time.sleep(0.1)
        
        # 记录最终位置
        final_pos = ctrl.odo_msg.xyz
        distance_moved = ((final_pos[0] - initial_pos[0])**2 + (final_pos[1] - initial_pos[1])**2)**0.5
        info(f"最终位置: x={final_pos[0]:.3f}, y={final_pos[1]:.3f}", "位置")
        info(f"移动距离: {distance_moved:.3f}米", "距离")
        
        success("俯身步态测试完成", "成功")
        
    except Exception as e:
        error(f"俯身步态测试失败: {str(e)}", "失败")
    finally:
        # 恢复正常模式
        msg.mode = 7  # PureDamper
        msg.gait_id = 0
        msg.duration = 0
        msg.life_count += 1
        ctrl.Send_cmd(msg)

def test_climbing_gait(ctrl, msg, observe=True):
    """测试爬坡步态"""
    section('测试：爬坡步态', "步态测试")
    info('开始测试爬坡步态...', "测试")
    
    try:
        # 加载爬坡步态配置
        gait_config_path = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), 
                                       'task_3', 'Gait_Params_up.toml')
        gait_def_path = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), 
                                    'task_3', 'Gait_Def_up.toml')
        
        if not os.path.exists(gait_config_path):
            error(f"爬坡步态配置文件不存在: {gait_config_path}", "错误")
            return
        
        steps = toml.load(gait_config_path)
        full_steps = {'step':[robot_cmd]}
        
        k = 0
        for i in steps['step']:
            cmd = copy.deepcopy(robot_cmd)
            cmd['duration'] = i['duration']
            if i['type'] == 'usergait':                
                cmd['mode'] = 11  # LOCOMOTION
                cmd['gait_id'] = 110  # USERGAIT
                cmd['vel_des'] = i['body_vel_des']
                cmd['rpy_des'] = i['body_pos_des'][0:3]
                cmd['pos_des'] = i['body_pos_des'][3:6]
                cmd['foot_pose'][0:2] = i['landing_pos_des'][0:2]
                cmd['foot_pose'][2:4] = i['landing_pos_des'][3:5]
                cmd['foot_pose'][4:6] = i['landing_pos_des'][6:8]
                cmd['ctrl_point'][0:2] = i['landing_pos_des'][9:11]
                cmd['step_height'][0] = math.ceil(i['step_height'][0] * 1e3) + math.ceil(i['step_height'][1] * 1e3) * 1e3
                cmd['step_height'][1] = math.ceil(i['step_height'][2] * 1e3) + math.ceil(i['step_height'][3] * 1e3) * 1e3
                cmd['acc_des'] = i['weight']
                cmd['value'] = i['use_mpc_traj']
                cmd['contact'] = math.floor(i['landing_gain'] * 1e1)
                cmd['ctrl_point'][2] = i['mu']
            if k == 0:
                full_steps['step'] = [cmd]
            else:
                full_steps['step'].append(cmd)
            k = k + 1
        
        # 保存完整步态文件
        full_gait_path = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), 
                                     'task_3', 'Gait_Params_up_full.toml')
        with open(full_gait_path, 'w') as f:
            f.write("# Gait Params\n")
            f.writelines(toml.dumps(full_steps))
        
        # 发送步态文件
        usergait_msg = file_send_lcmt()
        
        # 发送步态定义
        with open(gait_def_path, 'r') as f:
            usergait_msg.data = f.read()
        ctrl.lc_s.publish("user_gait_file", usergait_msg.encode())
        time.sleep(0.5)
        
        # 发送步态参数
        with open(full_gait_path, 'r') as f:
            usergait_msg.data = f.read()
        ctrl.lc_s.publish("user_gait_file", usergait_msg.encode())
        time.sleep(0.1)
        
        # 记录初始位置
        initial_pos = ctrl.odo_msg.xyz
        info(f"初始位置: x={initial_pos[0]:.3f}, y={initial_pos[1]:.3f}", "位置")
        
        # 执行爬坡步态
        msg.mode = 62
        msg.value = 0
        msg.contact = 15
        msg.gait_id = 110
        msg.duration = 1000
        msg.life_count += 1
        
        info("开始执行爬坡步态，持续3秒", "执行")
        start_time = time.time()
        
        while time.time() - start_time < 3.0:  # 执行3秒
            ctrl.Send_cmd(msg)
            # 监测高度变化
            current_height = ctrl.odo_msg.xyz[2]
            if time.time() - start_time > 0.5:  # 0.5秒后开始监测
                height_change = current_height - initial_pos[2]
                if observe and int((time.time() - start_time) * 10) % 10 == 0:
                    info(f"高度变化: {height_change:.3f}米", "监测")
            time.sleep(0.1)
        
        # 记录最终位置
        final_pos = ctrl.odo_msg.xyz
        distance_moved = ((final_pos[0] - initial_pos[0])**2 + (final_pos[1] - initial_pos[1])**2)**0.5
        height_change = final_pos[2] - initial_pos[2]
        info(f"最终位置: x={final_pos[0]:.3f}, y={final_pos[1]:.3f}, z={final_pos[2]:.3f}", "位置")
        info(f"水平移动距离: {distance_moved:.3f}米", "距离")
        info(f"高度变化: {height_change:.3f}米", "高度")
        
        success("爬坡步态测试完成", "成功")
        
    except Exception as e:
        error(f"爬坡步态测试失败: {str(e)}", "失败")
    finally:
        # 恢复正常模式
        msg.mode = 7  # PureDamper
        msg.gait_id = 0
        msg.duration = 0
        msg.life_count += 1
        ctrl.Send_cmd(msg)

def test_stone_road_gait(ctrl, msg, observe=True):
    """测试石板路步态（快步跑）"""
    section('测试：石板路步态（快步跑）', "步态测试")
    info('开始测试石板路步态...', "测试")
    
    try:
        # 记录初始位置
        initial_pos = ctrl.odo_msg.xyz
        info(f"初始位置: x={initial_pos[0]:.3f}, y={initial_pos[1]:.3f}", "位置")
        
        # 设置石板路步态参数
        msg.mode = 11  # 运动模式
        msg.gait_id = 3  # 步态ID（快步跑）
        msg.vel_des = [0.35, 0, 0]  # 期望速度
        msg.pos_des = [0, 0, 0]
        msg.duration = 0  # 零时长表示持续运动
        msg.step_height = [0.21, 0.21]  # 持续运动时摆动腿的离地高度
        msg.life_count += 1
        
        info("开始执行石板路步态，持续3秒", "执行")
        start_time = time.time()
        
        while time.time() - start_time < 3.0:  # 执行3秒
            ctrl.Send_cmd(msg)
            time.sleep(0.1)
        
        # 停止运动
        msg.vel_des = [0, 0, 0]
        msg.life_count += 1
        ctrl.Send_cmd(msg)
        
        # 记录最终位置
        final_pos = ctrl.odo_msg.xyz
        distance_moved = ((final_pos[0] - initial_pos[0])**2 + (final_pos[1] - initial_pos[1])**2)**0.5
        info(f"最终位置: x={final_pos[0]:.3f}, y={final_pos[1]:.3f}", "位置")
        info(f"移动距离: {distance_moved:.3f}米", "距离")
        info(f"平均速度: {distance_moved/3.0:.3f}米/秒", "速度")
        
        success("石板路步态测试完成", "成功")
        
    except Exception as e:
        error(f"石板路步态测试失败: {str(e)}", "失败")
    finally:
        # 恢复正常模式
        msg.mode = 7  # PureDamper
        msg.gait_id = 0
        msg.duration = 0
        msg.life_count += 1
        ctrl.Send_cmd(msg)

def test_normal_walking_gait(ctrl, msg, observe=True):
    """测试正常行走步态"""
    section('测试：正常行走步态', "步态测试")
    info('开始测试正常行走步态...', "测试")
    
    try:
        # 记录初始位置
        initial_pos = ctrl.odo_msg.xyz
        info(f"初始位置: x={initial_pos[0]:.3f}, y={initial_pos[1]:.3f}", "位置")
        
        # 使用基础直线移动函数测试正常步态
        go_straight(ctrl, msg, distance=1.0, speed=0.5, observe=observe)
        
        # 记录最终位置
        final_pos = ctrl.odo_msg.xyz
        distance_moved = ((final_pos[0] - initial_pos[0])**2 + (final_pos[1] - initial_pos[1])**2)**0.5
        info(f"最终位置: x={final_pos[0]:.3f}, y={final_pos[1]:.3f}", "位置")
        info(f"移动距离: {distance_moved:.3f}米", "距离")
        
        success("正常行走步态测试完成", "成功")
        
    except Exception as e:
        error(f"正常行走步态测试失败: {str(e)}", "失败")

def test_lie_down_stand_up(ctrl, msg, observe=True):
    """测试躺下和站起动作"""
    section('测试：躺下和站起动作', "姿态测试")
    info('开始测试躺下和站起动作...', "测试")
    
    try:
        # 记录初始姿态
        initial_pos = ctrl.odo_msg.xyz
        info(f"初始位置: x={initial_pos[0]:.3f}, y={initial_pos[1]:.3f}, z={initial_pos[2]:.3f}", "位置")
        
        # 躺下
        info("机器人躺下", "动作")
        ctrl.base_msg.lie_down(wait_time=3000)  # 躺下3秒
        
        # 记录躺下后的位置
        lying_pos = ctrl.odo_msg.xyz
        info(f"躺下后位置: x={lying_pos[0]:.3f}, y={lying_pos[1]:.3f}, z={lying_pos[2]:.3f}", "位置")
        height_change = lying_pos[2] - initial_pos[2]
        info(f"高度变化: {height_change:.3f}米", "高度")
        
        # 站起来
        info("机器人站起来", "动作")
        ctrl.base_msg.stand_up()
        
        # 等待站起完成
        time.sleep(2)
        
        # 记录站起后的位置
        standing_pos = ctrl.odo_msg.xyz
        info(f"站起后位置: x={standing_pos[0]:.3f}, y={standing_pos[1]:.3f}, z={standing_pos[2]:.3f}", "位置")
        height_recovery = standing_pos[2] - lying_pos[2]
        info(f"高度恢复: {height_recovery:.3f}米", "高度")
        
        success("躺下和站起动作测试完成", "成功")
        
    except Exception as e:
        error(f"躺下和站起动作测试失败: {str(e)}", "失败")

def run_custom_gait_tests(ctrl, msg):
    """运行所有自定义步态测试"""
    section('自定义步态测试套件', "开始")
    
    tests = [
        ("正常行走步态", test_normal_walking_gait),
        ("石板路步态（快步跑）", test_stone_road_gait),
        ("躺下和站起动作", test_lie_down_stand_up),
        ("俯身步态", test_stoop_gait),
        ("爬坡步态", test_climbing_gait),
    ]
    
    for test_name, test_func in tests:
        try:
            info(f"开始执行测试: {test_name}", "测试")
            test_func(ctrl, msg)
            success(f"测试 {test_name} 成功完成", "成功")
        except Exception as e:
            error(f"测试 {test_name} 失败: {str(e)}", "失败")
        
        # 每个测试之间暂停
        time.sleep(3)
    
    success("所有自定义步态测试完成", "完成")

if __name__ == "__main__":
    # 这里可以添加独立运行的代码
    print("自定义步态测试脚本")
    print("使用方法：")
    print("from single_test.test_custom_gait import run_custom_gait_tests")
    print("run_custom_gait_tests(ctrl, msg)")