import numpy as np
import matplotlib.pyplot as plt
import copy

class ForcePlanner:
    def __init__(self):
        self.force_vel_max = np.float64(20) # [0,0, 10, 0, 0, 0] # 10 N/s
        self.force_acc_max = np.float64(10) #[0,0, 20, 0, 0, 0] # 20 N/s²
        self.force_jerk_max = 90 #[0,0, 30, 0, 0, 0] # 30 N/s³
        self.Ta = self.force_vel_max / self.force_acc_max

    def linear_wrench_interpolate(self,wrench=None, time_points=None, time_step=0.1):
        if wrench is None:
            raise ValueError("至少需要输入力矩的序列")
        if time_points is None or len(time_points) < 2:
            raise ValueError("需要提供至少两个时间点")
        
        # 由于浮点数精度问题，time_step要乘0.9
        times = np.arange(time_points[0], time_points[-1] + time_step * 0.9, time_step)

        if wrench is not None:
            all_wrenchs = np.zeros((len(times), wrench.shape[1]))
            current_idx = 0
            start_wrench =  np.array([0, 0, 0, 0, 0, 0])
            end_wrench =  wrench

            for i in range(len(times)):
                # 计算 phase
                phase = i / len(times)  # 使用 i 计算

                # 计算 segment_wrenchs
                segment_wrenchs = start_wrench + phase * (end_wrench - start_wrench)

                # 将 segment_wrenchs 填充到 all_wrenchs
                all_wrenchs[current_idx:current_idx + len(segment_wrenchs)] = segment_wrenchs
                current_idx += len(segment_wrenchs)

            # 确保最后一个位置被处理
            if current_idx < len(times):
                all_wrenchs[current_idx:] = wrench

        all_wrenchs = np.clip(all_wrenchs, np.array([0, 0, -50, 0, 0, 0]), np.array([0, 0, 0, 0, 0, 0]))

        if wrench is not None:
            return np.array(all_wrenchs)

    def oscillation_wrench_interpolate(self,wrench=None, time_points=None, time_step=0.1):
        if wrench is None:
            raise ValueError("至少需要输入力矩的序列")
        if time_points is None or len(time_points) < 2:
            raise ValueError("需要提供至少两个时间点")
        
        # 由于浮点数精度问题，time_step要乘0.9
        times = np.arange(time_points[0], time_points[-1] + time_step * 0.9, time_step)

        if wrench is not None:
            all_wrenchs = np.zeros((len(times), wrench.shape[1]))
            current_idx = 0


            up_wrench = np.array([0, 0, 0, 0, 0, 0])
            down_wrench = np.array([0, 0, 0, 0, 0, 0])
            up_wrench[2] = wrench[0][2] + 15 if wrench[0][2] + 15 < 0 else 0
            down_wrench[2] = wrench[0][2] - 15 if wrench[0][2] - 15 > -50 else -50


            for i in range(len(times)):
                # 计算 phase
                phase = np.sin(20 * np.pi * i / len(times))  # 使用 i 计算

                phase = 1 if phase >= 0 else -1 

                # 计算 amplitude
                # amplitude = np.array([0, 0, -50, 0, 0, 0]) if phase < 0 else np.array([0, 0, 0, 0, 0, 0])
                amplitude = down_wrench if phase < 0 else up_wrench
                # 计算 segment_wrenchs
                segment_wrenchs = wrench + phase * (amplitude - wrench)

                # 将 segment_wrenchs 填充到 all_wrenchs
                all_wrenchs[current_idx:current_idx + len(segment_wrenchs)] = segment_wrenchs
                current_idx += len(segment_wrenchs)

            # 确保最后一个位置被处理
            if current_idx < len(times):
                all_wrenchs[current_idx:] = wrench

        all_wrenchs = np.clip(all_wrenchs, np.array([0, 0, -70, 0, 0, 0]), np.array([0, 0, 0, 0, 0, 0]))

        if wrench is not None:
            return np.array(all_wrenchs)


    def S_shaped_wrench_interpolate(self,start=None, wrench=None, time_points=None, time_step=0.1):
        if wrench is None:
            raise ValueError("At least a wrench sequence must be provided.")
        if time_points is None or len(time_points) < 2:
            raise ValueError("At least two time points are required.")
        if start is None:
            raise ValueError("At least a start wrench must be provided.")
        Tave = time_points[-1] / len(wrench)
        print("Tave:",Tave)

        start_wrench = copy.deepcopy(np.array(start))
        print("start_wrench:",start_wrench)
        
        # Time points array creation
        times = np.arange(time_points[0], time_points[-1] + time_step * 0.9, time_step)
        # wrench = [wrench]
        wrench = np.array(wrench)

        force_vel = 0
        num = 0

        all_wrenchs = np.zeros((len(times), 6))  # Use wrench.shape[0] to get the number of components
        # 时间点的插值计算
        for count in range(len(times)):
            # 计算当前阶段的进度
            phase = count / len(times)
            plant = phase * times[-1]
            # print("plant:",plant)

            if plant >= num * Tave:
                # print("plant:",plant,num,Tave)
                if wrench[num][2] != 0:
                    # Convert wrench to numpy array if it's a list
                    coeff = 1
                    end_wrench = copy.deepcopy(wrench[num])  # No need to convert again since wrench is now a numpy array
                    print("end_wrench:",end_wrench,num)
                    num += 1

                    if np.abs(end_wrench[2]-start_wrench[2]) > self.force_vel_max*self.Ta:
                        Ta = np.floor(self.Ta /time_step) * time_step
                        self.force_vel = Ta * self.force_acc_max
                        Tj = (np.abs(end_wrench[2]-start_wrench[2])-self.force_vel*Ta)/self.force_vel
                        coeff = np.abs(end_wrench[2] - start_wrench[2]) / (np.floor((Ta + Tj)/time_step)*time_step * self.force_vel)
                    else:
                        Ta = np.floor(np.sqrt(np.abs(end_wrench[2]-start_wrench[2])/self.force_acc_max)/time_step)*time_step
                        Tj = 0
                        coeff = np.abs(end_wrench[2] - start_wrench[2]) / ((self.force_acc_max*time_step) * (np.floor(Ta/time_step)*np.floor(Ta/time_step))*time_step)

                    if Tave < 2*Ta + Tj:
                        Ta = np.floor(Tave/time_step/2)*time_step
                        Tj = 0
                        coeff = np.abs(end_wrench[2] - start_wrench[2]) / ((self.force_acc_max*time_step) * (np.floor(Ta/time_step)*np.floor(Ta/time_step))*time_step)
                    
                    if np.abs(end_wrench[2]-start_wrench[2]) < 0.1:
                        coeff = 1
                    
                    start_wrench = np.array(start_wrench).astype(float)  # Ensure it's an array
                    last_wrench = copy.deepcopy(start_wrench[2])
                    force_vel = 0
                    plant = 0
                    segment_wrenchs = copy.deepcopy(start_wrench)
                    dir = (end_wrench[2]-start_wrench[2])/np.abs(end_wrench[2]-start_wrench[2])
                    t_start = copy.deepcopy(start_wrench)
                else:

                    coeff = 1
                    end_wrench = copy.deepcopy(wrench[num]) 
                    num += 1
                    segment_wrenchs = copy.deepcopy(start_wrench)
                    if Tave > 2:
                        Ta = 1
                    else:
                        Ta = np.floor(Tave*10/2)/10
                    Tj = 0
                    force_vel = 0
                    plant = 0
                    coeff = np.abs(end_wrench[2] - start_wrench[2]) / ((self.force_acc_max*time_step) * (np.floor(Ta/time_step)*np.floor(Ta/time_step))*time_step)
                    dir = (end_wrench[2]-start_wrench[2])/np.abs(end_wrench[2]-start_wrench[2])
                    last_wrench = copy.deepcopy(start_wrench[2])
                    force_vel = 0
                    t_start = copy.deepcopy(start_wrench)

                # print("all_wrenchs:", all_wrenchs.shape)

            plant = phase * times[-1] - (num-1) * Tave

            if end_wrench[2] != 0:
                if plant == 0:
                    segment_wrenchs = copy.deepcopy(start_wrench)
                elif 0 < plant <= Ta:
                    # 加速阶段
                    force_vel += self.force_acc_max * time_step
                    if force_vel > self.force_vel_max:
                        force_vel = copy.deepcopy(self.force_vel_max)
                    last_wrench += dir * force_vel * time_step
                    segment_wrenchs[2] = last_wrench
                elif Ta < plant <= Ta + Tj:
                    # 恒速阶段
                    last_wrench += dir * force_vel * time_step
                    segment_wrenchs[2] = last_wrench
                elif Ta + Tj < plant <= 2 * Ta + Tj:
                    # 减速阶段
                    force_vel -= self.force_acc_max * time_step
                    if force_vel < 0:
                        force_vel = 0
                    last_wrench += dir * force_vel * time_step
                    segment_wrenchs[2] = last_wrench
                else:
                    # 超过所有阶段后，使用结束力矩
                    segment_wrenchs = segment_wrenchs
                
                all_wrenchs[count] = copy.deepcopy(segment_wrenchs) 
                all_wrenchs[count][2] = (all_wrenchs[count][2] - t_start[2]) * coeff + t_start[2]
                start_wrench = copy.deepcopy(all_wrenchs[count])
            else:
                if plant-(Tave-Tj-2*Ta) == 0:
                    segment_wrenchs = copy.deepcopy(start_wrench)
                elif 0 < plant-(Tave-Tj-2*Ta) <= Ta:
                    # 加速阶段
                    force_vel += self.force_acc_max * time_step
                    if force_vel > self.force_vel_max:
                        force_vel = copy.deepcopy(self.force_vel_max)
                    last_wrench += dir * force_vel * time_step
                    segment_wrenchs[2] = last_wrench
                elif Ta < plant-(Tave-Tj-2*Ta) <= Ta + Tj:
                    # 恒速阶段
                    last_wrench += dir * force_vel * time_step
                    segment_wrenchs[2] = last_wrench
                elif Ta + Tj < plant-(Tave-Tj-2*Ta) <= 2 * Ta + Tj:
                    # 减速阶段
                    force_vel -= self.force_acc_max * time_step
                    if force_vel < 0:
                        force_vel = 0
                    last_wrench += dir * force_vel * time_step
                    segment_wrenchs[2] = last_wrench
                else:
                    # 超过所有阶段后，使用结束力矩
                    segment_wrenchs = segment_wrenchs
                
                all_wrenchs[count] = copy.deepcopy(segment_wrenchs) 
                all_wrenchs[count][2] = (all_wrenchs[count][2] - t_start[2]) * coeff + t_start[2]
                start_wrench = copy.deepcopy(all_wrenchs[count])

        
        all_wrenchs = np.clip(all_wrenchs, np.array([0, 0, -70, 0, 0, 0]), np.array([0, 0, 0, 0, 0, 0]))

        return all_wrenchs


FP = ForcePlanner()  

# Example input for testing
start_wrench = [0, 0, 0, 0, 0, 0]  # Starting wrench
#end_wrench = [0, 0, 47, 0, 0, 0]
# end_wrench = [[0, 0, -29, 0, 0, 0]]  # Ending wrench
end_wrench = [[0, 0, -20, 0, 0, 0], [0, 0, -5, 0, 0, 0], [0, 0, -35, 0, 0, 0],[0, 0, -60, 0, 0, 0], [0, 0, 0, 0, 0, 0]]  # Ending wrench

time_points = [0, 1]  # From time 0 to 1
#time_points = None
time_step = 0.1  # 0.1 second steps

# Call the function with example values
result = FP.S_shaped_wrench_interpolate(start=start_wrench, wrench=end_wrench, time_points=time_points, time_step=0.0083333)

# z_values = result[:, 2]
# # for t, z in zip(times, z_values):
# #     print(f"time = {t:.4f} s,\tZ = {z:.4f}")

# plt.figure(figsize=(6, 4))
# plt.plot(times, z_values, label='Z (third component)')
# plt.xlabel('Time (s)')
# plt.ylabel('Force Z')
# plt.title('S-shaped 插值 —— 第三分量 (Z)')
# plt.grid(True)
# plt.legend()
# plt.tight_layout()
# plt.show()