"""LiDAR adapters for the V1 standard intermediate format.

Accepted input: a directory containing ``frames_index.csv`` and per-frame NPZ files.

frames_index.csv
----------------
frame_id,file,t_start,t_end

Each NPZ referenced by ``file`` must contain:
- points: float array shaped (N, 3) in LiDAR Cartesian coordinates (metres)
"""

from __future__ import annotations

from pathlib import Path

import numpy as np

from .contracts import LidarFrame


def _read_frames_index(root: Path) -> tuple[np.ndarray, str]:
    index_path = root / "frames_index.csv"
    if not index_path.exists():
        raise FileNotFoundError(f"missing frames_index.csv under {root}")

    rows = np.genfromtxt(index_path, delimiter=",", names=True, dtype=None, encoding="utf-8")
    if rows.ndim == 0:
        rows = np.array([rows])
    names = set(rows.dtype.names or ())
    # NumPy may rename reserved name ``file`` to ``file_``.
    file_key = "filename" if "filename" in names else ("file_" if "file_" in names else "file")
    required = {"frame_id", "t_start", "t_end"}
    if not required.issubset(names) or file_key not in names:
        raise ValueError(
            f"frames_index.csv must contain frame_id,{file_key}/filename,t_start,t_end; got {sorted(names)}"
        )
    return rows, file_key


def list_lidar_frame_entries(path: Path | str) -> list[tuple[str, float, float, Path]]:
    """Return ``(frame_id, t_start, t_end, npz_path)`` sorted by mid time (same as ``load_lidar_frames``)."""

    root = Path(path)
    rows, file_key = _read_frames_index(root)
    entries: list[tuple[str, float, float, Path]] = []
    for row in rows:
        t0 = float(row["t_start"])
        t1 = float(row["t_end"])
        entries.append((str(row["frame_id"]), t0, t1, root / str(row[file_key])))
    entries.sort(key=lambda item: 0.5 * (item[1] + item[2]))
    return entries


def load_lidar_frame_at(root: Path | str, index: int) -> LidarFrame:
    """Load one frame by index in mid-time-sorted order (matches motion-pair ``i``/``j``)."""

    entries = list_lidar_frame_entries(root)
    if index < 0 or index >= len(entries):
        raise IndexError(f"frame index {index} outside [0, {len(entries) - 1}] for {root}")
    frame_id, t0, t1, npz_path = entries[index]
    with np.load(npz_path) as payload:
        if "points" not in payload.files:
            raise ValueError(f"{npz_path} must contain array 'points'")
        points = np.asarray(payload["points"], dtype=float)
    if points.ndim != 2 or points.shape[1] < 3:
        raise ValueError(f"{npz_path}: points must have shape (N, 3[+])")
    return LidarFrame(
        frame_id=frame_id,
        t_start_s=t0,
        t_end_s=t1,
        points_xyz=points[:, :3],
        path=npz_path,
    )


def lidar_frame_count(path: Path | str) -> int:
    return len(list_lidar_frame_entries(path))


def load_lidar_frames(path: Path | str) -> list[LidarFrame]:
    """Load all LiDAR frames listed by ``frames_index.csv`` under ``path``."""

    root = Path(path)
    rows, file_key = _read_frames_index(root)

    frames: list[LidarFrame] = []
    for row in rows:
        frame_id = str(row["frame_id"])
        rel = str(row[file_key])
        npz_path = root / rel
        with np.load(npz_path) as payload:
            if "points" not in payload.files:
                raise ValueError(f"{npz_path} must contain array 'points'")
            points = np.asarray(payload["points"], dtype=float)
        if points.ndim != 2 or points.shape[1] < 3:
            raise ValueError(f"{npz_path}: points must have shape (N, 3[+])")
        frames.append(
            LidarFrame(
                frame_id=frame_id,
                t_start_s=float(row["t_start"]),
                t_end_s=float(row["t_end"]),
                points_xyz=points[:, :3],
                path=npz_path,
            )
        )
    frames.sort(key=lambda frame: frame.t_mid_s)
    return frames


def save_lidar_session(
    root: Path | str,
    frames: list[LidarFrame],
    *,
    points_dirname: str = "frames",
) -> None:
    """Write a LiDAR session directory in the standard intermediate format."""

    destination = Path(root)
    frames_dir = destination / points_dirname
    frames_dir.mkdir(parents=True, exist_ok=True)
    index_rows: list[str] = ["frame_id,filename,t_start,t_end"]
    for index, frame in enumerate(frames):
        relative = f"{points_dirname}/frame_{index:05d}.npz"
        np.savez_compressed(destination / relative, points=np.asarray(frame.points_xyz, dtype=float))
        index_rows.append(f"{frame.frame_id},{relative},{frame.t_start_s:.9f},{frame.t_end_s:.9f}")
    (destination / "frames_index.csv").write_text("\n".join(index_rows) + "\n", encoding="utf-8")