from math import acos
from math import asin
from math import atan2
from math import cos
from math import fabs
from math import pi
from math import sin
from math import sqrt

from compas.geometry import Frame

from compas_fab.utilities import sign


def forward_kinematics_offset_wrist(joint_values, params):
    """Forward kinematics function for offset wrist 6-axis robots.

    Parameters
    ----------
    joint_values : list of float
        List of 6 joint values in radians.
    params : list of float
        The offset wrist parameters that specify the robot.

    Returns
    -------
    :class:`compas.geometry.Frame`

    Notes
    -----
    Code adapted from https://github.com/ros-industrial/universal_robot/blob/indigo-devel/ur_kinematics/src/ur_kin.cpp
    """
    d1, a2, a3, d4, d5, d6 = params
    q = joint_values

    s1, c1 = sin(q[0]), cos(q[0])
    q23, q234, s2, c2 = q[1], q[1], sin(q[1]), cos(q[1])
    q23 += q[2]
    q234 += q[2]
    s4, c4 = sin(q[3]), cos(q[3])
    q234 += q[3]
    s5, c5 = sin(q[4]), cos(q[4])
    s6, c6 = sin(q[5]), cos(q[5])
    s23, c23 = sin(q23), cos(q23)
    s234, c234 = sin(q234), cos(q234)

    T = [0.0 for _ in range(4 * 4)]

    T[0] = c234 * c1 * s5 - c5 * s1
    T[1] = c6 * (s1 * s5 + c234 * c1 * c5) - s234 * c1 * s6
    T[2] = -s6 * (s1 * s5 + c234 * c1 * c5) - s234 * c1 * c6
    T[3] = d6 * c234 * c1 * s5 - a3 * c23 * c1 - a2 * c1 * c2 - d6 * c5 * s1 - d5 * s234 * c1 - d4 * s1
    T[4] = c1 * c5 + c234 * s1 * s5
    T[5] = -c6 * (c1 * s5 - c234 * c5 * s1) - s234 * s1 * s6
    T[6] = s6 * (c1 * s5 - c234 * c5 * s1) - s234 * c6 * s1
    T[7] = d6 * (c1 * c5 + c234 * s1 * s5) + d4 * c1 - a3 * c23 * s1 - a2 * c2 * s1 - d5 * s234 * s1
    T[8] = -s234 * s5
    T[9] = -c234 * s6 - s234 * c5 * c6
    T[10] = s234 * c5 * s6 - c234 * c6
    T[11] = d1 + a3 * s23 + a2 * s2 - d5 * (c23 * c4 - s23 * s4) - d6 * s5 * (c23 * s4 + s23 * c4)
    T[15] = 1.0

    frame = Frame((T[3], T[7], T[11]), (T[0], T[4], T[8]), (T[1], T[5], T[9]))

    return frame


def inverse_kinematics_offset_wrist(frame, params, q6_des=0.0):
    """Inverse kinematics function for offset wrist 6-axis robots.

    Parameters
    ----------
    frame : :class:`compas.geometry.Frame`
        The frame we want to calculate the inverse kinematics for.
    params : list of float
        The offset wrist parameters that specify the robot.

    Returns
    -------
    list of list

    Notes
    -----
    Code adapted from https://github.com/ros-industrial/universal_robot/blob/indigo-devel/ur_kinematics/src/ur_kin.cpp
    """

    ZERO_THRESH = 0.00000001
    d1, a2, a3, d4, d5, d6 = params

    solutions = []

    T02, T12, T22 = frame.xaxis
    T00, T10, T20 = frame.yaxis
    T01, T11, T21 = frame.zaxis
    T03, T13, T23 = frame.point
    T02 *= -1
    T03 *= -1
    T12 *= -1
    T13 *= -1
    T20 *= -1
    T21 *= -1

    # shoulder rotate joint (q1)
    # q1[2]
    q1 = [0, 0]
    A = d6 * T12 - T13
    B = d6 * T02 - T03
    R = A * A + B * B
    if fabs(A) < ZERO_THRESH:
        div = 0.0
        if fabs(fabs(d4) - fabs(B)) < ZERO_THRESH:
            div = -sign(d4) * sign(B)
        else:
            div = -d4 / B
        arcsin = asin(div)
        if fabs(arcsin) < ZERO_THRESH:
            arcsin = 0.0
        if arcsin < 0.0:
            q1[0] = arcsin + 2.0 * pi
        else:
            q1[0] = arcsin
        q1[1] = pi - arcsin

    elif fabs(B) < ZERO_THRESH:
        div = 0.0
        if fabs(fabs(d4) - fabs(A)) < ZERO_THRESH:
            div = sign(d4) * sign(A)
        else:
            div = d4 / A
        arccos = acos(div)
        q1[0] = arccos
        q1[1] = 2.0 * pi - arccos

    elif d4 * d4 > R:
        raise ValueError("No solutions")
    else:
        arccos = acos(d4 / sqrt(R))
        arctan = atan2(-B, A)
        pos = arccos + arctan
        neg = -arccos + arctan
        if fabs(pos) < ZERO_THRESH:
            pos = 0.0
        if fabs(neg) < ZERO_THRESH:
            neg = 0.0
        if pos >= 0.0:
            q1[0] = pos
        else:
            q1[0] = 2.0 * pi + pos
        if neg >= 0.0:
            q1[1] = neg
        else:
            q1[1] = 2.0 * pi + neg

    # wrist 2 joint (q5)
    q5 = [[0, 0], [0, 0]]
    for i in range(2):
        numer = T03 * sin(q1[i]) - T13 * cos(q1[i]) - d4
        div = 0.0
        if fabs(fabs(numer) - fabs(d6)) < ZERO_THRESH:
            div = sign(numer) * sign(d6)
        else:
            div = numer / d6
        arccos = acos(div)
        q5[i][0] = arccos
        q5[i][1] = 2.0 * pi - arccos

    for i in range(2):
        for j in range(2):
            c1 = cos(q1[i])
            s1 = sin(q1[i])
            c5 = cos(q5[i][j])
            s5 = sin(q5[i][j])
            q6 = 0.0

            # wrist 3 joint (q6)
            if fabs(s5) < ZERO_THRESH:
                q6 = q6_des
            else:
                q6 = atan2(sign(s5) * -(T01 * s1 - T11 * c1), sign(s5) * (T00 * s1 - T10 * c1))
            if fabs(q6) < ZERO_THRESH:
                q6 = 0.0
            if q6 < 0.0:
                q6 += 2.0 * pi

            # RRR joints (q2,q3,q4)
            q2, q3, q4 = [0, 0], [0, 0], [0, 0]

            c6 = cos(q6)
            s6 = sin(q6)
            x04x = -s5 * (T02 * c1 + T12 * s1) - c5 * (s6 * (T01 * c1 + T11 * s1) - c6 * (T00 * c1 + T10 * s1))
            x04y = c5 * (T20 * c6 - T21 * s6) - T22 * s5
            p13x = (
                d5 * (s6 * (T00 * c1 + T10 * s1) + c6 * (T01 * c1 + T11 * s1))
                - d6 * (T02 * c1 + T12 * s1)
                + T03 * c1
                + T13 * s1
            )
            p13y = T23 - d1 - d6 * T22 + d5 * (T21 * c6 + T20 * s6)

            c3 = (p13x * p13x + p13y * p13y - a2 * a2 - a3 * a3) / (2.0 * a2 * a3)
            if fabs(fabs(c3) - 1.0) < ZERO_THRESH:
                c3 = sign(c3)
            elif fabs(c3) > 1.0:
                solutions.extend([None, None])
                continue

            arccos = acos(c3)
            q3[0] = arccos
            q3[1] = 2.0 * pi - arccos
            denom = a2 * a2 + a3 * a3 + 2 * a2 * a3 * c3
            s3 = sin(arccos)
            A = a2 + a3 * c3
            B = a3 * s3
            q2[0] = atan2((A * p13y - B * p13x) / denom, (A * p13x + B * p13y) / denom)
            q2[1] = atan2((A * p13y + B * p13x) / denom, (A * p13x - B * p13y) / denom)
            c23_0 = cos(q2[0] + q3[0])
            s23_0 = sin(q2[0] + q3[0])
            c23_1 = cos(q2[1] + q3[1])
            s23_1 = sin(q2[1] + q3[1])
            q4[0] = atan2(c23_0 * x04y - s23_0 * x04x, x04x * c23_0 + x04y * s23_0)
            q4[1] = atan2(c23_1 * x04y - s23_1 * x04x, x04x * c23_1 + x04y * s23_1)

            for k in range(2):
                if fabs(q2[k]) < ZERO_THRESH:
                    q2[k] = 0.0
                elif q2[k] < 0.0:
                    q2[k] += 2.0 * pi
                if fabs(q4[k]) < ZERO_THRESH:
                    q4[k] = 0.0
                elif q4[k] < 0.0:
                    q4[k] += 2.0 * pi

                solutions.append([q1[i], q2[k], q3[k], q4[k], q5[i][j], q6])

    return solutions


if __name__ == "__main__":
    from compas.geometry import allclose

    params = [0.089159, -0.42500, -0.39225, 0.10915, 0.09465, 0.0823]  # ur5
    q = [0.2, 5.5, 1.4, 1.3, 2.6, 3.6]
    frame = forward_kinematics_offset_wrist(q, params)
    sol = inverse_kinematics_offset_wrist(frame, params)
    assert allclose(sol[0], q)
