"""Kinematics calculations for a two-link serial chain (e.g., double-pendulum).

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.

Q0:     shoulder angle in radians, zero pointing down, positive is CCW
Q1:     elbow angle in radians relative to upper link, zero pointing down, positive is CCW

"""
import sympy as sym

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

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

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

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 - q1))
soln = sym.solve(elbowEqn, [q1])

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 - q1), sym.sqrt(R_squared) * sym.sin(theta - q0))

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

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)