"""Esplora l'envelope di visibilità del tag ArUco sul polso H1 (sx) usando
la posizione reale (actual), non il target. Poi calibra con sample random
nell'envelope, scartando duplicati e arm stuck. Salva workspace JSON, PNG
e calibration.

Uso: python3 scripts/explore_and_calibrate.py

Note H1: 4 DoF, niente polso articolato → la diversità di orientamento per la
calibrazione viene dall'IK sui (X,Y,Z) variati, non da delta_quat. Si comanda
sempre quat=identità e si lascia l'IK scegliere la branca.
"""

from __future__ import annotations

import json
import math
import random
import time
from pathlib import Path
from typing import Any

import requests

ORC = "http://localhost:8890"
H1 = "http://localhost:8767"
SIDE = "left"
QUAT_IDENT = [0.0, 0.0, 0.0, 1.0]

# Geometria di sondaggio
PROBE_STEP_M = 0.030
PROBE_MAX_M = 0.250
REFINE_HALVINGS = 2
Z_RANGE_HALF_M = 0.050  # ±50 mm da home (richiesta utente)
Z_PROBE_OFFSETS_M = [0.0, 0.025, -0.025, 0.050, -0.050]
ENVELOPE_MARGIN_M = 0.020  # restringi i bounds di 20 mm dai punti reachable

# Movimento — alta precisione: precise closed-loop, velocità bassa, waypoint fitti
PRECISE = True
SPEED_MMS = 40
TRAJ_STEP_MM = 10
SETTLE_S = 1.2
MOVE_TIMEOUT_S = 25.0
STUCK_DELTA_M = 0.030  # se actual cambia meno di 30 mm e il target era lontano → stuck
STUCK_TARGET_DIST_M = 0.100

# Calibrazione
N_SAMPLES = 24
RANDOM_SEED = 17
MIN_SAMPLE_SEP_M = 0.030  # distanza minima tra actual di sample accettati
REUSE_WORKSPACE_IF_FRESH_S = 3600  # se workspace_left.json è < 1h, riusalo (skip probes)

DATA_DIR = Path("/opt/docker/robotserver/presa-mano/data")


def post(url: str, body: dict[str, Any] | None = None, timeout: float = 15.0) -> dict[str, Any]:
    r = requests.post(url, json=body or {}, timeout=timeout)
    try:
        return r.json()
    except Exception:
        return {"status_code": r.status_code, "text": r.text[:500]}


def get(url: str, timeout: float = 5.0) -> dict[str, Any]:
    r = requests.get(url, timeout=timeout)
    return r.json()


def go_home() -> dict[str, Any]:
    return post(f"{H1}/arms/go_neutral", {})


def read_left_pos() -> list[float]:
    return get(f"{ORC}/api/wrist")["left_wrist"]["position"]


def read_left_quat() -> list[float]:
    return get(f"{ORC}/api/wrist")["left_wrist"]["quaternion"]


def move_abs(pos: list[float], quat: list[float] = QUAT_IDENT) -> dict[str, Any]:
    """Bypassa l'orchestratore (la sua AbsoluteReq scarta speed_mm_s/traj_step_mm).
    Parla direttamente al bridge H1 per ottenere precise + slow + waypoint fitti."""
    body = {
        "side": SIDE,
        "position": list(pos),
        "quaternion": list(quat),
        "wait": True,
        "timeout_s": MOVE_TIMEOUT_S,
        "precise": PRECISE,
        "speed_mm_s": SPEED_MMS,
        "traj_step_mm": TRAJ_STEP_MM,
    }
    return post(f"{H1}/wrist/absolute", body, timeout=MOVE_TIMEOUT_S + 5)


def detect_marker() -> dict[str, Any]:
    return post(f"{ORC}/api/calibration/preview", {}, timeout=10)


def take_calib_sample() -> dict[str, Any]:
    return post(f"{ORC}/api/calibration/sample", {"side": SIDE}, timeout=15)


def solve_calib() -> dict[str, Any]:
    return post(f"{ORC}/api/calibration/solve", {"side": SIDE}, timeout=30)


def clear_samples() -> dict[str, Any]:
    r = requests.delete(f"{ORC}/api/calibration/samples", params={"side": SIDE}, timeout=10)
    return r.json() if r.headers.get("content-type", "").startswith("application/json") else {"text": r.text}


def dist3(a: list[float], b: list[float]) -> float:
    return math.sqrt(sum((x - y) ** 2 for x, y in zip(a, b)))


def move_and_observe(target: list[float], prev_actual: list[float] | None = None) -> dict[str, Any]:
    """Esegue il move_abs, restituisce dict con fields:
       target, actual, found, marker_id, center_px, stuck (bool), move_resp.
    """
    move_resp = move_abs(target)
    time.sleep(SETTLE_S)
    actual = read_left_pos()
    det = detect_marker()
    found = bool(det.get("found"))
    stuck = False
    if prev_actual is not None:
        target_dist = dist3(target, prev_actual)
        actual_change = dist3(actual, prev_actual)
        if target_dist > STUCK_TARGET_DIST_M and actual_change < STUCK_DELTA_M:
            stuck = True
    return {
        "target": target,
        "actual": actual,
        "found": found,
        "center_px": det.get("center_px"),
        "marker_id": det.get("marker_id"),
        "stuck": stuck,
        "move_stalled": bool(move_resp.get("stalled")),
        "move_reached": bool(move_resp.get("reached_target")),
    }


def probe_axis(home: list[float], direction: str) -> list[dict]:
    """Cammina dall'home finché il tag scompare oppure il move fallisce o l'arm si blocca.
    Poi raffina binario tra l'ultimo visibile e il primo invisibile."""
    axis = 0 if direction[0] == "x" else 1
    sign = 1 if direction[1] == "+" else -1
    points: list[dict] = []
    last_visible = list(home)
    first_invisible: list[float] | None = None
    prev_actual = list(home)

    for k in range(1, int(PROBE_MAX_M / PROBE_STEP_M) + 1):
        target = list(home)
        target[axis] = home[axis] + sign * k * PROBE_STEP_M
        print(f"  step {k} {direction}: target={[round(v,3) for v in target]}", flush=True)
        rec = move_and_observe(target, prev_actual=prev_actual)
        rec.update({"phase": "axial", "direction": direction})
        points.append(rec)
        print(
            f"    actual={[round(v,3) for v in rec['actual']]} "
            f"found={rec['found']} stuck={rec['stuck']} stalled={rec['move_stalled']}",
            flush=True,
        )
        prev_actual = rec["actual"]
        if rec["stuck"]:
            print(f"    → arm stuck su {direction}, fermo qui", flush=True)
            break
        if rec["found"]:
            last_visible = list(rec["actual"])  # uso ACTUAL come edge
        else:
            first_invisible = list(rec["actual"])
            break

    if first_invisible is not None and not rec.get("stuck", False):
        for i in range(REFINE_HALVINGS):
            mid = [(a + b) / 2.0 for a, b in zip(last_visible, first_invisible)]
            print(f"  refine #{i+1} {direction}: target={[round(v,3) for v in mid]}", flush=True)
            rec = move_and_observe(mid, prev_actual=prev_actual)
            rec.update({"phase": "refine", "direction": direction})
            points.append(rec)
            print(
                f"    actual={[round(v,3) for v in rec['actual']]} "
                f"found={rec['found']} stuck={rec['stuck']}",
                flush=True,
            )
            prev_actual = rec["actual"]
            if rec["stuck"]:
                break
            if rec["found"]:
                last_visible = list(rec["actual"])
            else:
                first_invisible = list(rec["actual"])

    return points


def probe_z(home: list[float]) -> list[dict]:
    """Probe Z a home XY per misurare il range Z effettivamente reachable+visibile."""
    points: list[dict] = []
    prev_actual = list(home)
    for dz in Z_PROBE_OFFSETS_M:
        target = [home[0], home[1], home[2] + dz]
        print(f"  z probe dz={dz*1000:+.0f}mm: target={[round(v,3) for v in target]}", flush=True)
        rec = move_and_observe(target, prev_actual=prev_actual)
        rec.update({"phase": "z_axial", "direction": f"z{'+' if dz>=0 else '-'}", "dz_m": dz})
        points.append(rec)
        print(
            f"    actual={[round(v,3) for v in rec['actual']]} found={rec['found']}",
            flush=True,
        )
        prev_actual = rec["actual"]
    return points


def visualize(workspace: dict[str, Any]) -> Path:
    import matplotlib
    matplotlib.use("Agg")
    import matplotlib.pyplot as plt
    from matplotlib.patches import Rectangle

    fig, ax = plt.subplots(figsize=(9, 9))
    for p in workspace["probes"]:
        x, y = p["actual"][0], p["actual"][1]
        c = "tab:green" if p["found"] else "tab:red"
        marker = "o" if p["phase"] == "axial" else ("s" if p["phase"] == "refine" else "D")
        ax.scatter(x, y, c=c, s=80, alpha=0.7, marker=marker, edgecolors="black", linewidths=0.4)
    for s in workspace.get("calib_samples_xy", []):
        ax.scatter(s[0], s[1], c="tab:blue", s=60, marker="x", linewidths=2.0)
    xr = workspace["x_range_m"]
    yr = workspace["y_range_m"]
    rect = Rectangle((xr[0], yr[0]), xr[1] - xr[0], yr[1] - yr[0],
                     fill=False, edgecolor="tab:blue", linewidth=2, linestyle="--", label="envelope")
    ax.add_patch(rect)
    raw_xr = workspace["x_range_raw_m"]
    raw_yr = workspace["y_range_raw_m"]
    rect_raw = Rectangle((raw_xr[0], raw_yr[0]), raw_xr[1] - raw_xr[0], raw_yr[1] - raw_yr[0],
                         fill=False, edgecolor="gray", linewidth=1, linestyle=":", label="reachable raw")
    ax.add_patch(rect_raw)
    h = workspace["home"]
    ax.scatter(h[0], h[1], c="black", s=250, marker="*", zorder=5, label="home")
    ax.set_xlabel("X (m, base frame)")
    ax.set_ylabel("Y (m, base frame)")
    title = (
        f"Tag visibility envelope — left wrist (H1)\n"
        f"X: [{xr[0]:.3f}, {xr[1]:.3f}] m   Y: [{yr[0]:.3f}, {yr[1]:.3f}] m   "
        f"Z: [{workspace['z_range_m'][0]:.3f}, {workspace['z_range_m'][1]:.3f}] m\n"
        f"verde=visibile, rosso=non visibile, blu x=calib samples, *=home"
    )
    ax.set_title(title)
    ax.set_aspect("equal")
    ax.grid(True, alpha=0.3)
    ax.legend(loc="best")
    out_path = DATA_DIR / "workspace_left.png"
    plt.savefig(out_path, dpi=120, bbox_inches="tight")
    plt.close()
    return out_path


def shrink(lo: float, hi: float, margin: float) -> tuple[float, float]:
    if hi - lo <= 2 * margin:
        # envelope troppo stretto → niente shrink
        return lo, hi
    return lo + margin, hi - margin


def main() -> int:
    print("=== presa-mano: explore + calibrate (left, H1) — v2 (actual-based) ===", flush=True)

    print("\n[0] go_neutral", flush=True)
    print(go_home(), flush=True)
    time.sleep(2.0)
    home = read_left_pos()
    home_quat = read_left_quat()
    print(f"home pos = {[round(v,4) for v in home]}", flush=True)
    print(f"home quat = {[round(v,4) for v in home_quat]}", flush=True)

    ws_path = DATA_DIR / "workspace_left.json"
    workspace: dict[str, Any] | None = None
    if ws_path.exists():
        age = time.time() - ws_path.stat().st_mtime
        if age < REUSE_WORKSPACE_IF_FRESH_S:
            try:
                workspace = json.loads(ws_path.read_text())
                print(f"\n[1*] riuso workspace già mappato (age={age:.0f}s) — skip probes", flush=True)
                print(f"     X={workspace['x_range_m']} Y={workspace['y_range_m']} Z={workspace['z_range_m']}", flush=True)
            except Exception:
                workspace = None

    if workspace is None:
        print("\n[1] axial probes per envelope visibilità tag (X,Y)", flush=True)
        all_probes: list[dict] = []
        for direction in ["x+", "x-", "y+", "y-"]:
            print(f"\n--- {direction} ---", flush=True)
            go_home()
            time.sleep(1.5)
            all_probes.extend(probe_axis(home, direction))

        print("\n[1.5] Z probe a home XY", flush=True)
        go_home()
        time.sleep(1.5)
        all_probes.extend(probe_z(home))

        visible_actuals = [p["actual"] for p in all_probes if p["found"]]
        if not visible_actuals:
            print("Nessun probe visibile — abort.", flush=True)
            return 1
        xs = [a[0] for a in visible_actuals]
        ys = [a[1] for a in visible_actuals]
        zs = [a[2] for a in visible_actuals]
        raw_x = (min(xs), max(xs))
        raw_y = (min(ys), max(ys))
        raw_z = (min(zs), max(zs))

        x_lo, x_hi = shrink(*raw_x, ENVELOPE_MARGIN_M)
        y_lo, y_hi = shrink(*raw_y, ENVELOPE_MARGIN_M)
        z_lo_user = home[2] - Z_RANGE_HALF_M
        z_hi_user = home[2] + Z_RANGE_HALF_M
        z_lo = max(z_lo_user, raw_z[0] + 0.005)
        z_hi = min(z_hi_user, raw_z[1] - 0.005)
        if z_lo >= z_hi:
            z_lo = z_hi = home[2]

        workspace = {
            "side": SIDE,
            "robot": "h1",
            "home": home,
            "home_quat_xyzw": home_quat,
            "x_range_raw_m": list(raw_x),
            "y_range_raw_m": list(raw_y),
            "z_range_raw_m": list(raw_z),
            "x_range_m": [x_lo, x_hi],
            "y_range_m": [y_lo, y_hi],
            "z_range_m": [z_lo, z_hi],
            "probes": all_probes,
            "probe_step_m": PROBE_STEP_M,
            "envelope_margin_m": ENVELOPE_MARGIN_M,
            "refine_halvings": REFINE_HALVINGS,
            "timestamp": time.strftime("%Y-%m-%dT%H:%M:%S"),
        }
        ws_path.write_text(json.dumps(workspace, indent=2))
        print(f"\n[2] workspace salvato → {ws_path}", flush=True)
        print(f"    raw reachable+visibile (actual): X={raw_x} Y={raw_y} Z={raw_z}", flush=True)
        print(f"    envelope (con margin {ENVELOPE_MARGIN_M*1000:.0f}mm): X=[{x_lo:.3f},{x_hi:.3f}] Y=[{y_lo:.3f},{y_hi:.3f}] Z=[{z_lo:.3f},{z_hi:.3f}]", flush=True)
    else:
        x_lo, x_hi = workspace["x_range_m"]
        y_lo, y_hi = workspace["y_range_m"]
        z_lo, z_hi = workspace["z_range_m"]

    print(f"\n[3] random calibration: target {N_SAMPLES} sample, dedup ≥{MIN_SAMPLE_SEP_M*1000:.0f}mm", flush=True)
    go_home()
    time.sleep(2.0)
    print(f"clear samples: {clear_samples()}", flush=True)

    rng = random.Random(RANDOM_SEED)
    accepted_actuals: list[list[float]] = []
    accepted_records: list[dict] = []
    attempts = 0
    max_attempts = N_SAMPLES * 4
    prev_actual = read_left_pos()

    while len(accepted_actuals) < N_SAMPLES and attempts < max_attempts:
        attempts += 1
        x = rng.uniform(x_lo, x_hi)
        y = rng.uniform(y_lo, y_hi)
        z = rng.uniform(z_lo, z_hi) if z_hi > z_lo else home[2]
        target = [x, y, z]
        print(f"\n[sample {len(accepted_actuals)+1}/{N_SAMPLES}, attempt {attempts}] target={[round(v,3) for v in target]}", flush=True)
        rec = move_and_observe(target, prev_actual=prev_actual)
        actual = rec["actual"]
        if rec["stuck"]:
            print("  → stuck, skip (no sample API call)", flush=True)
            continue
        if not rec["found"]:
            print("  → marker non visibile, skip", flush=True)
            prev_actual = actual
            continue
        # dedup
        too_close = [a for a in accepted_actuals if dist3(a, actual) < MIN_SAMPLE_SEP_M]
        if too_close:
            min_d = min(dist3(a, actual) for a in too_close) * 1000
            print(f"  → duplicato (dist {min_d:.1f}mm da sample esistente), skip", flush=True)
            prev_actual = actual
            continue
        sres = take_calib_sample()
        if not sres.get("ok", True):
            print(f"  → sample API fallito: {sres}", flush=True)
            prev_actual = actual
            continue
        accepted_actuals.append(actual)
        accepted_records.append({"target": target, "actual": actual, "result": sres})
        prev_actual = actual
        print(f"  ✓ accept actual={[round(v,3) for v in actual]} count={sres.get('count','?')}", flush=True)

    print(f"\n→ {len(accepted_actuals)}/{N_SAMPLES} samples raccolti in {attempts} tentativi", flush=True)
    workspace["calib_samples_xy"] = [[a[0], a[1]] for a in accepted_actuals]
    workspace["calib_attempts"] = attempts
    workspace["calib_accepted"] = len(accepted_actuals)
    ws_path.write_text(json.dumps(workspace, indent=2))

    print("\n[4] solve calibration", flush=True)
    sol = solve_calib()
    summary = {k: v for k, v in sol.items() if k not in ("samples",)}
    print(json.dumps(summary, indent=2)[:2500], flush=True)

    print("\n[5] visualize", flush=True)
    png_path = visualize(workspace)
    print(f"PNG → {png_path}", flush=True)

    print("\n[6] go home final", flush=True)
    go_home()

    print("\n=== DONE ===", flush=True)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())