"""Planar 2R pick/place state machine with idealized joint dynamics.

No contact/friction/gripper physics. Standard library only.
"""
import json
import math

LENGTHS, DT = (1.0, 0.8), 0.005
OBSTACLE = (0.2, -0.4, 0.12)
LINK_RADIUS = 0.03


def fk(q):
    a, b = LENGTHS
    elbow = a * math.cos(q[0]), a * math.sin(q[0])
    tip = elbow[0] + b * math.cos(sum(q)), elbow[1] + b * math.sin(sum(q))
    return elbow, tip


def ik(target):
    x, y = target
    a, b = LENGTHS
    c = (x*x + y*y - a*a - b*b)/(2*a*b)
    if not -1 <= c <= 1:
        raise ValueError("unreachable target")
    q2 = -math.acos(c)
    return [math.atan2(y, x) - math.atan2(b*math.sin(q2), a+b*math.cos(q2)), q2]


def segment_distance(a, b, point):
    dx, dy = b[0] - a[0], b[1] - a[1]
    t = ((point[0]-a[0])*dx + (point[1]-a[1])*dy)/(dx*dx + dy*dy)
    t = max(0., min(1., t))
    return math.hypot(a[0]+t*dx-point[0], a[1]+t*dy-point[1])


def clearance(q, obstacle=OBSTACLE):
    elbow, tip = fk(q)
    point = obstacle[:2]
    return min(segment_distance((0., 0.), elbow, point),
               segment_distance(elbow, tip, point)) - obstacle[2] - LINK_RADIUS


def quintic(u):
    return 10*u**3 - 15*u**4 + 6*u**5, 30*u*u - 60*u**3 + 30*u**4


def main():
    phases = [("approach", (1.3, 0.3)), ("grasp", (1.1, 0.4)),
              ("lift", (1.0, 0.7)), ("place", (0.4, 1.2))]
    q, velocity = ik(phases[0][1]), [0., 0.]
    minimum, max_error, attached = clearance(q), 0., False
    events = []
    for name, target in phases:
        begin, end, duration = q[:], ik(target), 2.0
        for step in range(1, int(duration/DT)+1):
            u = step*DT/duration
            scale, derivative = quintic(u)
            desired = [begin[j]+(end[j]-begin[j])*scale for j in range(2)]
            desired_velocity = [(end[j]-begin[j])*derivative/duration for j in range(2)]
            for j in range(2):
                command = 60*(desired[j]-q[j]) + 14*(desired_velocity[j]-velocity[j])
                command = max(-4., min(4., command))
                # Unit diagonal joint inertia, viscous friction, no gravity/coupling.
                acceleration = command - 0.1*velocity[j]
                q[j] += velocity[j]*DT + 0.5*acceleration*DT*DT
                velocity[j] += acceleration*DT
            minimum = min(minimum, clearance(q))
            if minimum < 0:
                raise AssertionError("collision in numerical model")
        _, tip = fk(q)
        error = math.dist(tip, target)
        max_error = max(max_error, error)
        assert error < 0.02
        if name == "grasp":
            attached = True
        if name == "place":
            assert attached
            attached = False
        events.append({"phase": name, "tip_error_m": error, "object_attached": attached})
    blocked_target_detected = clearance(ik((1.1, 0.4)), (1.1, 0.4, 0.1)) < 0
    assert blocked_target_detected
    print(json.dumps({"simulation": "planar-2R-unit-inertia", "dt_s": DT,
                      "minimum_clearance_m": minimum, "max_keyframe_error_m": max_error,
                      "blocked_target_detected": blocked_target_detected, "events": events}, indent=2))


if __name__ == "__main__":
    main()