"""Local ArUco pose estimation with Nuwa RGB intrinsics.

The marker-finder service bakes D435i intrinsics into its solvePnP, which
produces systematic pose errors on Nuwa images. We recompute the pose from the
2D corners (still trustworthy) using the real Nuwa intrinsics, disambiguating
the planar-marker ambiguity via IPPE_SQUARE + smallest reprojection error.
"""

from __future__ import annotations

from typing import Any

import cv2
import numpy as np


def nuwa_intrinsics_to_matrix(intr: dict[str, Any]
                              ) -> tuple[np.ndarray, np.ndarray]:
    """Accept the `data` payload from Nuwa /api/intrinsics and return
    (K_3x3, dist_5x1) for RGB. Distortion fields default to 0 if absent."""
    rgb = intr.get("rgb") if isinstance(intr, dict) else None
    if not isinstance(rgb, dict):
        raise ValueError(f"Nuwa intrinsics missing 'rgb' block: {intr!r}")
    K = np.array([
        [float(rgb["fx"]), 0.0,              float(rgb["cx"])],
        [0.0,              float(rgb["fy"]), float(rgb["cy"])],
        [0.0,              0.0,              1.0],
    ], dtype=np.float64)
    d = rgb.get("distortion") or {}
    dist = np.array([
        float(d.get("k1", 0.0)),
        float(d.get("k2", 0.0)),
        float(d.get("p1", 0.0)),
        float(d.get("p2", 0.0)),
        float(d.get("k3", 0.0)),
    ], dtype=np.float64)
    return K, dist


def nuwa_intrinsics_to_opencv_payload(intr: dict[str, Any]) -> dict[str, list]:
    """Build the wire-format dict that marker-finder's /detect (and /intrinsics
    POST) expects: `{camera_matrix: 3x3, dist_coeffs: [k1,k2,p1,p2,k3]}`.
    Plain Python lists so it serialises through JSON."""
    K, dist = nuwa_intrinsics_to_matrix(intr)
    return {"camera_matrix": K.tolist(), "dist_coeffs": dist.tolist()}


def solve_square_marker(corners_px: np.ndarray, size_m: float,
                        K: np.ndarray, dist: np.ndarray
                        ) -> dict[str, Any]:
    """Recover (R, t) for a square planar marker of side `size_m` from its
    4 image corners in the order TL, TR, BR, BL (ArUco convention).

    Uses solvePnPGeneric with SOLVEPNP_IPPE_SQUARE, which returns both planar
    solutions; we pick the one with smaller reprojection error. The runner-up
    error is also returned so the caller can detect ambiguous detections
    (both errors similar → the marker is nearly fronto-parallel and the pose
    is fundamentally underdetermined by a single view).
    """
    s = size_m / 2.0
    obj = np.array([
        [-s,  s, 0.0],
        [ s,  s, 0.0],
        [ s, -s, 0.0],
        [-s, -s, 0.0],
    ], dtype=np.float64)
    img = np.asarray(corners_px, dtype=np.float64).reshape(-1, 2)
    if img.shape[0] != 4:
        raise ValueError(f"expected 4 corners, got shape {img.shape}")

    retval, rvecs, tvecs, errs = cv2.solvePnPGeneric(
        obj, img, K, dist, flags=cv2.SOLVEPNP_IPPE_SQUARE
    )
    if not retval or rvecs is None or len(rvecs) == 0:
        raise RuntimeError("solvePnP (IPPE_SQUARE) failed to converge")

    errs_arr = np.asarray(errs).flatten() if errs is not None else np.array([0.0])
    best = int(np.argmin(errs_arr)) if errs_arr.size > 0 else 0
    R, _ = cv2.Rodrigues(rvecs[best])
    t = np.asarray(tvecs[best]).flatten()

    # runner-up error for ambiguity reporting
    if errs_arr.size > 1:
        runner = float(np.partition(errs_arr, 1)[1])
    else:
        runner = float("inf")
    best_err = float(errs_arr[best])
    # ratio close to 1 means the two solutions are almost equally plausible.
    ambiguity = float(best_err / runner) if runner > 0 else 0.0
    return {
        "R": R,
        "t": t,
        "reprojection_error_px": best_err,
        "runnerup_error_px": runner,
        "ambiguity_ratio": ambiguity,
        "n_solutions": int(len(rvecs)),
    }