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

from __future__ import annotations

from collections.abc import Callable
from dataclasses import dataclass
from time import perf_counter
from typing import Any

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,
    min_registration_fitness: float = 0.5,
    max_imu_gap_s: float = 0.05,
    max_lidar_gap_s: float = 1.0,
    all_frame_times_s: np.ndarray | None = None,
    max_index_span: int = 4,
    progress_callback: Callable[[dict[str, Any]], None] | None = None,
) -> 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] = []
    rejected_fitness = 0
    rejected_imu_gap = 0
    rejected_lidar_gap = 0
    frame_times = (
        None
        if all_frame_times_s is None
        else np.asarray(all_frame_times_s, dtype=float).reshape(-1)
    )
    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",))

    total_candidates = sum(max(n - span, 0) for span in range(1, max_index_span + 1))
    processed_candidates = 0
    started_at = perf_counter()
    last_progress_at = started_at

    def report_progress(*, event: str, span: int, force: bool = False) -> None:
        nonlocal last_progress_at
        if progress_callback is None:
            return
        now = perf_counter()
        if not force and processed_candidates > 1 and now - last_progress_at < 10.0:
            return
        last_progress_at = now
        progress_callback(
            {
                "event": event,
                "processed_candidates": processed_candidates,
                "total_candidates": total_candidates,
                "progress_pct": 100.0 * processed_candidates / max(total_candidates, 1),
                "current_span": span,
                "max_span": max_index_span,
                "accepted_pairs": len(pairs),
                "rejected_fitness": rejected_fitness,
                "rejected_imu_gap": rejected_imu_gap,
                "rejected_lidar_gap": rejected_lidar_gap,
                "elapsed_s": now - started_at,
            }
        )

    report_progress(event="start", span=1, force=True)

    for span in range(1, max_index_span + 1):
        for start in range(0, n - span):
            processed_candidates += 1
            report_progress(event="running", span=span)
            i = start
            j = start + span
            frame_i = keyframes[i]
            frame_j = keyframes[j]
            source_i = int(keyframe_indices[i])
            source_j = int(keyframe_indices[j])
            if frame_times is not None:
                lo = min(source_i, source_j)
                hi = max(source_i, source_j)
                local_times = frame_times[lo : hi + 1]
                if local_times.size >= 2 and np.any(np.diff(local_times) > max_lidar_gap_s):
                    rejected_lidar_gap += 1
                    continue
            reg = register_lidar_pair(frame_j.points_xyz, frame_i.points_xyz)
            if not reg.ok:
                continue
            if reg.fitness < min_registration_fitness:
                rejected_fitness += 1
                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
            imu_lo = max(int(np.searchsorted(imu.t_s, t_i_imu, side="right")) - 1, 0)
            imu_hi = min(
                int(np.searchsorted(imu.t_s, t_j_imu, side="left")) + 1,
                imu.t_s.size,
            )
            if imu_hi - imu_lo >= 2 and np.any(
                np.diff(imu.t_s[imu_lo:imu_hi]) > max_imu_gap_s
            ):
                rejected_imu_gap += 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,
                        "gyro_bias0_rad_s": bias_g.tolist(),
                        "accel_bias0_m_s2": bias_a.tolist(),
                        "time_offset_s": float(delta_t_s),
                        "keyframe_span": int(span),
                        "is_consecutive": bool(span == 1),
                        "modeling": "imu_preintegration_factor_phase_c",
                    },
                )
            )

    report_progress(event="complete", span=max_index_span, force=True)

    notes.append(
        f"built {len(pairs)} motion pairs (Phase-C preintegration: ΔR/Δv/Δp, Σ9, J_bg/J_ba)"
    )
    notes.append(
        "quality rejects: "
        f"fitness<{min_registration_fitness:.2f}: {rejected_fitness}, "
        f"IMU gap>{max_imu_gap_s:.3f}s: {rejected_imu_gap}, "
        f"LiDAR gap>{max_lidar_gap_s:.3f}s: {rejected_lidar_gap}"
    )
    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
