"""Scan a bag for points inside the train gauge ahead of the sensor.

Sensor frame: forward = -Y, lateral = X, up = Z.
"""
from __future__ import annotations

import sys

import numpy as np

sys.path.insert(0, r"C:\Games\Study\AI_Lidar\tools")
from probe_bag import frames  # noqa: E402


def main(db, n=None, half_width=1.5, z_lo=-0.2, z_hi=2.2, d_min=5.0, d_max=250.0):
    print("frame  ts        floor_z   n_gauge  nearest  clusters(d[m] x[m] z[m] npts)")
    for k, (ts, m) in enumerate(frames(db, limit=n)):
        p = m["points"]
        x = p["x"].astype(np.float32)
        y = p["y"].astype(np.float32)
        z = p["z"].astype(np.float32)
        v = ~((x == 0) & (y == 0) & (z == 0))
        d = -y  # forward distance

        # estimate floor/rail level from near-field points under the sensor
        near = v & (d > 4) & (d < 25) & (np.abs(x) < 2.0)
        floor = np.percentile(z[near], 2) if near.sum() > 100 else np.nan

        sel = v & (d > d_min) & (d < d_max) & (np.abs(x) < half_width) \
            & (z > floor + z_lo) & (z < floor + z_hi)
        idx = np.flatnonzero(sel)
        info = ""
        if idx.size:
            ds = d[idx]
            order = np.argsort(ds)
            ds = ds[order]
            idx = idx[order]
            # crude 1-D clustering along the corridor
            splits = np.flatnonzero(np.diff(ds) > 1.0) + 1
            parts = np.split(np.arange(ds.size), splits)
            chunks = []
            for pr in parts:
                if pr.size < 8:
                    continue
                ii = idx[pr]
                chunks.append((float(ds[pr].min()), float(np.median(x[ii])),
                               float(np.median(z[ii]) - floor), int(pr.size)))
            chunks.sort(key=lambda c: -c[3])
            info = "  ".join(f"[{c[0]:6.1f} {c[1]:+5.2f} {c[2]:+5.2f} {c[3]:5d}]" for c in chunks[:4])
        print(f"{k:5d} {ts/1e9:.2f} {floor:8.2f} {idx.size:8d} "
              f"{(d[idx].min() if idx.size else float('nan')):8.1f}  {info}")


if __name__ == "__main__":
    main(sys.argv[1], int(sys.argv[2]) if len(sys.argv) > 2 else None)