"""Sort: each object goes to ITS OWN region, chosen per object rather than one shared bin."""
import numpy as np

from ..motion.arm import OPEN
from .base import approach_and_grasp, place_at


COMFORT_R = 0.25          # reach at which this arm tracks well; measured from the passing tasks


def _best_arm(env, *points):
    """Pick the arm that is COMFORTABLE for every point, not the one that is nearest.

    Nearest is actively wrong: the left pad sat 9 cm from the right arm's own base, so "nearest"
    chose the right arm and it folded into a configuration it could not servo out of (0.54 m
    tracking error). Scoring on |reach - COMFORT_R| picks the arm that has to neither fold up nor
    stretch out.
    """
    best, bscore = None, 1e9
    for side, arm in env.arms.items():
        root = np.asarray(arm.root, float)[:2]
        score = max(abs(float(np.linalg.norm(root-np.asarray(p, float)[:2]))-COMFORT_R)
                    for p in points)
        if score < bscore:
            best, bscore = side, score
    return best or "right"


def solve(env, assignment, arm="auto", jaw="auto", lift=0.16, on_top=False, max_gap=0.045):
    """`assignment` maps object name -> region name.

    Unlike multi_pick there is no drop spread: each region holds one object, so the drop point is
    the region centre and spreading would only push it toward the rim.

    `arm="auto"` picks the near arm per destination, which is what makes a two-sided sort work
    at all: one arm cannot serve pads on both sides of the table.
    """
    out = {}
    items = list(assignment.items())
    prev = None
    for i, (name, region) in enumerate(items):
        pre = f"[{i+1}/{len(items)}] "
        side = arm if arm != "auto" else _best_arm(
            env, env.scene.object_pos(name), env.scene.regions[region]["xy"])
        a = env.arms[side]
        print(f"[solver] {pre}{name} -> {region} with the {side} arm", flush=True)
        if prev is not None:
            # Regroup the arm that JUST finished: it is left stretched out over its pad, and
            # planning the next reach from there leaves it stuck. If the next item uses the other
            # arm this also clears the previous one out of the workspace.
            env.recorder.phase = f"{pre}0. REGROUP"
            prev.flow([prev._seg_start()+np.array([0, 0, 0.18], np.float32)], OPEN)
        prev = a
        res, ext = approach_and_grasp(env, a, name, jaw=jaw, max_gap=max_gap, prefix=pre)
        if not res.ok:
            out[name] = {"grasped": False, "reason": res.reason}
            continue
        out[name] = {"grasped": True, "region": region,
                     "place_err": place_at(env, a, region, float(ext[2])/2.0,
                                           on_top=on_top, lift=lift, prefix=pre)}
    return out