"""Load the parquet files of the Robot Localization Failure Prediction Dataset.

Expected layout (as extracted from ``preprocessed_data.zip``)::

    parquets/
        raw/<run>_id_01.parquet, <run>_id_02.parquet
        processed/<run>_id_01.processed.parquet, <run>_id_02.processed.parquet

Each run is split into two halves (``_id_01``/``_id_02``); the functions below
merge them again.

Only dependency: ``polars`` (``pip install polars``).

Examples
--------
>>> from load_parquets import load_run, load_dataset, load_map
>>> df = load_run("rec_20250821_104113", root="parquets")
>>> df.columns[:4]
['run', 'segment', 'time', '/scan/ranges']
>>> data = load_dataset(root="parquets", info_csv="info.csv")
>>> occupancy_map = load_map("rec_20250821_104113", root="parquets")

Command line summary of all runs::

    python load_parquets.py --root parquets --info info.csv

Notes on the LiDAR data
-----------------------
* ``-1.0`` in ``/scan/ranges`` marks a beam without a valid return
  (below ``range_min``). Treat these beams as invalid / out of range.
* In six runs (``rec_20250822_132819``, ``rec_20250822_135308``,
  ``rec_20250825_133605``, ``rec_20250825_135829``, ``rec_20250825_153216``,
  ``rec_20250825_161940``) the simulated LiDAR published at 70 Hz, but only
  every 7th message holds a full scan (10 Hz). The remaining messages contain
  only ``-1.0`` ("empty scans"). This is a simulation artifact, not a sensor
  dropout: valid scans arrive at a steady 10 Hz in these runs as well.
  Use ``empty_scans`` to drop or forward-fill these samples.
"""

from __future__ import annotations

import argparse
from pathlib import Path
from typing import Literal

import polars as pl

EmptyScans = Literal["keep", "drop", "ffill"]

SCAN_FIELDS = [
    "ranges",
    "angle_min",
    "angle_max",
    "angle_increment",
    "range_min",
    "range_max",
]


def _run_files(root: str | Path, run: str, kind: str) -> list[Path]:
    if kind == "processed":
        pattern = f"{run}_id_*.processed.parquet"
    elif kind == "raw":
        pattern = f"{run}_id_*.parquet"
    else:
        msg = f"kind must be 'processed' or 'raw', got {kind!r}"
        raise ValueError(msg)
    files = sorted((Path(root) / kind).glob(pattern))
    if not files:
        msg = f"No {kind} parquet files for run {run!r} in {Path(root) / kind}"
        raise FileNotFoundError(msg)
    return files


def is_empty_scan(column: str = "/scan/ranges") -> pl.Expr:
    """Expression that is True if every range of a scan is -1.0."""
    ranges = pl.col(column)
    return (ranges.list.min() == -1.0) & (ranges.list.max() == -1.0)


def list_runs(root: str | Path = "parquets") -> list[str]:
    """Return the names of all runs found in ``<root>/processed``."""
    files = (Path(root) / "processed").glob("*_id_*.processed.parquet")
    return sorted({f.name.split("_id_")[0] for f in files})


def load_run(
    run: str,
    root: str | Path = "parquets",
    empty_scans: EmptyScans = "drop",
) -> pl.DataFrame:
    """Load the processed samples of one run as a flat table.

    One row per labeled sample. Columns:

    * ``run``: run name, e.g. ``rec_20250821_104113``
    * ``segment``: index of the segment within the run. A new segment starts
      after every localization reset, so time is only continuous within a
      segment.
    * ``time``: simulation time in nanoseconds
    * ``/scan/*``: LaserScan (ranges, angles, range limits)
    * ``/particle_cloud/particles``: AMCL particles (pose + weight)
    * ``/momo/pose/*``: ground-truth pose
    * ``/amcl_pose/*``: AMCL pose estimate
    * ``position_error`` [m], ``heading_error`` [rad]
    * ``is_delocalized``: label (True = localization failure)

    Args:
        run: Run name as listed in ``info.csv``.
        root: Folder containing ``raw/`` and ``processed/``.
        empty_scans: How to handle samples whose scan contains only -1.0:
            ``"drop"`` removes them, ``"ffill"`` replaces them with the last
            valid scan of the same segment (samples before the first valid
            scan of a segment are removed), ``"keep"`` leaves them as is.
    """
    frames = []
    segment_offset = 0
    for path in _run_files(root, run, "processed"):
        part = (
            pl.read_parquet(path, columns=["measurements"])
            .with_row_index("segment", offset=segment_offset)
            .explode("measurements")
            .unnest("measurements")
            .unnest("value")
        )
        segment_offset += part["segment"].n_unique()
        frames.append(part)

    df = pl.concat(frames, how="vertical").select(
        pl.lit(run).alias("run"),
        pl.col("segment").cast(pl.Int32),
        pl.all().exclude("segment"),
    )
    return _handle_empty_scans(df, empty_scans)


def _handle_empty_scans(
    df: pl.DataFrame,
    empty_scans: EmptyScans,
) -> pl.DataFrame:
    if empty_scans == "keep":
        return df
    if empty_scans == "drop":
        return df.filter(~is_empty_scan())
    if empty_scans == "ffill":
        scan_columns = [f"/scan/{field}" for field in SCAN_FIELDS]
        empty = is_empty_scan()
        return (
            df.with_columns(
                pl.when(~empty).then(pl.col(c)).otherwise(None).alias(c)
                for c in scan_columns
            )
            .with_columns(
                pl.col(scan_columns).forward_fill().over("segment"),
            )
            .drop_nulls("/scan/ranges")
        )
    msg = f"empty_scans must be 'keep', 'drop' or 'ffill', got {empty_scans!r}"
    raise ValueError(msg)


def load_map(run: str, root: str | Path = "parquets") -> dict:
    """Return the occupancy map of a run.

    ``data`` is a flat list of booleans (True = occupied) in row-major order
    with ``info.width`` columns and ``info.height`` rows.
    """
    path = _run_files(root, run, "processed")[0]
    return pl.read_parquet(path, columns=["/map"])["/map"][0]


def load_dataset(
    runs: list[str] | None = None,
    root: str | Path = "parquets",
    info_csv: str | Path | None = "info.csv",
    empty_scans: EmptyScans = "drop",
) -> pl.DataFrame:
    """Load several runs into one table.

    If ``info_csv`` is given, the columns ``map``, ``dynamic`` and ``static``
    are added from it.
    """
    runs = runs or list_runs(root)
    df = pl.concat(
        [load_run(run, root, empty_scans) for run in runs],
        how="vertical",
    )
    if info_csv is not None:
        df = df.join(read_info(info_csv), on="run", how="left")
    return df


def read_info(path: str | Path = "info.csv") -> pl.DataFrame:
    """Read ``info.csv`` (without the total row) with typed columns."""
    return (
        pl.read_csv(path)
        .filter(pl.col("Name").str.starts_with("rec_"))
        .select(
            pl.col("Name").alias("run"),
            pl.col("Map").alias("map"),
            (pl.col("Dynamic") == "✓").alias("dynamic"),
            (pl.col("Static") == "✓").alias("static"),
        )
    )


def load_raw_run(
    run: str,
    root: str | Path = "parquets",
    empty_scans: Literal["keep", "drop"] = "drop",
) -> dict[str, pl.DataFrame]:
    """Load the raw (unsynchronized) topics of one run.

    Returns one table per topic (``/amcl_pose``, ``/momo/pose``, ``/scan``,
    ``/particle_cloud``, ``/delocalizations``) with a ``time`` column in
    nanoseconds plus the message fields. Topics have different rates and are
    not aligned; use e.g. ``DataFrame.join_asof`` on ``time`` to combine them.
    ``/map`` is returned as a single-row table.
    """
    topics: dict[str, list[pl.DataFrame]] = {}
    for path in _run_files(root, run, "raw"):
        raw = pl.read_parquet(path)
        for topic in raw.columns:
            messages = (
                raw.select(pl.col(topic).explode())
                .drop_nulls()
                .unnest(topic)
                .unnest("value")
            )
            topics.setdefault(topic, []).append(messages)

    tables = {
        topic: _merge_halves(parts).sort("time", maintain_order=True)
        for topic, parts in topics.items()
    }
    if empty_scans == "drop" and "/scan" in tables:
        tables["/scan"] = tables["/scan"].filter(~is_empty_scan("ranges"))
    return tables


def _merge_halves(parts: list[pl.DataFrame]) -> pl.DataFrame:
    """Concatenate the halves of a topic.

    Latched topics (e.g. ``/amcl_pose``, ``/map``) were replayed at the start
    of the second half: all earlier messages appear again, re-stamped with the
    first timestamp of that half. These replayed copies are removed.
    """
    merged = parts[0]
    for part in parts[1:]:
        value_columns = [c for c in part.columns if c != "time"]
        key = pl.struct(value_columns).hash().alias("_key")
        seen = merged.select(key)["_key"].implode()
        part = part.with_columns(key)
        replayed = (pl.col("time") == part["time"].min()) & pl.col(
            "_key",
        ).is_in(seen)
        merged = pl.concat(
            [merged, part.filter(~replayed).drop("_key")],
            how="vertical",
        )
    return merged


def _summary(root: str | Path, info_csv: str | Path | None) -> None:
    rows = []
    for run in list_runs(root):
        df = load_run(run, root, empty_scans="keep")
        rows.append(
            {
                "run": run,
                "segments": df["segment"].n_unique(),
                "samples": df.height,
                "failures": int(df["is_delocalized"].sum()),
                "empty_scans": int(df.select(is_empty_scan().sum()).item()),
            },
        )
    summary = pl.DataFrame(rows)
    if info_csv is not None:
        summary = read_info(info_csv).join(summary, on="run", how="right")
    with pl.Config(tbl_rows=-1, tbl_cols=-1):
        print(summary)


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description=__doc__.split("\n")[0])
    parser.add_argument("--root", default="parquets", help="parquets folder")
    parser.add_argument("--info", default=None, help="path to info.csv")
    args = parser.parse_args()
    _summary(args.root, args.info)