import numpy as np
from functools import reduce
from mpl_toolkits.mplot3d import Axes3D
import matplotlib.pyplot as plt
from matplotlib import cm
import math


#旋转矩阵
def rotate(axis,deg):
    AXIS = ('X', 'Y', 'Z')
    axis = str(axis).upper()
    if axis not in AXIS:
        print(f"{axis} is unknown axis, should be one of {AXIS}")
        return
    rot_x = (axis == 'X')
    rot_y = (axis == 'Y')
    rot_z = (axis == 'Z')
    rot_mat = np.array([[(np.cos(deg), 1)[rot_x], (0, -np.sin(deg))[rot_z], (0, np.sin(deg))[rot_y], 0],
                        [(0, np.sin(deg))[rot_z], (np.cos(deg), 1)[rot_y], (0, -np.sin(deg))[rot_x], 0],
                        [(0, -np.sin(deg))[rot_y], (0, np.sin(deg))[rot_x], (np.cos(deg), 1)[rot_z], 0],
                        [0, 0, 0, 1]], dtype=np.float32)
    rot_mat = np.where(np.abs(rot_mat) < 1e-10, 0, rot_mat)  # get a small value when np.cos(np.pi/2)
    return rot_mat

#平移矩阵
def trans(axis, dis):
    AXIS = ('X', 'Y', 'Z')
    axis = str(axis).upper()
    if axis not in AXIS:
        print(f"{axis} is unknown axis, should be one of {AXIS}")
        return
    trans_mat = np.eye(4)
    trans_mat[AXIS.index(axis), 3] = dis
    return trans_mat
#根据DH参数计算总的转换矩阵
def DHmartix(thea, l, d, a):
    T=np.mat(  [[np.cos(thea), -np.sin(thea), 0, l],
    [np.sin(thea)*np.cos(a), np.cos(thea)*np.cos(a), -np.sin(a), -np.sin(a)*d],
    [np.sin(thea)*np.sin(a), np.cos(thea)*np.sin(a), np.cos(a), np.cos(a)*d],
    [0, 0, 0, 1]], dtype=np.float16)
    return np.around(T,4)
def fk(joints):
    thea_1, thea_2, thea_3, thea_4, thea_5= joints
    # DH_pramater: [theta,l_i-1,d_i,a_i-1]，注意这里的单位是m
    DH = [[thea_1, 0, 0, 0],
          [thea_2, 0, 0, -np.pi/2],
          [thea_3, 0.105, 0, 0],
          [thea_4, 0.098, 0, 0],
          [thea_5, 0.063,0, -np.pi/2]]
    T = [rotate('z', thea_i).dot(trans('z', d_i)).dot(trans('x', l_i)).dot(rotate('x', a_i))
         for thea_i, l_i, d_i, a_i in DH]
    T_formula = [DHmartix(thea_j, l_j, d_j, a_j)
                    for thea_j, l_j, d_j, a_j in DH]
    T50 = reduce(np.dot,T_formula)
    return  T50,T_formula

# joint = [0,np.pi/2,0,0,-np.pi/2]
# T50,T = fk(joint)
# print(T)
# print(T50)

def get_jac(joints: np.ndarray):
    delta = 0.0001
    jac = np.zeros((16, joints.shape[0]))
    for i, joint in enumerate(joints):
        joints_m = joints.copy()
        joints_p = joints.copy()
        joints_m[i] -= delta
        joints_p[i] += delta
        Tm,_ = fk(joints_m)
        Tp,_ = fk(joints_p)
        jac[:, i] = (Tp - Tm).flatten() / (2 * delta)
    return jac

#Jacobian 迭代求解IK
def ik(T_tar, joints_init = np.zeros(5), tolerance = 1e-7):
    itertime = 0
    step = 0.5
    joints = joints_init.copy()
    while itertime < 1000:
        T_cur,_ = fk(joints)
        deltaT = (T_tar - T_cur).flatten()
        error = np.linalg.norm(deltaT)
        if error < tolerance:
            return joints
        jac = get_jac(joints)
        deltaq = np.linalg.pinv(jac) @ deltaT
        joints = joints + step * deltaq
        itertime += 1
    return False

#逆运动学解析解，仅对本机械臂适用
def ana_ik(p_5,pitch = None):
        x_t,y_t,z_t = p_5
        theta = math.atan2(z_t,math.sqrt(x_t**2+y_t**2))
        l1, l2 = 0.105, 0.098
        l3 = 0.063
        ik = []
        r=math.sqrt(x_t**2+y_t**2+z_t**2)
        if(pitch == None):
            beta=(r**2-0.03724)/(0.126*r)
            if(beta>1):
                print("目的点不可达")
                
                x_t *= 0.26/r
                y_t *= 0.26/r
                z_t *= 0.26/r
                p_5 = [x_t ,y_t ,z_t]
                beta = 1
                print("距离最近的可达点为：",p_5)
                # return None
            elif(beta < 0.2964 and beta >-1):#所计算的关节角令关节大于75.76°，以及beta不超过180°，则选取小于75.76°的值
                beta = 0.342                 #选取beta = 70°
                # beta=180*math.acos(beta)/np.pi
            elif(beta < -1):                 #如果beta大于180° ， 那就利用下面的公式重新选取beta角
                beta = math.atan2(l2, l3) + math.acos((0.1165**2 + r**2 - l1**2)/(2*0.1165*l1))
            else :                           #如果beta介于 0 ~ 75.76° ，选取0.8倍的beta角
                beta = 180*math.acos(beta)/np.pi 
                beta *= 0.8
            # beta = 2*int(beta)/3
        else :
            beta = - pitch + theta 
        #beta = int(beta)
        print('beta:',beta)
        print('theta:',theta)
        # alpha = beta/180*np.pi
        alpha = beta
        if z_t == 0:
            z_t +=1e-10
        if (x_t**2 + y_t**2) == 0:
            z = l3*math.cos(alpha)
            x = -l3*math.sin(alpha)
            y = 1e-7
        else:
            
            h2 = z_t - (l3 * math.cos(theta - alpha) * math.tan(theta))
            H = math.sqrt(x_t ** 2 + y_t ** 2 + z_t ** 2) / math.sqrt(x_t ** 2 + y_t ** 2) * math.sin(alpha) * l3
            z = H + h2
            x = x_t / z_t * h2
            y = y_t / z_t * h2
            if y == 0:
                y +=1e-7

        theta_1 = math.atan2(y, x)
        theta_3 = math.acos((x ** 2 + y ** 2 + z ** 2 - (l1 ** 2 + l2 ** 2)) / 2 * l1 * l2)

        k1 = math.cos(theta_3) * l2 + l1
        k2 = math.sin(theta_3) * l2

        theta_2 = - math.atan2(k2, k1) - math.atan2(z, math.sqrt(x ** 2 + y ** 2))
        theta_4 = - theta + alpha -theta_2 - theta_3
        if x < 0 :
            theta_1 = theta_1 -np.sign(y)*np.pi
            theta_2 = - theta_2 - np.pi
            theta_3 = - theta_3
            theta_4 = - theta_4
            

        ik=[theta_1,theta_2,theta_3,theta_4]
        pitch = alpha - theta
        return ik,pitch


def ana_ik_Special_1(p_5) :
    x_t,y_t,z_t = p_5
    l1, l2 = 0.105, 0.098
    l3 = 0.063
    ik = []
    x = x_t - x_t * math.sqrt(x_t**2+y_t**2) * l3
    y = y_t - y_t * math.sqrt(x_t**2+y_t**2) * l3
    z = z_t
    theta_1 = math.atan2(y, x)
    theta_3 = math.acos((x ** 2 + y ** 2 + z ** 2 - (l1 ** 2 + l2 ** 2)) / 2 * l1 * l2)
    k1 = math.cos(theta_3) * l2 + l1
    k2 = math.sin(theta_3) * l2
    theta_2 = - math.atan2(k2, k1) - math.atan2(z, math.sqrt(x ** 2 + y ** 2))
    theta_4 =  -theta_2 - theta_3
    ik=[theta_1,theta_2,theta_3,theta_4]
    return ik

def ana_ik_forward(p_r, p_t, forward):  # 分别沿x,y,z三个方向移动
    x_t, y_t, z_t = p_t
    x_r, y_r, z_r = p_r
    theta = math.atan2(z_t, math.sqrt(x_t ** 2 + y_t ** 2))
    l1, l2 = 0.105, 0.098
    l3 = 0.063
    ik = []
    r = math.sqrt(x_t ** 2 + y_t ** 2 + z_t ** 2)

    if (forward == ('x' or 'X')):
        r_min = math.sqrt(y_t ** 2 + z_t ** 2)
        if r_min < 0.1:
            return None
        else:
            for i in range(4):
                ik.append([x_r + (i + 1) * x_t / 4, y_t, z_t])
            return ik
    elif (forward == ('y' or 'Y')):
        r_min = math.sqrt(x_t ** 2 + z_t ** 2)
        if r_min < 0.1:
            return None
        else:
            for i in range(4):
                ik.append([x_t, y_r + (i + 1) * y_t / 4, z_t])
            return ik
    elif forward == ('z' or 'Z'):
        r_min = math.sqrt(x_t ** 2 + y_t ** 2)
        if r_min < 0.1:
            return None
        else:
            for i in range(4):
                ik.append([x_t, y_t, z_r + (i + 1) * z_t / 4])
            return ik
    else:
        return None


def workspace(p_t):
    x_t, y_t, z_t = p_t
    r = math.sqrt(x_t ** 2 + y_t ** 2 + z_t ** 2)
    if r > 0.266:
        print("目标不可达")
        x_t *= 0.26/r
        y_t *= 0.26/r
        z_t *= 0.26/r
        p_t = [x_t ,y_t ,z_t]
        
    return p_t
def Geometric(ik):
    l1, l2 = 0.105, 0.098
    l3 = 0.063
    theta1 ,theta2 ,theta3 ,theta4 =ik
    r = l1*np.cos(theta2) + l2*np.cos(theta2 + theta3) + l3*np.cos(theta2 + theta3 + theta4)
    x = r*np.cos(theta1)
    y = r*np.sin(theta1)
    z = -l1*np.sin(theta2) - l2*np.sin(theta2 + theta3) + l3*np.sin(theta2 + theta3 + theta4)
    Descartes=[x,y,z]
    return Descartes



# def jtra_plan(start_position,end_position,num_point):