from sympy import var, sin, cos, pi, Matrix, Function, Rational, simplify
from sympy.physics.mechanics import mechanics_printing
mechanics_printing()

"""
Explanation: Práctica 2 - Cinemática directa y dinámica de manipuladores
En este documento se describe el proceso de obtención de la dinámica de un robot manipulador (pendulo doble), por medio de la ecuación de Euler-Lagrange, empecemos importando las librerias necesarias:
End of explanation
"""


var("l1:3")

var("m1:3")

var("J1:3")

var("g t")

"""
Explanation: Definimos de una vez todas las variables necesarias:
End of explanation
"""


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

"""
Explanation: Y definimos las variables que dependen de otra variable, especificamente en este calculo, todo lo anterior es constante y solo $q_1$ es una variable dependiente del tiempo:
End of explanation
"""


def DH(params):
    from sympy import Matrix, sin, cos
    a, d, α, θ = params
    # ESCRIBE TU CODIGO AQUI
    raise NotImplementedError
    return A

from nose.tools import assert_equal
from sympy import var, Matrix, sin, cos, eye
var("l1")
q1 = Function("q1")(t)
A = Matrix([[cos(q1), -sin(q1), 0, l1*cos(q1)],
            [sin(q1), cos(q1), 0, l1*sin(q1)],
            [0, 0, 1, 0],
            [0, 0, 0, 1]])
assert_equal(DH([l1, 0, 0, q1]), A)
assert_equal(DH([0, 0, 0, 0]), eye(4))
print("Sin errores")

"""
Explanation: Esta vez vamos a hacer algo ligeramente diferente, vamos a automatizar un poco el proceso de obtención de la posición generalizada del manipulador, para esto vamos a apoyarnos de nuestros conocimientos de cinemática directa, para lo cual, primero necesitamos definir una función DH la cual tome una lista de parametros, en un orden especifico, y nos devuelva la matriz de transformación homogénea asociada a este eslabón:

Ejercicio
Define una función DH que tome los parametros Denavit-Hartenberg de un eslabón de un manipulador y devuelva la matriz de transformación homogénea asociada a este eslabón.
End of explanation
"""


A1 = DH([l1, 0, 0, q1])
A2 = DH([l2, 0, 0, q2])

A1

"""
Explanation: Una vez que tenemos la función DH para calcular las matrices de transformación homogéneas, ahora procedemos a calcular la posición de cada articulación por medio de estas matrices:
End of explanation
"""


p1 = A1[0:3, 3:4]
p1

"""
Explanation: Recordemos que la posición de cada articulación se obtendrá por medio del ultimo vector de esta matriz, por lo que podemos:
End of explanation
"""


p2 = (A1*A2)[0:3, 3:4]
p2

"""
Explanation: Para la posición de la segunda articulación necesitamos multiplicar las primeras dos matrices:
End of explanation
"""


p2 = simplify(p2)
p2

"""
Explanation: Aunque en este caso, podemos simplificar mas estas expresiones:
End of explanation
"""


v1 = p1.diff("t")
v1

v2 = p2.diff("t")
v2

"""
Explanation: Teniendo estas posiciones, para obtener la velocidad, necesitamos obtener su derivada:
End of explanation
"""


v1c = (v1.T*v1)[0]
v1c = v1c.simplify()
v1c

v2c = (v2.T*v2)[0]
v2c = v2c.simplify()
v2c

"""
Explanation: Una vez que tenemos la velocidad, obtener el cuadrado de esta velocidad es facil, para un vector podemos decir que:
$$
\left| v \right|^2 = v^T \cdot v
$$
End of explanation
"""


h1, h2 = p1[1], p2[1]
ω1, ω2 = q1.diff(t), q1.diff(t) + q2.diff(t)

"""
Explanation: Y calculando la altura y velocidad rotacional del eslabon:
End of explanation
"""


def ener_cin(params):
    from sympy import Rational
    m, v, J, ω = params
    # ESCRIBE TU CODIGO AQUI
    raise NotImplementedError
    return K

from nose.tools import assert_equal
from sympy import var, Matrix, sin, cos, Rational
var("m1 J1 l1 ω1")
q1 = Function("q1")(t)
v1 = Matrix([[l1*cos(q1)], [l1*sin(q1)], [0]])
assert_equal(ener_cin([m1, v1, J1, ω1]), Rational(1,2)*m1*l1**2 + Rational(1,2)*J1*ω1**2)
assert_equal(ener_cin([0, Matrix([[0],[0],[0]]), 0, 0]), 0)
print("Sin errores")

"""
Explanation: Ejercicio
Define una función ener_cin, la cual tome los parametros m, v, J y ω y devuelva la energía cinética del eslabon.
End of explanation
"""


h1, h2 = p1[1], p2[1]
ω1, ω2 = q1.diff(t), q1.diff(t) + q2.diff(t)

K1 = ener_cin([m1, v1, J1, ω1])
K1

K2 = ener_cin([m2, v2, J2, ω2])
K2

"""
Explanation: Si ahora calculamos las energías cinéticas, tenemos:
End of explanation
"""


def ener_pot(params):
    m, h = params
    # ESCRIBE TU CODIGO AQUI
    raise NotImplementedError
    return U

from nose.tools import assert_equal
from sympy import var
var("m1 m2 g h1 h2")
assert_equal(ener_pot([m1, h1]), m1*g*h1)
assert_equal(ener_pot([m2, h2]), m2*g*h2)
print("Sin errores")

"""
Explanation: Ejercicio
Defina una función ener_pot, la cual tome los parametros m y h y devuelva la energía potencial del eslabon.
End of explanation
"""


h1, h2 = p1[1], p2[1]
ω1, ω2 = q1.diff(t), q1.diff(t) + q2.diff(t)

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

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


K = K1 + K2
U = U1 + U2

"""
Explanation: Una vez que tenemos las energías cinéticas y potenciales de cada eslabón, podemos calcular la energía cinética total y la energía potencial total del manipulador:
End of explanation
"""


L = K - U

L

"""
Explanation: Con estas energias se puede calcular el Lagrangiano:
End of explanation
"""


τ1 = (L.diff(q1.diff(t)).diff(t) - L.diff(q1)).expand().collect(q1.diff(t).diff(t)).collect(q2.diff(t).diff(t))

τ1

τ2 = (L.diff(q2.diff(t)).diff(t) - L.diff(q2)).expand().collect(q1.diff(t).diff(t)).collect(q2.diff(t).diff(t))

τ2

"""
Explanation: Ya con el Lagrangiano, podemos calcular la ecuación de Euler-Lagrange para cada grado de libertad del manipulador:
End of explanation
"""


