"""Pour: grasp a vessel, carry it over a target, and ROTATE the wrist to tip the contents out."""
import numpy as np

from ..motion.arm import OPEN, CLOSE, grasp_quat, ease, fillet
from ..envs.scene import TABLE_TOP
from .base import approach_and_grasp


def _slerp(q0, q1, t):
    q0 = q0/(np.linalg.norm(q0)+1e-9); q1 = q1/(np.linalg.norm(q1)+1e-9)
    d = float(np.dot(q0, q1))
    if d < 0:
        q1, d = -q1, -d
    if d > 0.9995:
        q = q0+t*(q1-q0); return q/(np.linalg.norm(q)+1e-9)
    th = np.arccos(d); q2 = q1-q0*d; q2 /= (np.linalg.norm(q2)+1e-9)
    return q0*np.cos(th*t)+q2*np.sin(th*t)


def solve(env, obj, target, arm="right", pour_deg=105.0, beads=None, max_gap=0.045,
          jaw="auto", grasp_offset=(0.0, 0.0), grasp_top=None):
    a = env.arms[arm]
    rec = env.recorder
    upright = a.quat.copy()
    res, ext = approach_and_grasp(env, a, obj, max_gap=max_gap, jaw=jaw,
                                  grasp_offset=grasp_offset, grasp_top=grasp_top)
    if not res.ok:
        rec.phase = "ABORT: the vessel was never grasped"
        return {"grasped": False, "reason": res.reason}
    upright = a.quat.copy()
    tipped = grasp_quat("y", tilt_deg=pour_deg, tilt_sign=1.0)
    reg = env.scene.regions[target]
    over = a.to_root(np.array([reg["xy"][0], reg["xy"][1]+0.03, TABLE_TOP+0.20], np.float32))
    rec.phase = "4. CARRY over the target"
    a.flow(fillet([a._seg_start(), res.hold_pose+np.array([0, 0, 0.16], np.float32), over])[1:], CLOSE)
    rec.phase = "5. POUR (rotate the wrist in place)"
    hold = a._seg_start()
    for k in range(130):                     # position held, orientation slerped: a real pour
        a.quat = _slerp(upright, tipped, ease((k+1)/130.0))
        a._drive(hold, CLOSE)
    rec.phase = "6. HOLD (contents fall out)"
    for _ in range(70):
        a._drive(hold, CLOSE)
    rec.phase = "7. RETURN upright"
    for k in range(90):
        a.quat = _slerp(tipped, upright, ease((k+1)/90.0))
        a._drive(hold, CLOSE)
    out = {"grasped": True}
    if beads:
        n = sum(1 for b in beads if b in env.scene.objects
                and abs(env.scene.object_pos(b)[0]-reg["xy"][0]) < 0.12
                and abs(env.scene.object_pos(b)[1]-reg["xy"][1]) < 0.12
                and env.scene.object_pos(b)[2] < TABLE_TOP+0.10)
        out["poured"] = n
        print(f"[solver] poured {n}/{len(beads)} beads into the target area", flush=True)
    return out