"""Lever-arm calibration from RTK positions and full IMU preintegration."""

from __future__ import annotations

from dataclasses import dataclass

import numpy as np
from scipy.sparse import coo_matrix, csr_matrix, eye
from scipy.sparse.linalg import lsqr, splu
from scipy.spatial.transform import Rotation, Slerp

from imu_lidar.geometry import make_transform, orthonormalize_rotation, so3_exp
from .rtk_imu_rotation import RotationCalibrationResult, RotationSession, _attitude_rows


@dataclass(frozen=True)
class TranslationCalibrationResult:
    lever_IMU_to_RTK_in_IMU_m: np.ndarray
    t_RTK_IMU_m: np.ndarray
    T_RTK_IMU: np.ndarray
    translation_std_m: np.ndarray
    lever_information_singular_values: np.ndarray
    lever_precision_rank: int
    position_residual_rms_xyz_m: np.ndarray
    velocity_residual_rms_xyz_m_s: np.ndarray
    accel_bias_by_session_m_s2: dict[str, np.ndarray]
    knot_count_by_session: dict[str, int]
    loo_delta_m: dict[str, np.ndarray]
    ok: bool
    notes: tuple[str, ...]


@dataclass(frozen=True)
class _SessionFactors:
    session: RotationSession
    knot_t_s: np.ndarray
    position_enu_m: np.ndarray
    R_ENU_IMU: np.ndarray
    delta_p: tuple[np.ndarray, ...]
    delta_v: tuple[np.ndarray, ...]
    J_p_ba: tuple[np.ndarray, ...]
    J_v_ba: tuple[np.ndarray, ...]
    duration_s: np.ndarray


def _preintegrate_translation_interval(
    times_s: np.ndarray,
    gyro_rad_s: np.ndarray,
    acc_m_s2: np.ndarray,
    t0: float,
    t1: float,
    gyro_bias_rad_s: np.ndarray,
) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, float]:
    """Fast nominal ``delta_p/delta_v`` and accel-bias Jacobians.

    Rotation covariance and gyro-bias Jacobians are deliberately omitted here:
    rotation and gyro bias have already been fixed by Phase R1, while the
    translation linear system only consumes the accelerometer-bias Jacobians.
    """

    left = max(int(np.searchsorted(times_s, t0, side='left') - 1), 0)
    right = min(int(np.searchsorted(times_s, t1, side='right')), times_s.size - 1)
    delta_r = np.eye(3)
    delta_v = np.zeros(3)
    delta_p = np.zeros(3)
    j_v_ba = np.zeros((3, 3))
    j_p_ba = np.zeros((3, 3))
    for index in range(left, right):
        sample_t0 = float(times_s[index])
        sample_t1 = float(times_s[index + 1])
        if sample_t1 <= t0 or sample_t0 >= t1:
            continue
        segment_t0 = max(sample_t0, t0)
        segment_t1 = min(sample_t1, t1)
        dt = segment_t1 - segment_t0
        if dt <= 0.0:
            continue
        sample_dt = max(sample_t1 - sample_t0, 1e-12)
        u0 = (segment_t0 - sample_t0) / sample_dt
        u1 = (segment_t1 - sample_t0) / sample_dt
        gyro0 = (1.0 - u0) * gyro_rad_s[index] + u0 * gyro_rad_s[index + 1]
        gyro1 = (1.0 - u1) * gyro_rad_s[index] + u1 * gyro_rad_s[index + 1]
        acc0 = (1.0 - u0) * acc_m_s2[index] + u0 * acc_m_s2[index + 1]
        acc1 = (1.0 - u1) * acc_m_s2[index] + u1 * acc_m_s2[index + 1]
        omega = 0.5 * (gyro0 + gyro1) - gyro_bias_rad_s
        acc = 0.5 * (acc0 + acc1)
        r_i = delta_r
        delta_p = delta_p + delta_v * dt + 0.5 * r_i @ acc * dt**2
        delta_v = delta_v + r_i @ acc * dt
        j_p_ba = j_p_ba + j_v_ba * dt - 0.5 * r_i * dt**2
        j_v_ba = j_v_ba - r_i * dt
        delta_r = orthonormalize_rotation(delta_r @ so3_exp(omega * dt))
    return delta_p, delta_v, j_p_ba, j_v_ba, float(max(t1 - t0, 0.0))


def _make_session_factors(
    session: RotationSession,
    rotation: RotationCalibrationResult,
    *,
    knot_step_s: float,
) -> _SessionFactors:
    t_attitude, r_enu_rtk = _attitude_rows(session.rtk, rotation.convention)
    position_valid = session.rtk.position_valid
    t_position = session.rtk.t_s[position_valid]
    position = session.rtk.position_enu_m[position_valid]
    time_offset_s = rotation.applied_time_offset_s
    start = max(float(t_attitude[0]), float(t_position[0]), float(session.imu.t_s[0] - time_offset_s))
    end = min(float(t_attitude[-1]), float(t_position[-1]), float(session.imu.t_s[-1] - time_offset_s))
    if end - start < 5.0:
        raise ValueError(f"{session.session_id}: less than 5 s common RTK/IMU support")
    knot_t = np.arange(start + 0.25, end - 0.25, knot_step_s)
    if knot_t.size < 4:
        raise ValueError(f"{session.session_id}: not enough translation knots")
    position_knots = np.column_stack(
        [np.interp(knot_t, t_position, position[:, axis]) for axis in range(3)]
    )
    r_enu_rtk_knots = Slerp(t_attitude, Rotation.from_matrix(r_enu_rtk))(knot_t).as_matrix()
    r_enu_imu = r_enu_rtk_knots @ rotation.R_RTK_IMU
    bg = rotation.gyro_bias_by_session_rad_s[session.session_id]
    delta_p: list[np.ndarray] = []
    delta_v: list[np.ndarray] = []
    j_p_ba: list[np.ndarray] = []
    j_v_ba: list[np.ndarray] = []
    durations = []
    for t0, t1 in zip(knot_t[:-1], knot_t[1:]):
        dp, dv, jp, jv, duration = _preintegrate_translation_interval(
            session.imu.t_s,
            session.imu.gyro_rad_s,
            session.imu.acc_m_s2,
            float(t0 + time_offset_s),
            float(t1 + time_offset_s),
            bg,
        )
        delta_p.append(dp)
        delta_v.append(dv)
        j_p_ba.append(jp)
        j_v_ba.append(jv)
        durations.append(duration)
    return _SessionFactors(
        session=session,
        knot_t_s=knot_t,
        position_enu_m=position_knots,
        R_ENU_IMU=r_enu_imu,
        delta_p=tuple(delta_p),
        delta_v=tuple(delta_v),
        J_p_ba=tuple(j_p_ba),
        J_v_ba=tuple(j_v_ba),
        duration_s=np.asarray(durations),
    )


def _append_block(
    rows: list[int],
    cols: list[int],
    values: list[float],
    rhs: list[float],
    groups: list[int],
    matrix_blocks: list[tuple[int, np.ndarray]],
    vector: np.ndarray,
    sigma: np.ndarray,
    group: int,
) -> None:
    row0 = len(rhs)
    for axis in range(3):
        rhs.append(float(vector[axis] / sigma[axis]))
        groups.append(group)
        for col0, block in matrix_blocks:
            for local_col in range(block.shape[1]):
                value = float(block[axis, local_col] / sigma[axis])
                if value != 0.0:
                    rows.append(row0 + axis)
                    cols.append(col0 + local_col)
                    values.append(value)


def _build_system(
    factors: list[_SessionFactors],
    *,
    position_sigma_xyz_m: np.ndarray,
    velocity_sigma_xyz_m_s: np.ndarray,
) -> tuple[csr_matrix, np.ndarray, np.ndarray, dict[str, tuple[int, int]], list[tuple[str, int, str]]]:
    # x = [shared lever(3), per-session ba(3), per-knot velocities(3*K)]
    offsets: dict[str, tuple[int, int]] = {}
    variable_count = 3
    for item in factors:
        ba_offset = variable_count
        velocity_offset = ba_offset + 3
        offsets[item.session.session_id] = (ba_offset, velocity_offset)
        variable_count = velocity_offset + 3 * item.knot_t_s.size
    rows: list[int] = []
    cols: list[int] = []
    values: list[float] = []
    rhs: list[float] = []
    groups: list[int] = []
    factor_labels: list[tuple[str, int, str]] = []
    gravity = np.array([0.0, 0.0, -9.80665])
    group = 0
    for item in factors:
        ba_offset, velocity_offset = offsets[item.session.session_id]
        for index, dt in enumerate(item.duration_s):
            r_i = item.R_ENU_IMU[index]
            r_j = item.R_ENU_IMU[index + 1]
            dp_rtk = item.position_enu_m[index + 1] - item.position_enu_m[index]
            constant_p = dp_rtk - 0.5 * gravity * dt**2 - r_i @ item.delta_p[index]
            _append_block(
                rows,
                cols,
                values,
                rhs,
                groups,
                [
                    (0, r_i - r_j),
                    (ba_offset, -r_i @ item.J_p_ba[index]),
                    (velocity_offset + 3 * index, -dt * np.eye(3)),
                ],
                -constant_p,
                position_sigma_xyz_m,
                group,
            )
            factor_labels.append((item.session.session_id, group, "position"))
            group += 1
            constant_v = -gravity * dt - r_i @ item.delta_v[index]
            _append_block(
                rows,
                cols,
                values,
                rhs,
                groups,
                [
                    (ba_offset, -r_i @ item.J_v_ba[index]),
                    (velocity_offset + 3 * index, -np.eye(3)),
                    (velocity_offset + 3 * (index + 1), np.eye(3)),
                ],
                -constant_v,
                velocity_sigma_xyz_m_s,
                group,
            )
            factor_labels.append((item.session.session_id, group, "velocity"))
            group += 1
        # Loose physical bias prior. It prevents an unobservable constant
        # acceleration from masquerading as gravity while remaining data-led.
        _append_block(
            rows,
            cols,
            values,
            rhs,
            groups,
            [(ba_offset, np.eye(3))],
            np.zeros(3),
            np.full(3, 0.5),
            group,
        )
        factor_labels.append((item.session.session_id, group, "bias_prior"))
        group += 1
    matrix = coo_matrix((values, (rows, cols)), shape=(len(rhs), variable_count)).tocsr()
    return matrix, np.asarray(rhs), np.asarray(groups), offsets, factor_labels


def _irls(matrix: csr_matrix, rhs: np.ndarray, groups: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
    row_weights = np.ones(rhs.size)
    solution = np.zeros(matrix.shape[1])
    for _ in range(5):
        weighted = matrix.multiply(row_weights[:, None])
        solution = lsqr(weighted, rhs * row_weights, atol=1e-10, btol=1e-10, iter_lim=3000)[0]
        residual = matrix @ solution - rhs
        new_weights = np.ones_like(row_weights)
        for group in np.unique(groups):
            selection = groups == group
            norm = float(np.linalg.norm(residual[selection]))
            if norm > 3.0:
                new_weights[selection] = np.sqrt(3.0 / norm)
        if np.max(np.abs(new_weights - row_weights)) < 1e-3:
            row_weights = new_weights
            break
        row_weights = new_weights
    return solution, row_weights


def _solve_factors(
    factors: list[_SessionFactors],
    position_sigma: np.ndarray,
    velocity_sigma: np.ndarray,
) -> tuple[np.ndarray, np.ndarray, csr_matrix, np.ndarray, dict[str, tuple[int, int]], np.ndarray, np.ndarray]:
    matrix, rhs, groups, offsets, labels = _build_system(
        factors,
        position_sigma_xyz_m=position_sigma,
        velocity_sigma_xyz_m_s=velocity_sigma,
    )
    solution, row_weights = _irls(matrix, rhs, groups)
    weighted = matrix.multiply(row_weights[:, None]).tocsr()
    residual = matrix @ solution - rhs
    data_groups = {group for _, group, kind in labels if kind != "bias_prior"}
    data_rows = np.isin(groups, list(data_groups))
    variance = float(np.sum((residual[data_rows] * row_weights[data_rows]) ** 2) / max(np.count_nonzero(data_rows) - solution.size, 1))
    information = (weighted.T @ weighted).tocsc() + eye(weighted.shape[1], format="csc") * 1e-10
    h_ll = information[:3, :3].toarray()
    h_ln = information[:3, 3:]
    h_nn = information[3:, 3:]
    nuisance_solve = splu(h_nn).solve(h_ln.T.toarray())
    schur = h_ll - h_ln.toarray() @ nuisance_solve
    covariance_lever = np.linalg.pinv(schur, rcond=1e-10) * variance
    return solution, covariance_lever, matrix, rhs, offsets, groups, residual


def solve_rtk_imu_translation(
    sessions: list[RotationSession] | tuple[RotationSession, ...],
    rotation: RotationCalibrationResult,
    *,
    knot_step_s: float = 2.0,
    compute_loo: bool = True,
) -> TranslationCalibrationResult:
    """Estimate the shared IMU-to-RTK lever arm and return ``T_RTK_IMU``."""

    items = list(sessions)
    factors = [_make_session_factors(session, rotation, knot_step_s=knot_step_s) for session in items]
    position_sigma = np.array([0.025, 0.025, 0.060])
    velocity_sigma = np.array([0.08, 0.08, 0.12])
    solution, covariance_l, matrix, rhs, offsets, groups, residual = _solve_factors(
        factors, position_sigma, velocity_sigma
    )
    lever = solution[:3]
    t_rtk_imu = -rotation.R_RTK_IMU @ lever
    covariance_t = rotation.R_RTK_IMU @ covariance_l @ rotation.R_RTK_IMU.T
    std_t = np.sqrt(np.maximum(np.diag(covariance_t), 0.0))
    schur_information = np.linalg.pinv(covariance_l, rcond=1e-12)
    singular_values = np.linalg.svd(schur_information, compute_uv=False)
    threshold = max(float(singular_values[0]) * 1e-4, 1e-9)
    rank = int(np.count_nonzero(singular_values > threshold))

    # Recover physical residuals: system rows are grouped in XYZ triples and
    # alternate position/velocity, followed by one bias prior per session.
    position_errors: list[np.ndarray] = []
    velocity_errors: list[np.ndarray] = []
    cursor = 0
    for item in factors:
        for _ in range(item.knot_t_s.size - 1):
            position_errors.append(residual[cursor : cursor + 3] * position_sigma)
            cursor += 3
            velocity_errors.append(residual[cursor : cursor + 3] * velocity_sigma)
            cursor += 3
        cursor += 3
    pos_rms = np.sqrt(np.mean(np.asarray(position_errors) ** 2, axis=0))
    vel_rms = np.sqrt(np.mean(np.asarray(velocity_errors) ** 2, axis=0))
    biases = {
        item.session.session_id: solution[offsets[item.session.session_id][0] : offsets[item.session.session_id][0] + 3].copy()
        for item in factors
    }
    loo: dict[str, np.ndarray] = {}
    if compute_loo and len(factors) >= 3:
        for omitted in factors:
            kept = [item for item in factors if item.session.session_id != omitted.session.session_id]
            loo_solution, *_ = _solve_factors(kept, position_sigma, velocity_sigma)
            loo[omitted.session.session_id] = (-rotation.R_RTK_IMU @ loo_solution[:3]) - t_rtk_imu
    max_loo_xy = max((float(np.linalg.norm(value[:2])) for value in loo.values()), default=0.0)
    max_loo_z = max((abs(float(value[2])) for value in loo.values()), default=0.0)
    ok = bool(
        rank == 3
        and float(np.max(std_t[:2])) <= 0.05
        and float(std_t[2]) <= 0.10
        and float(np.max(pos_rms[:2])) <= 0.10
        and float(pos_rms[2]) <= 0.20
        and max_loo_xy <= 0.10
        and max_loo_z <= 0.20
    )
    notes = [
        "lever l is vector IMU-origin -> RTK-origin expressed in IMU",
        "transform translation uses t_RTK_IMU = -R_RTK_IMU @ l",
        "RTK position is never differentiated; position and velocity preintegration factors are solved jointly",
    ]
    if not rotation.ok:
        notes.append("upstream rotation is not accepted, so translation is diagnostic only")
        ok = False
    if not ok:
        notes.append("translation failed one or more strict acceptance gates")
    return TranslationCalibrationResult(
        lever_IMU_to_RTK_in_IMU_m=lever,
        t_RTK_IMU_m=t_rtk_imu,
        T_RTK_IMU=make_transform(t_rtk_imu, rotation.R_RTK_IMU),
        translation_std_m=std_t,
        lever_information_singular_values=singular_values,
        lever_precision_rank=rank,
        position_residual_rms_xyz_m=pos_rms,
        velocity_residual_rms_xyz_m_s=vel_rms,
        accel_bias_by_session_m_s2=biases,
        knot_count_by_session={item.session.session_id: int(item.knot_t_s.size) for item in factors},
        loo_delta_m=loo,
        ok=ok,
        notes=tuple(notes),
    )
