import asyncio
import math
import moteus
import time
import numpy as np

pr = moteus.PositionResolution()
pr.position = moteus.INT32
pr.velocity = moteus.INT32
pr.kp_scale = moteus.INT32
pr.kd_scale = moteus.INT32

qr = moteus.QueryResolution()
qr.mode = moteus.INT8
qr.position = moteus.INT32
qr.velocity = moteus.INT32

# qr.torque = moteus.F32
qr.trajectory_complete = True


def qr_offset_angle(a1, a2, _offset1, _offset2):
    a1 = a1 * 16 - _offset1
    a2 = a2 * 16 - _offset2
    return a1, a2


def rad_to_rotation(radians):
    rot = radians / (2 * math.pi)
    return rot


def qr_inv_kin(x, z):
    y = 0
    tmp = z
    z = x
    x = tmp
    l1 = 0
    l2 = 0.194
    l3 = 0.194
    d = ((x * x) + (y * y) + (z * z) - (l1 * l1) - l2 * l2 - l3 * l3) / (2 * l2 * l3)
    th3 = math.atan2(math.sqrt(1 - d * d), d)
    if x < 0:
        th2 = math.atan2(z, math.sqrt(x * x + y * y - l1 * l1)) - math.atan2(l3 * math.sin(th3),
                                                                             l2 + l3 * math.cos(th3))
    else:
        th2 = math.atan2(z, -math.sqrt(x * x + y * y - l1 * l1)) - math.atan2(l3 * math.sin(th3),
                                                                              l2 + l3 * math.cos(th3))
    return th2, th3


offset1 = 0
offset2 = 0
i = 0


async def main():
    first = True
    sign = -1
    x = 0
    z = -0.2
    # print("res:", qr_inv_kin(0.15, -0.3))

    fdcanusb1 = moteus.Fdcanusb(path='/dev/ttyACM0')
    router = moteus.Router([(fdcanusb1, [1, 2])])

    # transport = moteus.Fdcanusb()
    c1 = moteus.Controller(id=1, position_resolution=pr, query_resolution=qr)
    c2 = moteus.Controller(id=2, position_resolution=pr, query_resolution=qr)
    data = await router.cycle([
        c1.make_stop(query=True),
        c2.make_stop(query=True)]
    )

    while True:

        if first:
            data = await router.cycle([
                c1.make_position(position=math.nan, query=True),
                c2.make_position(position=math.nan, query=True)]
            )
            print("trajectory complete: ", data[0].values[11])
            print("trajectory complete: ", data[1].values[11])
        else:

            if sign == -1:
                z = z - 0.002
            else:
                z = z + 0.002

            if z < -0.35:
                sign = 1
            elif z > -0.15:
                sign = -1

            th1, th2 = qr_inv_kin(x, z)
            print("th1: ", th1, " th2: ", th2)
            rot1 = rad_to_rotation(th1)
            rot2 = rad_to_rotation(th2)
            rot1_off, rot2_off = qr_offset_angle(rot1, rot2, offset1, offset2)
            print("rot1: ", rot1, " rot2: ", rot2)

            data = await router.cycle([
                c1.make_position(position=rot1_off, accel_limit=1, velocity_limit=1, maximum_torque=0.8, query=True),
                c2.make_position(position=rot2_off, accel_limit=1, velocity_limit=1, maximum_torque=0.8, query=True)]
            )

        if (first):
            first = False
            offset1 = data[0].values[1]
            offset2 = data[1].values[1]
            print("offset1: ", offset1)
            print("offset2: ", offset2)

        print("positionMot1: ", round(data[0].values[1], 3), " positionMot2 ", round(data[1].values[1], 3))

        await asyncio.sleep(0.02)


asyncio.run(main())