from sympy.physics.mechanics import mechanics_printing
mechanics_printing()

from sympy import var, Function, pi

var("l1:4")

var("m1:4")

var("g t")

q1 = Function("q1")(t)
q2 = Function("q2")(t)
q3 = Function("q3")(t)

"""
Explanation: Práctica 3 - Dinámica de manipuladores
En esta práctica nuestro objetivo será simular el comportamiento de un manipulador tipo PUMA, empecemos importando las liberrias necesarias:
End of explanation
"""


def DH(params):
    from sympy import Matrix, sin, cos
    a, d, α, θ = params
    A = Matrix([[cos(θ), -sin(θ)*cos(α), sin(θ)*sin(α), a*cos(θ)],
                [sin(θ), cos(θ)*cos(α), -cos(θ)*sin(α), a*sin(θ)],
                [0, sin(α), cos(α), d],
                [0, 0, 0, 1]])
    return A

A1 = DH([0, l1, pi/2, q1])
A2 = DH([l2, 0, 0, q2])
A3 = DH([l3, 0, 0, q3])

As = [A1, A2, A3]
As

"""
Explanation: Y copiando la función para generar matrices de transformación homogéneas a partir de los parametros DH:
End of explanation
"""


def transf_art(transformaciones):
    from sympy import eye, simplify
    Hs = [eye(4)]
    for trans in transformaciones:
        Hs.append(simplify(Hs[-1]*trans))
        
    return Hs[1:]

Hs = transf_art(As)
Hs

"""
Explanation: He guardado todas las matrices de transformación homgénea en un solo arreglo, de tal manera que puedo hacer una función que tome todas las transformaciones de cada eslabon, y me devuelva las transformaciones a cada articulacion:
End of explanation
"""


ps = [H[0:3, 3:4] for H in Hs]
ps

"""
Explanation: Una vez obtenido esto, puedo obtener las posiciones de cada articulación con una List comprehension:
End of explanation
"""


# ESCRIBE TU CODIGO AQUI
raise NotImplementedError
Rs

from nose.tools import assert_equal
from sympy import Matrix, sin, cos, var

R1 = Matrix([[cos(q1), 0, sin(q1)],
             [sin(q1), 0, -cos(q1)],
             [0, 1, 0]])
R2 = Matrix([[cos(q1)*cos(q2), -sin(q2)*cos(q1), sin(q1)],
             [sin(q1)*cos(q2), -sin(q2)*sin(q1), -cos(q1)],
             [sin(q2), cos(q2), 0]])
R3 = Matrix([[cos(q1)*cos(q2+q3), -sin(q2+q3)*cos(q1), sin(q1)],
             [sin(q1)*cos(q2+q3), -sin(q2+q3)*sin(q1), -cos(q1)],
             [sin(q2+q3), cos(q2+q3), 0]])

assert_equal(Rs[0], R1)
assert_equal(Rs[1], R2)
assert_equal(Rs[2], R3)

"""
Explanation: Ejercicio

Genera una lista que contenga todas las matrices de rotación de cada articulación usando list comprehensions
End of explanation
"""


q = [q1, q2, q3]

"""
Explanation: Si ahora declaramos un vector con todos los grados de libertad:
End of explanation
"""


ps[1].jacobian(q)

"""
Explanation: podemos obtener el Jacobiano traslacional de cada articulacion con:
End of explanation
"""


# ESCRIBE TU CODIGO AQUI
raise NotImplementedError
Jvs

from nose.tools import assert_equal
assert_equal(Jvs[0], ps[0].jacobian(q))
assert_equal(Jvs[1], ps[1].jacobian(q))
assert_equal(Jvs[2], ps[2].jacobian(q))

"""
Explanation: Ejercicio

Genera una lista con los Jacobianos traslacionales
End of explanation
"""


o1 = Matrix([[0], [0], [q1]])
o1

o2 = Matrix([[0], [q2], [q1]])
o2

o3 = Matrix([[0], [q2 + q3], [q1]])
o3

"""
Explanation: Un paso que tenemos que hacer manualmente es definir los vectores de orientación (compuesto por $\phi$, $\theta$ y $\psi$) ya que se tiene un sistema sobrerestringido, pero son lo suficientemente faciles de obtener:
End of explanation
"""


os = [o1, o2, o3]

"""
Explanation: y si se guarda una lista con cada uno de estos vectores, se puede obtener el jacobiano rotacional de la misma manera que el traslacional:
End of explanation
"""


# ESCRIBE TU CODIGO AQUI
raise NotImplementedError
Jωs

from nose.tools import assert_equal
assert_equal(Jωs[0], os[0].jacobian(q))
assert_equal(Jωs[1], os[1].jacobian(q))
assert_equal(Jωs[2], os[2].jacobian(q))

"""
Explanation: Ejercicio

Genera una lista con los Jacobianos rotacionales
End of explanation
"""


def tens_iner(q):
    from sympy import Matrix
    Is = []
    for i in range(len(q)):
        Js = [var("J_{" + str(i+1) + "_" + eje + "}") for eje in "xyz"]
        I = Matrix([[Js[0], 0, 0], [0, Js[1], 0], [0, 0, Js[2]]])
        Is.append(I)
    return Is

Is = tens_iner(q)
Is

"""
Explanation: Otra cosa que podemos hacer en automatico es definir los tensores de inercia necesarios para el manipulador, ya que esto solo depende del numero de grados de libertad, defino la función que va a tomar el vector con el estado del sistema, $q$, y va a calcular una lista con los tensores:
End of explanation
"""


ms = [m1, m2, m3]

"""
Explanation: definiré una lista con todas las masas de los eslabones:
End of explanation
"""


def matriz_masas(ms, Jvs, Is, Jωs):
    from sympy import zeros, expand, simplify
    M = zeros(len(ms))
    for m, Jv, I, Jω in zip(ms, Jvs, Is, Jωs):
        M += simplify(expand(m*Jv.T*Jv + Jω.T*I*Jω))
    return M

M = matriz_masas(ms, Jvs, Is, Jωs)
M

"""
Explanation: De tal manera que podamos hacer una función que tome estas, los jacobianos y los tensores de inercia,para calcular la matriz de masas:
End of explanation
"""


from sympy import simplify

simplify(M)

"""
Explanation: mmm... un poco grande, tratemos de simplificar un poco:
End of explanation
"""


M[0].collect(l2**2).collect(l3**2).collect(m3).simplify()

"""
Explanation: mmm... un poco mejor, pero aun no es viable; los terminos del segundo y tercer grado de libertad son simples, el problema es el primero, tratemos de simplificar solo ese termino, intentaremos factorizar $l_2^2$ y $l_3^2$ y despues simplificar:
End of explanation
"""


M = simplify(M.applyfunc(lambda M: collect(M, l2**2)).applyfunc(lambda M: collect(M, l3**2)).applyfunc(lambda M: collect(M, m3)))
M

"""
Explanation: esto se ve aceptable, apliquemoslo a toda la matriz:
End of explanation
"""


def christoffel(M, q, i, j, k):
    from sympy import Rational, simplify
    # ESCRIBE TU CODIGO AQUI
    raise NotImplementedError
    return simplify(simbolo)

from nose.tools import assert_equal
from sympy import Rational, expand
assert_equal(christoffel(M, q, 0,0,1), expand(Rational(1,2)*((m2+m3)*l2**2*sin(2*q2) + m3*l3**2*sin(2*(q2+q3))) + m3*l2*l3*sin(2*q2+q3)))
assert_equal(christoffel(M, q, 0,0,0), 0)

"""
Explanation: Ejercicio

Escribe el código de una función, que dada la matriz de masas, una lista con los grados de libertad y los indices de posición, calcule el símbolo de Christoffel pedido, recuerde que la formula es:

$$
c_{ijk} = \frac{1}{2}\left{\frac{\partial M_{kj}}{\partial q_i} + \frac{\partial M_{ki}}{\partial q_j} - \frac{\partial M_{ij}}{\partial q_k}\right}
$$
End of explanation
"""


c113 = christoffel(M, q, 0,0,2)
c113

"""
Explanation: Con esta función podemos calcular cualquier simbolo de Christoffel (recordando que los indices en Python empiezan en $0$:
End of explanation
"""


def simbolos_chris(M, q):
    simbolos = []
    for i in range(len(q)):
        sim = []
        for j in range(len(q)):
            s = [christoffel(M, q, i, j, k) for k in range(len(q))]
            sim.append(s)
        simbolos.append(sim)
    return simbolos

simbolos_christoffel = simbolos_chris(M, q)

simbolos_christoffel[0][0][2]

"""
Explanation: y crear una función que calcule todos los simbolos de Christoffel a partir de esta función:
End of explanation
"""


def matriz_coriolis(simbolos, q̇):
    from sympy import Matrix
    coriolis = []
    for k in range(len(simbolos)):
        cor = []
        for j in range(len(simbolos)):
            c=0
            for i in range(len(simbolos)):
                c+= simbolos[i][j][k]*q̇[i]
            cor.append(c)
        coriolis.append(cor)
    return Matrix(coriolis)

C = simplify(matriz_coriolis(simbolos_christoffel, q̇))
C

"""
Explanation: Y ya con los simbolos de Christoffel, calcular la matriz de Coriolis:
End of explanation
"""


def ener_pot(params):
    m, h = params
    U = m*g*h
    return U

"""
Explanation: En este punto tenemos un resultado lo suficientemente compacto para copiarlo a las definiciones numéricas, por lo que seguimos al vector de gravedad:
End of explanation
"""


h1, h2, h3 = ps[0][2], ps[1][2], ps[2][2]

U1 = ener_pot([m1, h1])
U2 = ener_pot([m2, h2])
U3 = ener_pot([m3, h3])

U = U1 + U2 + U3

def vector_grav(U, q):
    from sympy import Matrix
    return Matrix([[U]]).jacobian(q).T

G = vector_grav(U, q)
G

"""
Explanation: Y calculando las energías potenciales:
End of explanation
"""


