"""Two arms pinch a rope at both ends and pull it taut.

The rope is a CHAIN: many short rigid links joined by universal joints. Isaac Lab does have FEM
deformables, but a rope you have to pinch and pull is far better served by a chain -- a soft body
slips out of a parallel jaw, and the thing being measured here (end-to-end span versus the rope's
own length) is exactly what a chain represents faithfully.

Nothing here goes through approach_and_grasp: that works on rigid bodies in `scene.objects`, and
a rope link lives in `scene.articulations`. The end links are located with link_pos() instead.
"""
import numpy as np

from ..motion.arm import OPEN, CLOSE, grasp_quat
from ..envs.scene import TABLE_TOP


def solve(env, rope, n_links, pull_to=0.30, lift=0.10, grip_inset=1):
    """Pinch link `grip_inset` from each end, lift, then draw the hands apart to `pull_to`."""
    rec = env.recorder
    L, R = env.arms["left"], env.arms["right"]
    lo, hi = grip_inset, n_links-1-grip_inset

    a_xyz = env.scene.link_pos(rope, lo)
    b_xyz = env.scene.link_pos(rope, hi)
    # whichever end is on each arm's side, so the arms never cross
    if a_xyz[1] < b_xyz[1]:
        a_xyz, b_xyz = b_xyz, a_xyz
    print(f"[solver] rope ends at {a_xyz.round(3)} (left) and {b_xyz.round(3)} (right), "
          f"apart {np.linalg.norm(a_xyz-b_xyz):.3f} m", flush=True)

    for arm in (L, R):
        arm.quat = grasp_quat("y")
    above = 0.11
    Lg = np.array([a_xyz[0], a_xyz[1], TABLE_TOP+0.012], np.float32)
    Rg = np.array([b_xyz[0], b_xyz[1], TABLE_TOP+0.012], np.float32)

    rec.phase = "1. BOTH HANDS reach over their end of the rope"
    for k in range(120):
        t = (k+1)/120.0
        env.drive_both(L.to_root(Lg+np.array([0, 0, above*(1-t)+0.10*t], np.float32)),
                       R.to_root(Rg+np.array([0, 0, above*(1-t)+0.10*t], np.float32)), OPEN, OPEN)
    rec.phase = "2. LOWER onto the rope ends"
    for k in range(140):
        t = (k+1)/140.0
        env.drive_both(L.to_root(Lg+np.array([0, 0, 0.10*(1-t)], np.float32)),
                       R.to_root(Rg+np.array([0, 0, 0.10*(1-t)], np.float32)), OPEN, OPEN)
    rec.phase = "3. PINCH both ends"
    for _ in range(70):
        env.drive_both(L.to_root(Lg), R.to_root(Rg), CLOSE, CLOSE)

    rec.phase = "4. LIFT the rope clear of the table"
    for k in range(110):
        t = (k+1)/110.0
        env.drive_both(L.to_root(Lg+np.array([0, 0, lift*t], np.float32)),
                       R.to_root(Rg+np.array([0, 0, lift*t], np.float32)), CLOSE, CLOSE)

    rec.phase = f"5. PULL APART until the rope is straight ({pull_to*100:.0f} cm)"
    mid_y = (Lg[1]+Rg[1])/2.0
    for k in range(200):
        t = (k+1)/200.0
        ly = mid_y+(pull_to/2.0)*t+(Lg[1]-mid_y)*(1-t)
        ry = mid_y-(pull_to/2.0)*t+(Rg[1]-mid_y)*(1-t)
        env.drive_both(L.to_root(np.array([Lg[0], ly, TABLE_TOP+lift], np.float32)),
                       R.to_root(np.array([Rg[0], ry, TABLE_TOP+lift], np.float32)), CLOSE, CLOSE)
    rec.phase = "6. HOLD it taut"
    for _ in range(50):
        env.drive_both(L._seg_start(), R._seg_start(), CLOSE, CLOSE)

    a2 = env.scene.link_pos(rope, 0)
    b2 = env.scene.link_pos(rope, n_links-1)
    span = float(np.linalg.norm(a2-b2))
    print(f"[solver] rope end-to-end span {span:.3f} m", flush=True)
    return {"span": span}