"""SO(3) rotation hand-eye solver for ``R_A R_X = R_X R_B``."""

from __future__ import annotations

from dataclasses import dataclass

import numpy as np
from scipy.optimize import least_squares

from .contracts import MotionPair
from .geometry import orthonormalize_rotation, rotation_angle_deg, skew, so3_exp, so3_log


@dataclass(frozen=True)
class RotationHandeyeResult:
    R_IMU_lidar: np.ndarray
    residual_rms_deg: float
    residual_median_deg: float
    pair_count: int
    ok: bool
    notes: tuple[str, ...] = ()


def _pair_weight(pair: MotionPair) -> float:
    weight = float(pair.metadata.get("weight", 1.0))
    if not np.isfinite(weight) or weight <= 0:
        return 1.0
    return weight


def _tsai_rotation_initial(pairs: list[MotionPair]) -> np.ndarray:
    """Closed-form rotation hand-eye initial guess (Tsai-style linear solve)."""

    rows: list[np.ndarray] = []
    rhs: list[np.ndarray] = []
    for pair in pairs:
        alpha = so3_log(pair.R_A)
        beta = so3_log(pair.R_B)
        if np.linalg.norm(alpha) < 1e-6 or np.linalg.norm(beta) < 1e-6:
            continue
        w = np.sqrt(_pair_weight(pair))
        rows.append(w * skew(alpha + beta))
        rhs.append(w * (beta - alpha))
    if len(rows) < 2:
        return np.eye(3)
    a = np.vstack(rows)
    b = np.concatenate(rhs)
    try:
        rotvec, *_ = np.linalg.lstsq(a, b, rcond=None)
    except np.linalg.LinAlgError:
        return np.eye(3)
    return orthonormalize_rotation(so3_exp(rotvec))


def _pair_residual_deg(r_x: np.ndarray, pair: MotionPair) -> float:
    err = so3_log(r_x.T @ pair.R_A @ r_x @ pair.R_B.T)
    return float(np.degrees(np.linalg.norm(err)))


def _rms_deg(r_x: np.ndarray, pairs: list[MotionPair]) -> float:
    if not pairs:
        return 1e9
    errs = np.asarray([_pair_residual_deg(r_x, pair) for pair in pairs], dtype=float)
    return float(np.sqrt(np.mean(errs**2)))


def solve_rotation_handeye(
    pairs: list[MotionPair] | tuple[MotionPair, ...],
    *,
    R_prior: np.ndarray | None = None,
    prior_sigma_deg: float | None = None,
) -> RotationHandeyeResult:
    """Solve ``R_A R_X = R_X R_B`` with weighted robust nonlinear refinement.

    Optional CAD / installation ``R_prior`` soft-constrains the extrinsic yaw that
    is weakly observable under near-planar motion.
    """

    usable = [pair for pair in pairs if rotation_angle_deg(pair.R_A) > 1.0 and rotation_angle_deg(pair.R_B) > 1.0]
    notes: list[str] = []
    if len(usable) < 3:
        return RotationHandeyeResult(
            R_IMU_lidar=np.eye(3),
            residual_rms_deg=1e9,
            residual_median_deg=1e9,
            pair_count=len(usable),
            ok=False,
            notes=("need at least 3 motion pairs with meaningful rotation",),
        )

    r0 = _tsai_rotation_initial(usable)
    r_prior = None
    if R_prior is not None:
        r_prior = orthonormalize_rotation(np.asarray(R_prior, dtype=float).reshape(3, 3))
        rms_tsai = _rms_deg(r0, usable)
        rms_prior = _rms_deg(r_prior, usable)
        if rms_prior <= rms_tsai * 1.25:
            r0 = r_prior
            notes.append(
                f"init from rotation prior (rms={rms_prior:.3f} deg vs Tsai {rms_tsai:.3f} deg)"
            )
        else:
            notes.append(
                f"init from Tsai (rms={rms_tsai:.3f} deg; prior {rms_prior:.3f} deg kept as soft constraint)"
            )

    weights = np.asarray([_pair_weight(pair) for pair in usable], dtype=float)
    notes.append(
        f"weighted hand-eye: weight median={float(np.median(weights)):.3g}, "
        f"min={float(np.min(weights)):.3g}, max={float(np.max(weights)):.3g}"
    )

    def pack(r: np.ndarray) -> np.ndarray:
        return so3_log(r)

    def unpack(vec: np.ndarray) -> np.ndarray:
        return orthonormalize_rotation(so3_exp(vec))

    sigma = 15.0 if prior_sigma_deg is None else float(prior_sigma_deg)
    prior_w = 0.0
    if r_prior is not None and sigma > 1e-6:
        # Scale prior to a few strong pairs so it regularizes yaw without dominating.
        prior_w = float(np.sqrt(np.median(weights)) / np.deg2rad(sigma))
        notes.append(f"rotation prior soft constraint sigma={sigma:.1f} deg, weight={prior_w:.3g}")

    def residual(vec: np.ndarray) -> np.ndarray:
        r_x = unpack(vec)
        residuals = []
        for pair, weight in zip(usable, weights):
            err = so3_log(r_x.T @ pair.R_A @ r_x @ pair.R_B.T)
            residuals.append(np.sqrt(weight) * err)
        if r_prior is not None and prior_w > 0:
            residuals.append(prior_w * so3_log(r_prior.T @ r_x))
        return np.concatenate(residuals)

    opt = least_squares(residual, pack(r0), loss="huber", f_scale=np.deg2rad(1.0), max_nfev=200)
    r_x = unpack(opt.x)
    errs = np.asarray([_pair_residual_deg(r_x, pair) for pair in usable], dtype=float)
    # Report unweighted RMS/median for interpretability.
    rms = float(np.sqrt(np.mean(errs**2)))
    med = float(np.median(errs))
    notes.append(f"optimized over {len(usable)} pairs")
    ok = rms < 5.0 and len(usable) >= 3
    if not ok:
        notes.append("rotation residual RMS too high or too few pairs")
    return RotationHandeyeResult(
        R_IMU_lidar=r_x,
        residual_rms_deg=rms,
        residual_median_deg=med,
        pair_count=len(usable),
        ok=ok,
        notes=tuple(notes),
    )