"""Characterise the Pandar128 scan geometry inside the metro tunnel bags."""
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

W, H = 7200, 128


def as_image(p):
    """Return (H,W) arrays: ordering is col-major with ring cycling fastest."""
    x = p["x"].reshape(W, H).T.astype(np.float32)
    y = p["y"].reshape(W, H).T.astype(np.float32)
    z = p["z"].reshape(W, H).T.astype(np.float32)
    i = p["intensity"].reshape(W, H).T.astype(np.float32)
    return x, y, z, i


def main(db, idx=100):
    for ts, m in frames(db, limit=1, start=idx):
        x, y, z, inten = as_image(m["points"])
        valid = ~((x == 0) & (y == 0) & (z == 0))
        r = np.sqrt(x * x + y * y + z * z)

        # azimuth measured in the sensor XY plane; forward is -Y
        az = np.degrees(np.arctan2(x, -y))  # 0 = forward, + = right
        el = np.degrees(np.arcsin(np.clip(z / np.maximum(r, 1e-6), -1, 1)))

        print(f"valid {valid.sum()}/{valid.size} = {valid.mean():.1%}")

        print("\n--- azimuth per column (median over valid rings) ---")
        azc = np.where(valid, az, np.nan)
        with np.errstate(all="ignore"):
            colaz = np.nanmedian(azc, axis=0)
        good = np.isfinite(colaz)
        print("columns with any return:", good.sum())
        cols = np.arange(W)
        for c in [0, 1, 2, 1800, 3599, 3600, 3601, 5400, 7198, 7199]:
            print(f"  col {c:5d}: az={colaz[c]:8.3f}")
        d = np.diff(colaz[good])
        d = d[np.abs(d) < 1.0]
        print(f"  median azimuth step: {np.median(d):.4f} deg")

        print("\n--- elevation per ring (median over valid columns) ---")
        elr = np.where(valid, el, np.nan)
        with np.errstate(all="ignore"):
            ringel = np.nanmedian(elr, axis=1)
        print("  ring0..9 :", np.round(ringel[:10], 2))
        print("  ring60..69:", np.round(ringel[60:70], 2))
        print("  ring118..127:", np.round(ringel[118:], 2))
        print(f"  elevation span: {np.nanmin(ringel):.2f} .. {np.nanmax(ringel):.2f}")

        print("\n--- valid-return fraction by azimuth sector ---")
        for lo, hi in [(-180, -90), (-90, -30), (-30, -10), (-10, 10), (10, 30), (30, 90), (90, 180)]:
            sel = (colaz >= lo) & (colaz < hi)
            if sel.sum() == 0:
                print(f"  [{lo:4d},{hi:4d}) : no columns")
                continue
            v = valid[:, sel]
            print(f"  [{lo:4d},{hi:4d}) : {sel.sum():5d} cols, valid {v.mean():.1%}")

        print("\n--- forward cone (|az|<3 deg) range distribution ---")
        fwd = np.abs(colaz) < 3.0
        rf = r[:, fwd][valid[:, fwd]]
        print(f"  columns {fwd.sum()}, valid pts {rf.size}")
        print("  pct:", np.round(np.percentile(rf, [50, 90, 99, 99.9, 100]), 2))

        print("\n--- points beyond 100 m, in train gauge (|x|<1.6, -1<z<1.4) ---")
        far = valid & (r > 100)
        gauge = far & (np.abs(x) < 1.6) & (z > -1.0) & (z < 1.4)
        print(f"  far {far.sum()}, of them in gauge {gauge.sum()}")
        if gauge.sum():
            print("  max distance in gauge:", np.round(r[gauge].max(), 1))


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