import sympy
import numpy
from math import cos, sin

import sympybotics
from sympybotics._compatibility_ import exec_





def test_scara_dh_sym_geo_kin():

    pi = sympy.pi
    q = sympybotics.robotdef.q

    a1, a2, d3, d4 = sympy.symbols('a1, a2, d3, d4')

    scara = sympybotics.robotdef.RobotDef(
        'SCARA - Spong',
        [( 0, a1,  0, q),
         (pi, a2,  0, q),
         ( 0,  0,  q, 0),
         ( 0,  0, d4, q)],
        dh_convention='standard')

    scara_geo = sympybotics.geometry.Geometry(scara)
    scara_kin = sympybotics.kinematics.Kinematics(scara, scara_geo)

    cos, sin = sympy.cos, sympy.sin
    q1, q2, q3, q4 = sympy.flatten(scara.q)

    T_spong = sympy.Matrix([
        [(-sin(q1)*sin(q2) + cos(q1)*cos(q2))*cos(q4) + (sin(q1)*cos(q2) +
         sin(q2)*cos(q1))*sin(q4), -(-sin(q1)*sin(q2) +
         cos(q1)*cos(q2))*sin(q4) + (sin(q1)*cos(q2) +
         sin(q2)*cos(q1))*cos(q4), 0, a1*cos(q1) - a2*sin(q1)*sin(q2) +
         a2*cos(q1)*cos(q2)],
        [(sin(q1)*sin(q2) - cos(q1)*cos(q2))*sin(q4) + (sin(q1)*cos(q2) +
         sin(q2)*cos(q1))*cos(q4), (sin(q1)*sin(q2) -
         cos(q1)*cos(q2))*cos(q4) - (sin(q1)*cos(q2) +
         sin(q2)*cos(q1))*sin(q4), 0, a1*sin(q1) + a2*sin(q1)*cos(q2) +
         a2*sin(q2)*cos(q1)],
        [0, 0, -1, -d4 - q3],
        [0, 0, 0, 1]])

    J_spong = sympy.Matrix([[-a1*sin(q1) - a2*sin(q1)*cos(q2) -
                             a2*sin(q2)*cos(q1), -a2*sin(q1)*cos(q2) -
                             a2*sin(q2)*cos(q1), 0, 0],
                            [a1*cos(q1) - a2*sin(q1)*sin(q2) +
                             a2*cos(q1)*cos(q2), -a2*sin(q1)*sin(q2) +
                             a2*cos(q1)*cos(q2), 0, 0],
                            [0, 0, -1, 0],
                            [0, 0, 0, 0],
                            [0, 0, 0, 0],
                            [1, 1, 0, -1]])

    assert (scara_geo.T[-1] - T_spong).expand() == sympy.zeros(4)
    assert (scara_kin.J[-1] - J_spong).expand() == sympy.zeros(6, 4)


def test_puma_dh_num_geo_kin_dyn():
    pi = sympy.pi
    q = sympybotics.robotdef.q

    puma560_def = sympybotics.RobotDef(
        'Puma 560 Robot',
        [(   pi/2,        0,         0,   q),
         (      0,   0.4318,         0,   q),
         ('-pi/2', '0.0203', '0.15005', 'q'),  # test sympify
         (   pi/2,        0,    0.4318,   q),
         (  -pi/2,        0,         0,   q),
         (      0,        0,         0,   q)],
        dh_convention='standard')

    puma560_def.frictionmodel = None

    puma560 = sympybotics.RobotDynCode(puma560_def)

    q_test = [0.7504516826728697, 0.8395156106908136, 0.16851233582594916, 0.3849629637427072, 0.5252993946810777, 0.6701207256444748]
    dq_test = [0.24721855939629367, 0.9805915670454258, 0.9895299755642817, 0.7861135739668947, 0.273842245476577, 0.17182358900767503]
    ddq_test = [0.707405815485141, 0.25295715193420953, 0.9763909835998361, 0.8412822676113918, 0.4867768296473465, 0.11480270540937143]

    dynparm_test = [0.0, -0, -0, 0.34999999999999998, -0, 0.0, 0.0, 0.0, 0.0, 0, 1.03118515, 0.037980719999999996, 1.4401022999999999,
                    3.7274564059999999, -0.023751000000000001, 2.8425240559999998, -6.33012, 0.10439999999999999, 3.9585, 17.4,
                    0.090474288000000014, -0.0013739039999999998, 0.0068208000000000001, 0.111498032, 0.0047375999999999998,
                    0.015432319999999999, -0.09743999999999998, -0.06767999999999999, 0.336, 4.8, 0.0020960200000000001, -0, -0,
                    0.0012999999999999999, -0, 0.0020960200000000001, 0.0, 0.015579999999999998, 0.0, 0.82, 0.00029999999999999997, -0, -0,
           
         0.00040000000000000002, -0, 0.00029999999999999997, 0.0, 0.0, 0.0, 0.34, 0.00024216, -0, -0, 0.00024216, -0,
                    4.0000000000000003e-05, 0.0, 0.0, 0.0028799999999999997, 0.09]

    dynparm_test2 = [0.47804562306292275, 0.5871876506259908, 0.9487349009746813, 0.35387185413632094, 0.28071604959871554, 0.368556182617345,
                     0.24355010647230801, 0.675346
3418802456, 0.9728151452953864, 0.6620741264406734, 0.34669638996014096, 0.01593886435340608,
                     0.3521748260592804, 0.5384045845183812, 0.021600503502885116, 0.4654003203805651, 0.5202014161122065, 0.33744920539722967,
                     0.052363297799702835, 0.07051826001770234, 0.7389222546505236, 0.771434543548951, 0.3652539269897015, 0.2603059367721896,
                     0.4310648491411889, 0.7071252186366607, 0.16320122542325732, 0.44948421506462655, 0.48085540250421277, 0.08408482356412372,
                     0.923593157615906, 0.46852453511703684, 0.6670004526434297, 0.573634975268657, 0.12665814747855264, 0.3152822779549781,
                     0.21725524421221942, 0.5727451645276381, 0.7984025459787872, 0.7630923700903489, 0.27670044727992893, 0.919667661316697,
                     0.8828363383223908, 0.3618378476984413, 0.20690997158181246, 0.8321536435628591, 0.556153477173109, 0.3985225289975399,
                     0.7128303168401632, 0.433898343199971, 0.35116090526732047, 0.551636779637059, 0.03397585970570116, 0.28333059714010744,
                     0.4785548774382852, 0.16258401709779136, 0.8246056496848533, 0.23027686557606952, 0.26655111443159285, 0.6087446254045791]

    q_num_subs = dict(zip(puma560_def.q, q_test))

    T = puma560.geo.T[-1]
    T = T.subs(q_num_subs)
    T = numpy.matrix(T).astype(numpy.float64)

    J = puma560.kin.J[-1]
    J = J.subs(q_num_subs)
    J = numpy.matrix(J).astype(numpy.float64)

    M_func_def = sympybotics.robotcodegen.robot_code_to_func(
        'python', puma560.M_code, 'M', 'M_puma560', puma560_def)
    c_func_def = sympybotics.robotcodegen.robot_code_to_func(
        'python', puma560.c_code, 'c', 'c_puma560', puma560_def)
    C_func_def = sympybotics.robotcodegen.robot_code_to_func(
        'python', puma560.C_code, 'C', 'C_puma560', puma560_def)
    g_func_def = sympybotics.robotcodegen.robot_code_to_func(
        'python', puma560.g_code, 'g', 'g_puma560', puma560_def)
    tau_func_def = sympybotics.robotcodegen.robot_code_to_func(
        'python', puma560.invdyn_code, 'tau', 'tau_puma560', puma560_def)
    H_func_def = sympybotics.robotcodegen.robot_code_to_func(
        'python', puma560.H_code, 'H', 'H_puma560', puma560_def)

    l = locals()
    exec_(M_func_def, globals(), l)
    exec_(c_func_def, globals(), l)
    exec_(C_func_def, globals(), l)
    exec_(g_func_def, globals(), l)
    exec_(tau_func_def, globals(), l)
    exec_(H_func_def, globals(), l)
    tau_puma560 = l['tau_puma560']
    g_puma560 = l['g_puma560']
    c_puma560 = l['c_puma560']
    C_puma560 = l['C_puma560']
    M_puma560 = l['M_puma560']
    H_puma560 = l['H_puma560']

    tau = tau_puma560(dynparm_test, q_test, dq_test, ddq_test)
    tau = numpy.matrix(tau).T.astype(numpy.float64)

    g = g_puma560(dynparm_test, q_test)
    g = numpy.matrix(g).T.astype(numpy.float64)

    c = c_puma560(dynparm_test, q_test, dq_test)
    c = numpy.matrix(c).T.astype(numpy.float64)

    C = C_puma560(dynparm_test, q_test, dq_test)
    C = numpy.matrix(C).reshape(puma560.dof, puma560.dof).astype(numpy.float64)
    C_dq = C * numpy.matrix(dq_test).T

    M = M_puma560(dynparm_test, q_test)
    M = numpy.matrix(M).reshape(puma560.dof, puma560.dof).astype(numpy.float64)

    H = H_puma560(q_test, dq_test, ddq_test)
    H = numpy.matrix(H).reshape(puma560.dof, puma560.dyn.n_dynparms
                                ).astype(numpy.float64)

    tau_t2 = tau_puma560(dynparm_test2, q_test, dq_test, ddq_test)
    tau_t2 = numpy.matrix(tau_t2).T.astype(numpy.float64)

    g_t2 = g_puma560(dynparm_test2, q_test)
    g_t2 = numpy.matrix(g_t2).T.astype(numpy.float64)

    c_t2 = c_puma560(dynparm_test2, q_test, dq_test)
    c_t2 = numpy.matrix(c_t2).T.astype(numpy.float64)

    C_t2 = C_puma560(dynparm_test2, q_test, dq_test)
    C_t2 = numpy.matrix(C_t2).reshape(puma560.dof, puma560.dof
                                      ).astype(numpy.float64)
    C_dq_t2 = C_t2 * numpy.matrix(dq_test).T

    M_t2 = M_puma560(dynparm_test2, q_test)
    M_t2 = numpy.matrix(M_t2).reshape(puma560.dof, puma560.dof
    