"""Observable RTK--IMU rotation stages using native asynchronous measurements."""

from __future__ import annotations

import csv
import json
from dataclasses import asdict, dataclass
from pathlib import Path

import numpy as np
from scipy.optimize import least_squares
from scipy.spatial.transform import Rotation

from imu_lidar.contracts import ImuSeries
from imu_lidar.geometry import so3_exp, so3_log
from imu_lidar.imu_preintegration import apply_bias_jacobian_correction, preintegrate_gyro
from .rtk_attitude import gnhpr_to_baseline_enu


@dataclass(frozen=True)
class UnifiedSession:
    session_id: str
    batch_id: str
    imu: ImuSeries
    imu_rpy_deg: np.ndarray
    imu_quaternion_wxyz: np.ndarray
    imu_host_receive_utc_s: np.ndarray
    rtk_by_type: dict[str, list[dict[str, str]]]


@dataclass(frozen=True)
class R1bResult:
    baseline_axis_imu: np.ndarray
    tilt_yz_deg: np.ndarray
    pair_count: int
    residual_rms_deg: float
    residual_p95_deg: float
    covariance_deg2: np.ndarray
    std_deg: np.ndarray
    information_singular_values: np.ndarray
    per_session_rms_deg: dict[str, float]
    gyro_bias_by_session_rad_s: dict[str, np.ndarray]
    ok: bool
    notes: tuple[str, ...]


@dataclass(frozen=True)
class CompletedRotationResult:
    method: str
    R_RTK_IMU: np.ndarray | None
    rpy_deg: np.ndarray | None
    sample_count: int
    session_count: int
    residual_rms_deg: float
    residual_p95_deg: float
    covariance_deg2: np.ndarray
    std_deg: np.ndarray
    per_session_rms_deg: dict[str, float]
    leave_one_session_delta_deg: dict[str, float]
    block_out_delta_deg: dict[str, float]
    convention: str
    ok: bool
    notes: tuple[str, ...]


@dataclass(frozen=True)
class R3Result:
    r1b: R1bResult
    r2v: CompletedRotationResult
    r2g: CompletedRotationResult
    r2v_r2g_delta_deg: float
    full_rotation_accepted: bool
    translation_unlocked: bool
    blockers: tuple[str, ...]


@dataclass(frozen=True)
class _BaselinePair:
    session_index: int
    session_id: str
    world_angle_rad: float
    delta_R: np.ndarray
    J_bg: np.ndarray


def _f(row: dict[str, str], key: str, default: float = np.nan) -> float:
    try:
        return float(row.get(key, ""))
    except (TypeError, ValueError):
        return default


def _truth(row: dict[str, str], key: str) -> bool:
    return str(row.get(key, "")).strip().lower() in {"1", "true", "yes"}


def load_unified_sessions(
    manifest_path: Path | str,
    *,
    selected_session_ids: set[str] | None = None,
) -> list[UnifiedSession]:
    manifest_source = Path(manifest_path)
    manifest = json.loads(manifest_source.read_text(encoding="utf-8"))
    sessions: list[UnifiedSession] = []
    for entry in manifest["sessions"]:
        session_id = str(entry["session_id"])
        if selected_session_ids and session_id not in selected_session_ids:
            continue
        directory = Path(entry["directory"])
        if not directory.is_absolute():
            directory = manifest_source.parent / directory
        with np.load(directory / "imu.npz") as payload:
            t = np.asarray(payload["system_time_s"], dtype=float)
            gyro = np.asarray(payload["gyro_rad_s"], dtype=float)
            accel = np.asarray(payload["accel_m_s2"], dtype=float)
            rpy = np.asarray(payload["rpy_deg"], dtype=float)
            quaternion = np.asarray(payload["quaternion_wxyz"], dtype=float)
            host = np.asarray(payload["host_receive_utc_s"], dtype=float)
        by_type: dict[str, list[dict[str, str]]] = {}
        with (directory / "rtk.csv").open("r", encoding="utf-8", newline="") as stream:
            for row in csv.DictReader(stream):
                by_type.setdefault(row["message_type"], []).append(row)
        sessions.append(
            UnifiedSession(
                session_id=session_id,
                batch_id=str(entry["batch_id"]),
                imu=ImuSeries(t_s=t, gyro_rad_s=gyro, acc_m_s2=accel),
                imu_rpy_deg=rpy,
                imu_quaternion_wxyz=quaternion,
                imu_host_receive_utc_s=host,
                rtk_by_type=by_type,
            )
        )
    return sessions


def _valid_hpr(session: UnifiedSession) -> tuple[np.ndarray, np.ndarray]:
    rows = [
        row for row in session.rtk_by_type.get("GNHPR", [])
        if _truth(row, "checksum_valid") and int(_f(row, "heading_quality", -1)) == 4
    ]
    if not rows:
        return np.zeros(0), np.zeros((0, 3))
    t = np.asarray([_f(row, "t_device_s") for row in rows])
    baseline = gnhpr_to_baseline_enu(
        np.asarray([_f(row, "heading_deg") for row in rows]),
        np.asarray([_f(row, "pitch_deg") for row in rows]),
    )
    finite = np.isfinite(t) & np.all(np.isfinite(baseline), axis=1)
    t, baseline = t[finite], baseline[finite]
    order = np.argsort(t)
    t, baseline = t[order], baseline[order]
    unique, indices = np.unique(t, return_index=True)
    return unique, baseline[indices]


def _baseline_pairs(sessions: list[UnifiedSession]) -> list[_BaselinePair]:
    pairs: list[_BaselinePair] = []
    for session_index, session in enumerate(sessions):
        t, baseline = _valid_hpr(session)
        if t.size < 3:
            continue
        dt = np.diff(t)
        jump = np.degrees(
            np.arccos(np.clip(np.sum(baseline[:-1] * baseline[1:], axis=1), -1.0, 1.0))
        )
        continuous = (dt >= 0.03) & (dt <= 0.25) & (jump / np.maximum(dt, 1e-6) <= 45.0)
        last_anchor = -np.inf
        for index, t0 in enumerate(t[:-1]):
            if t0 - last_anchor < 1.0:
                continue
            last_anchor = t0
            for duration in (0.75, 1.5, 3.0):
                target = t0 + duration
                end = int(np.searchsorted(t, target))
                candidates = [candidate for candidate in (end - 1, end) if index < candidate < t.size]
                if not candidates:
                    continue
                j = min(candidates, key=lambda candidate: abs(t[candidate] - target))
                if abs((t[j] - t0) - duration) > 0.12 or not np.all(continuous[index:j]):
                    continue
                try:
                    pre = preintegrate_gyro(
                        session.imu.t_s, session.imu.gyro_rad_s, float(t0), float(t[j])
                    )
                except ValueError:
                    continue
                world_angle = float(
                    np.arccos(np.clip(np.dot(baseline[index], baseline[j]), -1.0, 1.0))
                )
                if np.degrees(world_angle) < 0.4:
                    continue
                pairs.append(
                    _BaselinePair(
                        session_index=session_index,
                        session_id=session.session_id,
                        world_angle_rad=world_angle,
                        delta_R=pre.delta_R,
                        J_bg=pre.J_bg,
                    )
                )
    return pairs


def _axis_from_parameters(parameters: np.ndarray) -> np.ndarray:
    axis = np.asarray([1.0, parameters[0], parameters[1]], dtype=float)
    return axis / np.linalg.norm(axis)


def solve_r1b(sessions: list[UnifiedSession]) -> R1bResult:
    pairs = _baseline_pairs(sessions)
    if len(pairs) < 20:
        raise ValueError("R1b needs at least 20 continuous baseline/gyro motion pairs")
    session_count = len(sessions)
    pair_counts = {
        session.session_id: sum(pair.session_id == session.session_id for pair in pairs)
        for session in sessions
    }
    active_sessions = sum(count > 0 for count in pair_counts.values())
    target_count = len(pairs) / max(active_sessions, 1)

    def residual(parameters: np.ndarray) -> np.ndarray:
        axis = _axis_from_parameters(parameters[:2])
        biases = parameters[2:].reshape(session_count, 3)
        values = []
        for pair in pairs:
            corrected = apply_bias_jacobian_correction(
                pair.delta_R, pair.J_bg, biases[pair.session_index]
            )
            body_angle = np.arccos(
                np.clip(np.dot(axis, corrected @ axis), -1.0, 1.0)
            )
            weight = np.sqrt(target_count / pair_counts[pair.session_id])
            values.append(weight * (body_angle - pair.world_angle_rad))
        values.extend((biases / 0.01).reshape(-1))
        # Weak 20-degree installation prior selects the physically known +X hemisphere.
        values.extend(np.asarray(parameters[:2]) / np.tan(np.deg2rad(20.0)))
        return np.asarray(values)

    initial = np.zeros(2 + 3 * session_count)
    optimum = least_squares(
        residual, initial, loss="huber", f_scale=np.deg2rad(0.25), max_nfev=120
    )
    axis = _axis_from_parameters(optimum.x[:2])
    biases = optimum.x[2:].reshape(session_count, 3)
    errors = []
    per_session_values: dict[str, list[float]] = {}
    for pair in pairs:
        corrected = apply_bias_jacobian_correction(
            pair.delta_R, pair.J_bg, biases[pair.session_index]
        )
        body_angle = np.arccos(np.clip(np.dot(axis, corrected @ axis), -1.0, 1.0))
        error = float(np.degrees(body_angle - pair.world_angle_rad))
        errors.append(error)
        per_session_values.setdefault(pair.session_id, []).append(error)
    errors_array = np.asarray(errors)
    data_rows = len(pairs)
    jacobian = optimum.jac[:data_rows, :2]
    information = jacobian.T @ jacobian
    residual_variance = float(np.mean(np.deg2rad(errors_array) ** 2))
    covariance = residual_variance * np.linalg.pinv(information, rcond=1e-12)
    covariance_deg2 = np.degrees(1.0) ** 2 * covariance
    std_deg = np.sqrt(np.maximum(np.diag(covariance_deg2), 0.0))
    singular = np.linalg.svd(information, compute_uv=False)
    rms = float(np.sqrt(np.mean(errors_array**2)))
    p95 = float(np.percentile(np.abs(errors_array), 95.0))
    return R1bResult(
        baseline_axis_imu=axis,
        tilt_yz_deg=np.degrees(np.arctan(optimum.x[:2])),
        pair_count=len(pairs),
        residual_rms_deg=rms,
        residual_p95_deg=p95,
        covariance_deg2=covariance_deg2,
        std_deg=std_deg,
        information_singular_values=singular,
        per_session_rms_deg={
            key: float(np.sqrt(np.mean(np.asarray(value) ** 2)))
            for key, value in per_session_values.items()
        },
        gyro_bias_by_session_rad_s={
            session.session_id: biases[index] for index, session in enumerate(sessions)
        },
        ok=bool(rms <= 1.0 and p95 <= 2.0 and np.min(singular) >= 1e-3),
        notes=(
            "2DoF ANT1-to-ANT2 direction; +X hemisphere selected by installation knowledge",
            "weak 20 deg prior is reported and prevents sign/gauge branch switching",
        ),
    )


def _rotation_mean(matrices: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
    mean = Rotation.from_matrix(matrices).mean().as_matrix()
    errors = np.asarray(
        [np.degrees(np.linalg.norm(so3_log(mean.T @ matrix))) for matrix in matrices]
    )
    return mean, errors


def _result_from_samples(
    method: str,
    samples: list[tuple[str, str, np.ndarray]],
    convention: str,
    notes: tuple[str, ...],
) -> CompletedRotationResult:
    if not samples:
        return CompletedRotationResult(
            method, None, None, 0, 0, np.nan, np.nan, np.full((3, 3), np.nan),
            np.full(3, np.nan),
            {}, {}, {}, convention, False, notes + ("no qualifying samples",)
        )
    matrices = np.asarray([item[2] for item in samples])
    mean, errors = _rotation_mean(matrices)
    rotvec = np.asarray([so3_log(mean.T @ matrix) for matrix in matrices])
    rotvec_deg = np.degrees(rotvec)
    std = np.std(rotvec_deg, axis=0, ddof=1) if len(samples) > 1 else np.full(3, np.nan)
    covariance = (
        np.cov(rotvec_deg, rowvar=False, ddof=1) / len(samples)
        if len(samples) > 1 else np.full((3, 3), np.nan)
    )
    ids = sorted({item[0] for item in samples})
    per_session = {}
    loo = {}
    for session_id in ids:
        selected = [item[2] for item in samples if item[0] == session_id]
        _, local_errors = _rotation_mean(np.asarray(selected))
        per_session[session_id] = float(np.sqrt(np.mean(local_errors**2)))
        kept = np.asarray([item[2] for item in samples if item[0] != session_id])
        if kept.size:
            kept_mean, _ = _rotation_mean(kept)
            loo[session_id] = float(np.degrees(np.linalg.norm(so3_log(mean.T @ kept_mean))))
    block_groups = sorted({(item[0], item[1]) for item in samples})
    block_out = {}
    for session_id, block_id in block_groups:
        kept = np.asarray(
            [item[2] for item in samples if (item[0], item[1]) != (session_id, block_id)]
        )
        if kept.size:
            kept_mean, _ = _rotation_mean(kept)
            block_out[f"{session_id}:{block_id}"] = float(
                np.degrees(np.linalg.norm(so3_log(mean.T @ kept_mean)))
            )
    rms = float(np.sqrt(np.mean(errors**2)))
    p95 = float(np.percentile(errors, 95.0))
    max_loo = max(loo.values(), default=np.inf)
    ok = bool(
        len(samples) >= 10 and len(ids) >= 2 and rms <= 2.0 and p95 <= 3.0
        and np.nanmax(std) <= 1.0 and max_loo <= 1.0
    )
    return CompletedRotationResult(
        method=method,
        R_RTK_IMU=mean,
        rpy_deg=Rotation.from_matrix(mean).as_euler("xyz", degrees=True),
        sample_count=len(samples),
        session_count=len(ids),
        residual_rms_deg=rms,
        residual_p95_deg=p95,
        covariance_deg2=covariance,
        std_deg=std,
        per_session_rms_deg=per_session,
        leave_one_session_delta_deg=loo,
        block_out_delta_deg=block_out,
        convention=convention,
        ok=ok,
        notes=notes,
    )


def solve_r2g(
    sessions: list[UnifiedSession],
    r1b: R1bResult,
    *,
    level_static_session_ids: set[str],
    block_duration_s: float = 10.0,
) -> CompletedRotationResult:
    samples: list[tuple[str, str, np.ndarray]] = []
    right = r1b.baseline_axis_imu
    for session in sessions:
        if session.session_id not in level_static_session_ids:
            continue
        gyro_norm = np.linalg.norm(session.imu.gyro_rad_s, axis=1)
        accel_norm = np.linalg.norm(session.imu.acc_m_s2, axis=1)
        valid = (gyro_norm <= np.deg2rad(0.35)) & (np.abs(accel_norm - 9.80665) <= 0.15)
        block = np.floor(
            (session.imu.t_s - session.imu.t_s[0]) / block_duration_s
        ).astype(int)
        for block_id in np.unique(block[valid]):
            selected = valid & (block == block_id)
            if np.count_nonzero(selected) < 200:
                continue
            up = np.median(session.imu.acc_m_s2[selected], axis=0)
            up /= np.linalg.norm(up)
            up -= right * np.dot(up, right)
            if np.linalg.norm(up) < 0.9:
                continue
            up /= np.linalg.norm(up)
            forward = np.cross(up, right)
            forward /= np.linalg.norm(forward)
            C_IMU_RTK = np.column_stack([right, forward, up])
            samples.append((session.session_id, str(int(block_id)), C_IMU_RTK.T))
    return _result_from_samples(
        "R2G_baseline_plus_level_gravity",
        samples,
        "accelerometer specific-force points vehicle up on explicit level-static blocks",
        (
            "only caller-declared level-static sessions are eligible",
            "result is level/gravity-prior constrained, not dual-antenna-only",
        ),
    )


def _nearest_index(t: np.ndarray, value: float, tolerance: float) -> int | None:
    index = int(np.searchsorted(t, value))
    candidates = [item for item in (index - 1, index) if 0 <= item < t.size]
    if not candidates:
        return None
    best = min(candidates, key=lambda item: abs(t[item] - value))
    return best if abs(t[best] - value) <= tolerance else None


def solve_r2v(
    sessions: list[UnifiedSession],
    *,
    min_speed_m_s: float = 1.5,
    max_yaw_rate_deg_s: float = 3.0,
    max_baseline_course_error_deg: float = 15.0,
) -> CompletedRotationResult:
    candidates: dict[str, list[tuple[str, str, np.ndarray]]] = {
        "HI13_q_body_to_ENU": [],
        "NED_to_ENU_times_HI13_q": [],
        "HI13_q_inverse_as_body_to_ENU": [],
        "NED_to_ENU_times_HI13_q_inverse": [],
    }
    NED_TO_ENU = np.asarray([[0.0, 1.0, 0.0], [1.0, 0.0, 0.0], [0.0, 0.0, -1.0]])
    for session in sessions:
        hpr_t, hpr_baseline = _valid_hpr(session)
        if hpr_t.size == 0:
            continue
        quaternion = session.imu_quaternion_wxyz
        norm = np.linalg.norm(quaternion, axis=1)
        valid_quaternion = np.isfinite(norm) & (np.abs(norm - 1.0) <= 0.02)
        normalized = quaternion / np.maximum(norm[:, None], 1e-12)
        imu_rotations = Rotation.from_quat(normalized[:, [1, 2, 3, 0]]).as_matrix()
        for row_index, row in enumerate(session.rtk_by_type.get("BESTNAVA", [])):
            if not (
                _truth(row, "checksum_valid")
                and _truth(row, "position_fixed")
                and _truth(row, "doppler_velocity_valid")
                and _f(row, "horizontal_speed_m_s") >= min_speed_m_s
                and _f(row, "horizontal_speed_std_m_s") <= 0.25
            ):
                continue
            t = _f(row, "t_device_s")
            hpr_index = _nearest_index(hpr_t, t, 0.15)
            imu_index = _nearest_index(session.imu.t_s, t, 0.03)
            if hpr_index is None or imu_index is None or not valid_quaternion[imu_index]:
                continue
            if abs(np.degrees(session.imu.gyro_rad_s[imu_index, 2])) > max_yaw_rate_deg_s:
                continue
            right = hpr_baseline[hpr_index]
            forward = np.asarray(
                [_f(row, "velocity_east_m_s"), _f(row, "velocity_north_m_s"), 0.0]
            )
            forward /= np.linalg.norm(forward)
            course_error = np.degrees(
                np.arcsin(np.clip(abs(np.dot(right, forward)), 0.0, 1.0))
            )
            if course_error > max_baseline_course_error_deg:
                continue
            forward -= right * np.dot(right, forward)
            forward /= np.linalg.norm(forward)
            up = np.cross(right, forward)
            if up[2] < 0:
                forward = -forward
                up = -up
            up /= np.linalg.norm(up)
            R_ENU_RTK = np.column_stack([right, forward, up])
            q = imu_rotations[imu_index]
            world_candidates = {
                "HI13_q_body_to_ENU": q,
                "NED_to_ENU_times_HI13_q": NED_TO_ENU @ q,
                "HI13_q_inverse_as_body_to_ENU": q.T,
                "NED_to_ENU_times_HI13_q_inverse": NED_TO_ENU @ q.T,
            }
            block_id = str(row_index // 10)
            for name, R_ENU_IMU in world_candidates.items():
                candidates[name].append(
                    (session.session_id, block_id, R_ENU_RTK.T @ R_ENU_IMU)
                )
    diagnostics = {
        name: _result_from_samples(
            "R2V_baseline_plus_doppler_velocity",
            values,
            name,
            (
                "RTK fixed + Doppler velocity + speed + low-yaw + baseline/course gates",
                "HI13 absolute quaternion may contain magnetic/navigation yaw bias",
            ),
        )
        for name, values in candidates.items()
    }
    finite = [result for result in diagnostics.values() if result.sample_count]
    if not finite:
        return diagnostics["HI13_q_body_to_ENU"]
    return min(finite, key=lambda result: result.residual_rms_deg)


def solve_r3(
    sessions: list[UnifiedSession],
    *,
    level_static_session_ids: set[str],
) -> R3Result:
    dynamic_sessions = [
        session for session in sessions if session.session_id not in level_static_session_ids
    ]
    r1b = solve_r1b(dynamic_sessions)
    r2v = solve_r2v(dynamic_sessions)
    r2g = solve_r2g(sessions, r1b, level_static_session_ids=level_static_session_ids)
    if r2v.R_RTK_IMU is None or r2g.R_RTK_IMU is None:
        delta = np.nan
    else:
        delta = float(
            np.degrees(np.linalg.norm(so3_log(r2v.R_RTK_IMU.T @ r2g.R_RTK_IMU)))
        )
    blockers = []
    if not r1b.ok:
        blockers.append("R1b baseline direction failed residual/observability gates")
    if not r2v.ok:
        blockers.append("R2V velocity-completed rotation failed stability gates")
    if not r2g.ok:
        blockers.append("R2G gravity-completed rotation failed stability gates")
    if not np.isfinite(delta) or delta > 2.0:
        blockers.append("R2V and R2G disagree by more than 2 deg")
    accepted = not blockers
    return R3Result(
        r1b=r1b,
        r2v=r2v,
        r2g=r2g,
        r2v_r2g_delta_deg=delta,
        full_rotation_accepted=accepted,
        translation_unlocked=accepted,
        blockers=tuple(blockers),
    )


def result_to_jsonable(result: R3Result) -> dict:
    def convert(value):
        if isinstance(value, np.ndarray):
            return value.tolist()
        if isinstance(value, np.generic):
            return value.item()
        if hasattr(value, "__dataclass_fields__"):
            return {key: convert(item) for key, item in asdict(value).items()}
        if isinstance(value, dict):
            return {str(key): convert(item) for key, item in value.items()}
        if isinstance(value, (tuple, list)):
            return [convert(item) for item in value]
        return value
    return convert(result)