"""Build IMU/LiDAR relative-motion pairs for hand-eye calibration."""

from __future__ import annotations

from dataclasses import dataclass

import numpy as np

from .contracts import ImuSeries, LidarFrame, MotionPair
from .geometry import make_transform, rotation_angle_deg
from .imu_preintegration import preintegrate_imu
from .registration import register_lidar_pair
from .time_offset import lidar_time_to_imu_time


@dataclass(frozen=True)
class MotionPairSet:
    pairs: tuple[MotionPair, ...]
    notes: tuple[str, ...] = ()


def build_motion_pairs(
    *,
    session_id: str,
    keyframes: list[LidarFrame],
    keyframe_indices: list[int] | tuple[int, ...],
    imu: ImuSeries,
    delta_t_s: float,
    gyro_bias_rad_s: np.ndarray | None = None,
    acc_bias_m_s2: np.ndarray | None = None,
    min_rotation_deg: float = 3.0,
    min_translation_m: float = 0.3,
    max_index_span: int = 4,
) -> MotionPairSet:
    """Create A/B motion pairs between nearby keyframes.

    IMU side uses full Phase-C preintegration (``ΔR/Δv/Δp``, ``Σ9``, ``J_bg/J_ba``).
    Rotation hand-eye still consumes ``R_A = ΔR`` only.
    """

    notes: list[str] = []
    pairs: list[MotionPair] = []
    bias_g = np.zeros(3) if gyro_bias_rad_s is None else np.asarray(gyro_bias_rad_s, dtype=float)
    bias_a = np.zeros(3) if acc_bias_m_s2 is None else np.asarray(acc_bias_m_s2, dtype=float)
    n = len(keyframes)
    if n < 2:
        return MotionPairSet((), ("need at least two keyframes",))

    for span in range(1, max_index_span + 1):
        for start in range(0, n - span):
            i = start
            j = start + span
            frame_i = keyframes[i]
            frame_j = keyframes[j]
            reg = register_lidar_pair(frame_j.points_xyz, frame_i.points_xyz)
            if not reg.ok:
                continue
            if reg.rotation_deg < min_rotation_deg and reg.translation_m < min_translation_m:
                continue

            t_i_imu = lidar_time_to_imu_time(frame_i.t_mid_s, delta_t_s)
            t_j_imu = lidar_time_to_imu_time(frame_j.t_mid_s, delta_t_s)
            if t_j_imu <= t_i_imu:
                continue
            if t_i_imu < imu.t_s[0] or t_j_imu > imu.t_s[-1]:
                continue

            preint = preintegrate_imu(
                imu.t_s,
                imu.gyro_rad_s,
                imu.acc_m_s2,
                t_i_imu,
                t_j_imu,
                bias_g,
                bias_a,
            )
            r_a = preint.delta_R
            r_b = reg.transform[:3, :3]
            t_b = reg.transform[:3, 3]
            rot_a = rotation_angle_deg(r_a)
            if abs(rot_a - reg.rotation_deg) > max(15.0, 1.0 * max(rot_a, reg.rotation_deg)):
                continue

            pairs.append(
                MotionPair(
                    session_id=session_id,
                    i=int(keyframe_indices[i]),
                    j=int(keyframe_indices[j]),
                    t_i_s=frame_i.t_mid_s,
                    t_j_s=frame_j.t_mid_s,
                    R_A=r_a,
                    R_B=r_b,
                    t_A_m=np.asarray(preint.delta_p, dtype=float),
                    t_B_m=np.asarray(t_b, dtype=float),
                    fitness=reg.fitness,
                    metadata={
                        "backend": reg.backend,
                        "rotation_deg_B": reg.rotation_deg,
                        "translation_m_B": reg.translation_m,
                        "rotation_deg_A": rot_a,
                        "weight": preint.weight,
                        "duration_s": preint.duration_s,
                        "mean_gyro_norm": preint.mean_gyro_norm,
                        "preint_sigma_rad": preint.sigma_rad,
                        "cov": preint.cov[0:3, 0:3].tolist(),
                        "cov9": preint.cov.tolist(),
                        "J_bg": preint.J_bg[0:3, 0:3].tolist(),
                        "J_bg9": preint.J_bg.tolist(),
                        "J_ba": preint.J_ba.tolist(),
                        "delta_v": preint.delta_v.tolist(),
                        "delta_p": preint.delta_p.tolist(),
                        "t_i_imu_s": t_i_imu,
                        "t_j_imu_s": t_j_imu,
                        "modeling": "imu_preintegration_factor_phase_c",
                    },
                )
            )

    notes.append(
        f"built {len(pairs)} motion pairs (Phase-C preintegration: ΔR/Δv/Δp, Σ9, J_bg/J_ba)"
    )
    return MotionPairSet(pairs=tuple(pairs), notes=tuple(notes))


def pairs_to_transforms(pairs: tuple[MotionPair, ...]) -> tuple[list[np.ndarray], list[np.ndarray]]:
    """Helper returning SE(3) lists when translations are present."""

    a_list: list[np.ndarray] = []
    b_list: list[np.ndarray] = []
    for pair in pairs:
        if pair.t_B_m is None:
            continue
        t_a = np.zeros(3) if pair.t_A_m is None else pair.t_A_m
        a_list.append(make_transform(t_a, pair.R_A))
        b_list.append(make_transform(pair.t_B_m, pair.R_B))
    return a_list, b_list