"""Kinematics calculations for the two-link ZYY arm.

This model has a Z axis base rotation, followed by a two link arm with two Y
axis joints.

Given the link lengths, calculate the endpoint positions.

Note that since this is a serial drive mechanism, the forward kinematics is
simpler than the inverse kinematics.  The reference pose is straight up; each
joint is measured using the right-hand rule around the axis.

q0:     base Z-axis rotation angle in radians
q1:     shoulder Y-axis rotation angle in radians, zero pointing up
q2:     elbow Y-axis rotation angle in radians relative to upper link, zero pointing up

L0:     Z offset from floor to Y axis of joint1
L1:     length of link1, i.e., distance between joint1 and joint2 axes
L2:     length of link2

"""
import sympy as sym

##### forward kinematics #############
L0, L1, L2, x, y, q0, q1, q2, theta = sym.symbols('L0 L1 L2 x y q0 q1 q2 theta')

# The elbow location is uniquely by the shoulder angle q1.
elbow = sym.Matrix([L1 * sym.sin(q1), -L1 * sym.cos(q1)])

# The endpoint location locus is a circle centered on the elbow.
endpoint = elbow + sym.Matrix([L2 * sym.sin(q1+q2), -L2 * sym.cos(q1+q2)])

print("\n\nForward kinematics, i.e. endpoint location as a function of joint angles:")
print(endpoint)

##### inverse kinematics #############

# the square of the target distance
R_squared = x**2 + y**2

# apply the law of cosines to calculate the elbow angle based on the target distance and the link lengths
elbowEqn = sym.Eq(R_squared, L1**2 + L2**2 - 2*L1*L2*sym.cos(sym.pi - q2))
soln = sym.solve(elbowEqn, [q2])

print("\n\nElbow solution:", soln)

# the overall angle from the center to the endpoint
theta = sym.atan2(x, -y)

# apply the law of sines
alphaEqn = sym.Eq(L2 * sym.sin(sym.pi - q2), sym.sqrt(R_squared) * sym.sin(theta - q1))

soln = sym.solve([elbowEqn, alphaEqn], [q1, q2])

print("\n\nInverse kinematics, i.e. joint angles as a function of endpoint location:")
print(soln)

##### demonstrate a solution ##############

values = {L1:0.5, L2:0.5, x:0.01, y:-0.99}
angles = [[expr[0].evalf(subs=values), expr[1].evalf(subs=values)] for expr in soln]

print("\n\nangle solutions:", angles)