"""Push an object to a target WITHOUT grasping it: closed jaw as a finger, closed-loop servo."""
import numpy as np

from ..motion.arm import OPEN, CLOSE
from ..motion.planner import plan_path
from ..envs.scene import TABLE_TOP

# half-thickness of a closed gripper finger; the pusher contact point is offset by this much
FINGER_HALF = 0.014


def relay(env, obj, legs, steps=220, **kw):
    """Two arms push the same block in turn, each along ITS OWN leg.

    `legs` is [{"arm": "right", "target": "waypoint"}, {"arm": "left", "target": "goal"}].

    The earlier version pushed one straight line and simply handed over halfway, which does not
    work on this robot. To push a block in -y the hand must stand on its +y side, and the right
    arm's shoulder is at y=-0.2: it had to reach right over the block to get behind it, and the
    approach swiped the block off the table.

    An L-shaped route fixes it. Each arm only ever pushes AWAY from its own shoulder -- the right
    arm drives +x, the left arm drives -y -- so neither reaches past the block, and the corner
    between the legs is what makes two arms necessary rather than decorative.
    """
    rec = env.recorder
    out = {}
    for i, leg in enumerate(legs):
        side, tgt = leg["arm"], leg["target"]
        rec.phase = f"{chr(65+i)}. {side.upper()} ARM pushes to {tgt}"
        print(f"[solver] relay leg {i+1}/{len(legs)}: {side} arm -> {tgt}", flush=True)
        out[f"leg{i+1}"] = solve(env, obj, tgt, arm=side, steps=steps, **kw)
        a = env.arms[side]
        rec.phase = f"{chr(65+i)}b. {side.upper()} ARM clears out of the way"
        a.flow([a._seg_start()+np.array([0, 0, 0.18], np.float32)], CLOSE)
    last = out[f"leg{len(legs)}"]
    return {**out, "pushed": last.get("pushed", False), "err": last.get("err")}


def solve(env, obj, target, arm="right", steps=240, approach_gap=0.045, advance=0.009,
          stop_at_s=None, advance_max=0.010, push_z_frac=0.45, stop_dist=0.015):
    """Re-derive the pusher pose from the block's LIVE pose every step.

    Open-loop pushing fails the moment the block skids off the contact normal: the pusher keeps
    driving down its planned line while the block squirts away sideways. Servoing on the live
    pose keeps the finger on the block->target line the whole way.
    """
    a = env.arms[arm]
    rec = env.recorder
    reg = env.scene.regions[target]
    ext = env.scene.object_size(obj)
    # Contact height matters for anything that is not a cube: pushing a dome at 45% of its
    # height is above its centre of mass, so it tips and rolls instead of sliding.
    push_z = TABLE_TOP+float(ext[2])*float(push_z_frac)
    w0 = env.scene.object_pos(obj)

    rec.phase = "1. CLOSE THE JAW (used as a finger)"
    for _ in range(40):
        a._drive(a._seg_start(), CLOSE)

    d = np.array([reg["xy"][0]-w0[0], reg["xy"][1]-w0[1]], float)
    n = d/(np.linalg.norm(d)+1e-9)
    behind = np.array([w0[0]-n[0]*(float(ext[0])/2+FINGER_HALF+approach_gap),
                       w0[1]-n[1]*(float(ext[1])/2+FINGER_HALF+approach_gap), push_z], np.float32)
    # OVER THE TOP, then straight down. The contact point is on the far side of the block, so an
    # arm whose shoulder is on the target's side has to reach PAST the block to get there -- and
    # a planned path that cuts the corner swipes the block on the way in and launches it (it
    # ended 0.76 m off the table that way). A high traverse plus a vertical descent cannot.
    clear = a.to_root(behind+np.array([0, 0, 0.20], np.float32))
    rec.phase = "2. RISE clear of the block"
    a.flow([a._seg_start()+np.array([0, 0, 0.16], np.float32)], CLOSE)
    rec.phase = "3. TRAVERSE over the block to the far side"
    a.move_to(clear, CLOSE, tol=0.012, max_steps=150)
    rec.phase = "4. LOWER to push height"
    reach = a.move_to(a.to_root(behind), CLOSE, tol=0.008, max_steps=130)
    if reach > 0.03:
        # Say so plainly: a block that never moves looks identical to a block that would not
        # slide, and only this number tells them apart.
        print(f"[solver] WARNING: the {arm} arm could not reach the pushing stance "
              f"(off by {reach:.3f} m at {np.round(behind, 3)}) -- it is out of its workspace, "
              f"so nothing will move.", flush=True)

    rec.phase = "4. PUSH (swept finger, lateral correction from the live block pose)"
    # The finger SWEEPS: its commanded point marches from `behind` to just short of the target at
    # a fixed rate, so it keeps displacing the block. Re-deriving the command from the block's
    # live pose each step instead (the obvious closed-loop form) makes the finger trail the block
    # by a constant stand-off and it creeps -- 2.5 cm in 240 steps, because a commanded point
    # that is always just-touching never actually advances into anything.
    back = float(ext[0])/2+FINGER_HALF
    p0 = behind[:2].astype(float)
    goal = np.array([reg["xy"][0], reg["xy"][1]], float)-n*back
    # The sweep is parameterised by CONTROL STEPS, not by distance. Sizing the loop as
    # span/advance conflated the two: it gave 25 iterations = 25 sim steps for a 22 cm push, so
    # the loop exited before the arm had traversed anything and the block "crept" 3 cm.
    ramp = max(1, int(steps*0.85))
    for k in range(steps):
        w = env.scene.object_pos(obj)
        to_t = np.array([reg["xy"][0]-w[0], reg["xy"][1]-w[1]], float)
        dist = float(np.linalg.norm(to_t))
        # A rolling object needs the push to END EARLY and let it coast in: keep driving
        # until the centre is on the target and a dome has already rolled past it.
        if dist < float(stop_dist):
            break
        # relay handoff: stop once the block has crossed into the other arm's half
        if stop_at_s is not None and float(np.dot(w[:2]-w0[:2], n)) >= stop_at_s:
            print(f"[solver] handoff: block passed s={stop_at_s:.3f} m after {k} steps", flush=True)
            break
        # RE-AIM every step: the contact point is `back` behind the block along the LIVE
        # block->target line, driven `bite` past it. Sweeping a line fixed at t=0 instead lets a
        # block that skids sideways stay skidded -- it ended 13 cm off in x that way -- because a
        # fixed line has no authority to steer the block back.
        nn = to_t/(dist+1e-9)
        # penetration is capped: 17 mm into a 25 mm half-block launched it 1.8 m off the table
        bite = min(advance, dist, advance_max)
        cp = w[:2]-nn*(back-bite)
        # blend toward the swept ideal so the finger keeps net forward progress even while the
        # block is momentarily stuck against static friction
        swept = p0+(goal-p0)*min(1.0, (k+1)/ramp)
        cp = cp+n*max(0.0, float(np.dot(swept-cp, n)))*0.25
        a._drive(a.to_root(np.array([cp[0], cp[1], push_z], np.float32)), CLOSE)
        if k % 60 == 0:
            print(f"[solver]   push k={k:3d} block=({w[0]:+.3f},{w[1]:+.3f}) dist={dist:.3f}",
                  flush=True)
    wf = env.scene.object_pos(obj)
    err = float(np.hypot(wf[0]-reg["xy"][0], wf[1]-reg["xy"][1]))
    rec.phase = "5. RETREAT"
    a.flow([a._seg_start()+np.array([0, 0, 0.14], np.float32)], CLOSE)
    print(f"[solver] push done: block=({wf[0]:.3f},{wf[1]:.3f}) target="
          f"({reg['xy'][0]:.3f},{reg['xy'][1]:.3f}) err={err:.3f}", flush=True)
    return {"pushed": err < 0.07, "err": err}