import sympy as sp
import numpy as np
from sympy.vector import CoordSys3D
# This solution is for the first three joints of the Stanford arm
# Define symbolic variables
th1, th2, L1, L2, d3 = sp.symbols("th1 th2 L1 L2 d3") #d3 includes the joint variable and the link length L3
N = CoordSys3D('N')

# Define the transformation matrices for each link using given DH parameters
A1 = sp.Matrix([
    [sp.cos(th1), 0, -sp.sin(th1), 0],
    [sp.sin(th1), 0, sp.cos(th1), 0],
    [0, -1, 0, L1],
    [0, 0, 0, 1]
]) 
A2 = sp.Matrix([
    [sp.cos(th2), 0, sp.sin(th2), 0],
    [sp.sin(th2), 0, -sp.cos(th2), 0],
    [0, 1, 0, L2],
    [0, 0, 0, 1]
])

A3 = sp.Matrix([
    [1, 0, 0, 0],
    [0, 1, 0, 0],
    [0, 0, 1, d3],
    [0, 0, 0, 1]
])
# find z-coordinates
z_0 = N.k
z_1 = A1[0:3, 2]
z_1 = z_1[0]*N.i + z_1[1]*N.j + z_1[2]*N.k
z_2 = (A1 * A2)[0:3, 2]
z_2 = z_2[0]*N.i + z_2[1]*N.j + z_2[2]*N.k

def jacobian_v(A1,A2,A3,z0,z1,z2):
    T1 = A1
    T2 = A1 * A2
    T3 = sp.simplify(T2* A3)

    three_minus_null=(T3[0,3]*N.i)+(T3[1,3]*N.j)+(T3[2,3]*N.k)
    three_minus_one = (T3[0,3]-T1[0,3])*N.i+(T3[1,3]-T1[1,3])*N.j+(T3[2,3]-T1[2,3])*N.k

    z0xthree_minus_null = sp.simplify(z0.cross(three_minus_null))
    z1xthree_minus_one = sp.simplify(z1.cross(three_minus_one))

    z0xthree_minus_null_entries = [z0xthree_minus_null.dot(N.i), z0xthree_minus_null.dot(N.j), z0xthree_minus_null.dot(N.k)]
    z1xthree_minus_one_entries = [z1xthree_minus_one.dot(N.i), z1xthree_minus_one.dot(N.j), z1xthree_minus_one.dot(N.k)]
    z2x = [z2.dot(N.i), z2.dot(N.j), z2.dot(N.k)]

    Jv = sp.Matrix([z0xthree_minus_null_entries, z1xthree_minus_one_entries, z2x]).T
    print("Jv")
    sp.pretty_print(Jv)
    return Jv
  
def determinant(Jv):
    d = sp.simplify(Jv.det())
    print("Det")
    sp.pretty_print(d)
    d_solution = sp.solve(d)
    print("Solutions:")
    sp.pretty_print(d_solution)

Jvel = jacobian_v(A1, A2, A3, z_0,z_1,z_2)
determinant(Jvel)