"""Generate a tiny synthetic session for V1 smoke tests."""

from __future__ import annotations

from pathlib import Path

import numpy as np

from imu_lidar.contracts import ImuSeries, LidarFrame
from imu_lidar.geometry import so3_exp
from imu_lidar.imu_io import save_imu_csv
from imu_lidar.lidar_io import save_lidar_session


def _wall_cloud(rng: np.random.Generator, n: int = 800) -> np.ndarray:
    yz = rng.uniform([-5, -1], [5, 3], size=(n // 3, 2))
    wall_x = np.column_stack([np.full(n // 3, 8.0), yz[:, 0], yz[:, 1]])
    xz = rng.uniform([-5, -1], [5, 3], size=(n // 3, 2))
    wall_y = np.column_stack([xz[:, 0], np.full(n // 3, 6.0), xz[:, 1]])
    xy = rng.uniform([-5, -5], [5, 5], size=(n - 2 * (n // 3), 2))
    ground = np.column_stack([xy[:, 0], xy[:, 1], np.full(xy.shape[0], -1.0)])
    return np.vstack([wall_x, wall_y, ground])


def generate_synthetic_session(
    output_root: Path,
    *,
    delta_t_s: float = 0.17,
    yaw_extrinsic_deg: float = 25.0,
    seed: int = 0,
) -> dict[str, float]:
    """Write IMU CSV + LiDAR frames with known extrinsic rotation and time offset."""

    rng = np.random.default_rng(seed)
    output_root = Path(output_root)
    output_root.mkdir(parents=True, exist_ok=True)

    r_x = so3_exp(np.deg2rad(np.array([2.0, -1.5, yaw_extrinsic_deg])))
    map_points = _wall_cloud(rng)

    lidar_hz = 10.0
    duration = 8.0
    lidar_times = np.arange(0.0, duration, 1.0 / lidar_hz)
    # Non-yaw excitation is required for unique SO(3) hand-eye observability.
    yaw = 0.5 * np.sin(0.8 * lidar_times) + 0.12 * lidar_times
    pitch = 0.18 * np.sin(1.3 * lidar_times + 0.4)
    roll = 0.12 * np.sin(1.7 * lidar_times + 1.0)
    yaw_rate = np.gradient(yaw, lidar_times)
    pitch_rate = np.gradient(pitch, lidar_times)
    roll_rate = np.gradient(roll, lidar_times)

    frames: list[LidarFrame] = []
    for index, (t, yaw_i, pitch_i, roll_i) in enumerate(zip(lidar_times, yaw, pitch, roll)):
        r_wl = so3_exp(np.array([roll_i, pitch_i, yaw_i]))
        t_wl = np.array([0.4 * t, 0.05 * np.sin(0.5 * t), 0.0])
        points = (map_points - t_wl) @ r_wl
        points = points + rng.normal(0.0, 0.01, size=points.shape)
        frames.append(
            LidarFrame(
                frame_id=str(index),
                t_start_s=float(t),
                t_end_s=float(t + 0.08),
                points_xyz=points.astype(float),
            )
        )
    save_lidar_session(output_root / "lidar", frames)

    imu_hz = 100.0
    t_lidar_grid = np.arange(0.0, duration, 1.0 / imu_hz)
    omega_lidar = np.column_stack(
        [
            np.interp(t_lidar_grid, lidar_times, roll_rate),
            np.interp(t_lidar_grid, lidar_times, pitch_rate),
            np.interp(t_lidar_grid, lidar_times, yaw_rate),
        ]
    )
    omega_imu = omega_lidar @ r_x.T

    g_world = np.array([0.0, 0.0, 9.80665])
    acc_rows = []
    for yaw_i, pitch_i, roll_i in zip(
        np.interp(t_lidar_grid, lidar_times, yaw),
        np.interp(t_lidar_grid, lidar_times, pitch),
        np.interp(t_lidar_grid, lidar_times, roll),
    ):
        r_wl = so3_exp(np.array([roll_i, pitch_i, yaw_i]))
        g_in_lidar = r_wl.T @ g_world
        acc_rows.append(r_x @ g_in_lidar)
    acc = np.asarray(acc_rows, dtype=float)

    static_t = np.arange(-1.0, 0.0, 1.0 / imu_hz)
    static_gyro = np.zeros((static_t.size, 3))
    static_acc = np.tile(r_x @ g_world, (static_t.size, 1))

    t_imu = np.concatenate([static_t + delta_t_s, t_lidar_grid + delta_t_s])
    gyro = np.vstack([static_gyro, omega_imu]) + rng.normal(0.0, 0.001, size=(t_imu.size, 3))
    acc_all = np.vstack([static_acc, acc]) + rng.normal(0.0, 0.01, size=(t_imu.size, 3))
    imu = ImuSeries(t_s=t_imu, gyro_rad_s=gyro, acc_m_s2=acc_all)
    save_imu_csv(output_root / "imu.csv", imu)

    return {
        "delta_t_s": float(delta_t_s),
        "yaw_extrinsic_deg": float(yaw_extrinsic_deg),
    }


if __name__ == "__main__":
    import json

    out = Path("examples/synthetic_session")
    meta = generate_synthetic_session(out)
    (out / "meta.json").write_text(json.dumps(meta, indent=2), encoding="utf-8")
    print(f"wrote {out}")
    print(meta)