from hardware.dobot_nova5 import dobot_nova5

from algorithms.arm_state import ArmState
from algorithms.controller_manager import ControllerManager
from algorithms.admittance_controller import AdmittanceController

from tools.log import CustomLogger
from tools.yaml_operator import read_yaml
from tools.Rate import Rate
import numpy as np
import threading
import time
import os
class MassageRobot:
    def __init__(self,arm_config_path,is_log=False):
        self.logger = CustomLogger(
            log_name="测试日志",
            log_file="logs/MassageRobot_nova5_test.log",
            precise_time=True)
        if is_log:
            self.logger_data = CustomLogger(log_name="运动数据日志",
                                            log_file="logs/MassageRobot_kinematics_data.log",
                                            precise_time=True)
            # 日志线程    
            threading.Thread(target=self.log_thread,daemon=True).start()
        
        self.arm_state = ArmState()
        self.arm_config = read_yaml(arm_config_path)
        self.arm = dobot_nova5(arm_ip=self.arm_config['arm_ip'])
        self.force_sensor = None

        # controller
        self.controller_manager = ControllerManager(self.arm_state)
        self.controller_manager.add_controller(AdmittanceController,'admittance',self.arm_config['controller'][0])
        
        # massage heads
        massage_head_dir = self.arm_config['massage_head_dir']
        all_items = os.listdir(massage_head_dir)
        head_config_files = [f for f in all_items if os.path.isfile(os.path.join(massage_head_dir, f))]
        self.playload_dict = {}
        for file in head_config_files:
            file_address = massage_head_dir + '/' + file
            play_load = read_yaml(file_address)
            self.playload_dict[play_load['name']] = play_load
        self.current_head = 'none'

        # 频率
        self.control_rate = Rate(self.arm_config['control_rate'])
        self.sensor_rate = Rate(self.arm_config['sensor_rate'])
        self.command_rate = Rate(self.arm_config['command_rate'])

        # 低通滤波
        self.cutoff_freq = 80.0

        # flags 
        self.exit_event = threading.Event()
        self.exit_event.set() # 运行 True
        self.interrupt_event = threading.Event()
        self.interrupt_event.clear() # 中断 False
        self.pause_envent = threading.Event()
        self.pause_envent.clear() # 暂停 False
        self.skip_envent = threading.Event()
        self.skip_envent.clear() # 跳过 False

        self.is_waitting = False

        self.last_print_time = 0
        self.last_record_time = 0
        self.last_command_time = 0
        self.move_to_point_count = 0
        self.width_default = 5

        self.x_base = np.zeros(6)
        self.P_base  = np.eye(6)
        # 过程噪声协方差矩阵
        self.Q_base = np.eye(6) * 0.01
        # 测量噪声协方差矩阵
        self.R_base = np.eye(6) * 0.1

        self.x_tcp = np.zeros(6)
        self.P_tcp  = np.eye(6)
        # 过程噪声协方差矩阵
        self.Q_tcp = np.eye(6) * 0.01
        # 测量噪声协方差矩阵
        self.R_tcp = np.eye(6) * 0.1
        # 传感器故障计数器
        self.sensor_fail_count = 0 

        # init
        self.arm.__init__()
    
    # 预测步骤
    def kalman_predict(self,x, P, Q):
        # 预测状态（这里假设状态不变）
        x_predict = x
        # 预测误差协方差
        P_predict = P + Q
        return x_predict, P_predict

    # 更新步骤
    def kalman_update(self,x_predict, P_predict, z, R):
        # 卡尔曼增益
        K = P_predict @ np.linalg.inv(P_predict + R)
        # 更新状态
        x_update = x_predict + K @ (z - x_predict)
        # 更新误差协方差
        P_update = (np.eye(len(K)) - K) @ P_predict
        return x_update, P_update
    
    def arm_measure_loop(self):
        return

    def arm_command_loop(self):
        return

    def start(self):
        if self.exit_event.is_set():
            self.exit_event.clear()
            self.arm_measure_thread = threading.Thread(target=self.arm_measure_loop)
            self.arm_control_thread = threading.Thread(target=self.arm_command_loop)
            self.arm_measure_thread.start()
            position,quat_rot = self.arm.get_arm_position()
            # 初始值
            self.arm_state.desired_position = position
            self.arm_state.arm_position_command = position
            self.arm_state.desired_orientation = quat_rot
            self.arm_state.arm_orientation_command = quat_rot
            time.sleep(0.5)
            
            

    def stop(self):
        if not self.exit_event.is_set():
            self.exit_event.set()
            self.interrupt_event.clear()
            self.arm_control_thread.join()
            self.arm_measure_thread.join()
            
            return

    def init_hardwares(self,ready_pose):
        self.ready_pose = np.array(ready_pose)
        
    def switch_payload(self,name):
        if name in self.playload_dict:
            return
            

    def log_thread(self):
        while True:
            self.logger_data.log_info(f"机械臂位置:{self.arm_state.arm_position},机械臂姿态:{self.arm_state.arm_orientation}",is_print=False)
            self.logger_data.log_info(f"机械臂期望位置:{self.arm_state.desired_position},机械臂真实位置:{self.arm_state.arm_position}",is_print=False)
            self.logger_data.log_info(f"机械臂期望姿态:{self.arm_state.desired_orientation},机械臂真实姿态:{self.arm_state.arm_orientation}",is_print=False)
            self.logger_data.log_info(f"机械臂期望力矩:{self.arm_state.desired_wrench},机械臂真实力矩:{self.arm_state.external_wrench_base}",is_print=False)
            self.logger_data.log_info(f"当前按摩头:{self.current_head}",is_print=False)
            time.sleep(1)

if __name__ == "__main__":
    robot = MassageRobot(arm_config_path="MassageControl/config/robot_config.yaml")