#!/usr/bin/env python3
# SPDX-License-Identifier: Apache-2.0
"""Build the SHIPPED samples of frnet-p150: one synthetic LiDAR sweep per serve profile, generated here (no dataset
involved, Apache-2.0), and their stored fp32 CPU-reference outputs (the container smoke test's reference). CPU only,
no device; the reference needs ONNX Runtime (research venv).

    cd bundles/frnet-p150 && OMP_NUM_THREADS=4 PYTHONPATH=code \\
        /home/ubuntu/experiments/tt-models/tools/research-venv/bin/python code/scripts/make_sample.py \\
        [--weights-dir <lidar_frnet v2.0 files>] [--check-against <dir with the same files>]

The scan is a deterministic ray cast from a sensor 2 m above a flat road (seed 0) against 25 cars, 16 building
facades, 30 poles / trees and 10 pedestrians (axis-aligned boxes), with 5 % beam dropout, 2 cm range noise and
random 0-255 intensities, in the **sensor frame** (x forward, y left, z up). The beam layout follows the profile's
own sensor: 128 rows evenly spread over the vertical field of view (OT128 +15 / -25 deg, QT128 +-52.6 deg) and a
fixed number of columns over 360 deg, chosen so that the frame lies inside Autoware's profile (the node skips a frame
with more than 60,000 unique cells): OT128 512 columns, QT128 600 columns. The generator is a copy of the research
reference ``research/frnet/scripts/frnet_preproc_ref.py synthetic_scan`` (``--check-against`` compares the arrays
with an earlier output, e.g. the M5 bench inputs ``logs/frnet/m5_baseline/inputs``).

Writes, per profile ``<v>`` (``ot128``, ``qt128``):

- ``code/tt_frnet/samples/synthetic_<v>.npz``: ``points``, a structured (x, y, z, intensity) float32 array;
- ``code/tt_frnet/calib/synthetic_<v>.json``: the calibration preset of the same name (sensor -> base_link = a pure
  2 m lift: base_link on the road under the sensor), for the ego crop box;
- ``code/tt_frnet/samples/synthetic_<v>.<v>.reference.json``: the ``/predict`` body of the fp32 CPU reference
  (``tt_frnet.reference.pipeline`` with ONNX Runtime on the deployed ONNX, the reference of record), reduced to what
  the container smoke compares (``labels``, ``scores``, ``filter_keep``, npz_compressed) plus the counts.

Regenerate the references whenever the reference, the weights or the post-processing changes.
"""
from __future__ import annotations

import argparse
import json
import sys
from pathlib import Path

import numpy as np

HERE = Path(__file__).resolve()
CODE = HERE.parents[1]
sys.path.insert(0, str(CODE))

WS = Path("/home/ubuntu/experiments/tt-models")
KEEP = ("model", "frame_id", "num_points", "num_points_raw", "class_names", "class_counts")
# per profile: the sensor's vertical field of view (deg) and rows (the ml_package frustum height) and the columns
SENSORS = {"ot128": dict(fov_up_deg=15.0, fov_down_deg=-25.0, rows=128, cols=512),
           "qt128": dict(fov_up_deg=52.6, fov_down_deg=-52.6, rows=128, cols=600)}
SENSOR_HEIGHT = 2.0
SEED = 0


def synthetic_scan(fov_up_deg: float, fov_down_deg: float, rows: int, n_cols: int, seed: int = 0,
                   sensor_height: float = 2.0, max_range: float = 200.0):
    """Ray cast against a ground plane and boxes; returns ``(points (n, 4) float32 x, y, z, intensity, labels)`` in
    a seeded random order (the order a driver would not guarantee either). Same arithmetic and random stream as
    ``research/frnet/scripts/frnet_preproc_ref.py synthetic_scan`` (cfg = the profile's FOV, ``interp_h`` rows)."""
    rng = np.random.default_rng(seed)
    elev = np.deg2rad(np.linspace(fov_down_deg + 0.05, fov_up_deg - 0.05, rows))
    azim = np.deg2rad(np.arange(n_cols) * (360.0 / n_cols) - 180.0)
    E, A = np.meshgrid(elev, azim, indexing="ij")
    d = np.stack([np.cos(E) * np.cos(A), np.cos(E) * np.sin(A), np.sin(E)], axis=-1).reshape(-1, 3)
    t_best = np.full(len(d), np.inf)
    label = np.full(len(d), -1)
    with np.errstate(divide="ignore", invalid="ignore"):          # ground plane z = -sensor_height
        tg = np.where(d[:, 2] < -1e-6, -sensor_height / d[:, 2], np.inf)
    upd = tg < t_best
    t_best[upd] = tg[upd]
    label[upd] = 0
    boxes = []                                                    # (cx, cy, cz_bottom, lx, ly, lz, class)
    for _ in range(25):                                           # cars
        cx, cy = rng.uniform(-60, 60), rng.choice([-1, 1]) * rng.uniform(3, 12)
        boxes.append((cx, cy, -sensor_height, 4.5, 1.9, 1.6, 5))
    for side in (-1, 1):                                          # building facades
        for k in range(8):
            boxes.append((-80 + 20 * k, side * rng.uniform(18, 25), -sensor_height, 15, 3, rng.uniform(6, 20), 3))
    for _ in range(30):                                           # poles / trees
        boxes.append((rng.uniform(-70, 70), rng.choice([-1, 1]) * rng.uniform(13, 16), -sensor_height, .4, .4, 6, 4))
    for _ in range(10):                                           # pedestrians
        boxes.append((rng.uniform(-30, 30), rng.choice([-1, 1]) * rng.uniform(12, 14), -sensor_height, .6, .6, 1.7,
                      17))
    for (cx, cy, cz, lx, ly, lz, cls) in boxes:
        lo = np.array([cx - lx / 2, cy - ly / 2, cz])
        hi = np.array([cx + lx / 2, cy + ly / 2, cz + lz])
        with np.errstate(divide="ignore", invalid="ignore"):
            inv = 1.0 / d
            t1 = lo * inv
            t2 = hi * inv
        tmin = np.nanmax(np.minimum(t1, t2), axis=1)
        tmax = np.nanmin(np.maximum(t1, t2), axis=1)
        hit = (tmax >= np.maximum(tmin, 0)) & (tmin > 0.5)
        upd = hit & (tmin < t_best)
        t_best[upd] = tmin[upd]
        label[upd] = cls
    ok = np.isfinite(t_best) & (t_best < max_range)
    ok &= rng.random(len(d)) > 0.05                               # beam dropout
    xyz = d[ok] * (t_best[ok, None] + rng.normal(0, 0.02, size=(ok.sum(), 1)))   # range noise
    inten = np.clip(rng.normal(40, 25, size=ok.sum()), 0, 255).round()          # 0..255 intensity
    pts = np.concatenate([xyz, inten[:, None]], axis=1).astype(np.float32)
    perm = rng.permutation(len(pts))
    return pts[perm], label[ok][perm]


def structured(xyzi: np.ndarray) -> np.ndarray:
    pts = np.zeros(len(xyzi), dtype=[("x", "<f4"), ("y", "<f4"), ("z", "<f4"), ("intensity", "<f4")])
    for i, name in enumerate(("x", "y", "z", "intensity")):
        pts[name] = xyzi[:, i]
    return pts


def main() -> int:
    ap = argparse.ArgumentParser(description=__doc__.split("\n\n")[0])
    ap.add_argument("--weights-dir", default=str(WS / "assets" / "frnet" / "hf_lidar_frnet"))
    ap.add_argument("--threads", type=int, default=4)
    ap.add_argument("--variants", nargs="*", default=list(SENSORS))
    ap.add_argument("--check-against", type=Path, help="directory with synthetic_<v>.npz to compare the points with")
    ap.add_argument("--no-reference", action="store_true", help="write the scans and presets only")
    a = ap.parse_args()
    samples, calib = CODE / "tt_frnet" / "samples", CODE / "tt_frnet" / "calib"
    T = np.eye(4)
    T[2, 3] = SENSOR_HEIGHT
    for variant in a.variants:
        s = SENSORS[variant]
        xyzi, _ = synthetic_scan(s["fov_up_deg"], s["fov_down_deg"], s["rows"], s["cols"], seed=SEED,
                                 sensor_height=SENSOR_HEIGHT)
        pts = structured(xyzi)
        stem = f"synthetic_{variant}"
        if a.check_against is not None:
            old = np.load(a.check_against / f"{stem}.npz", allow_pickle=False)["points"]
            if not (old.dtype == pts.dtype and np.array_equal(old, pts)):
                sys.exit(f"{stem}: the points differ from {a.check_against}")
            print(f"{stem}: identical to {a.check_against / (stem + '.npz')}", flush=True)
        np.savez_compressed(samples / f"{stem}.npz", points=pts)
        preset = {"name": stem,
                  "description": f"synthetic {variant.upper()}-like scan (code/scripts/make_sample.py): the sensor "
                                 f"{SENSOR_HEIGHT:g} m above a flat road, base_link on the road under it (a pure lift)",
                  "source": "code/scripts/make_sample.py (generated, no dataset; Apache-2.0)",
                  "T_base_link_from_sensor": T.tolist()}
        (calib / f"{stem}.json").write_text(json.dumps(preset, indent=1) + "\n")
        print(f"{stem}: {len(pts)} points ({s['rows']} rows x {s['cols']} columns, seed {SEED})", flush=True)
        if a.no_reference:
            continue
        from tt_frnet.io import encode_array
        from tt_frnet.reference.pipeline import FRNetReference

        ref = FRNetReference(a.weights_dir, variant, backend="ort", threads=a.threads)
        run = ref.run(samples / f"{stem}.npz", calibration={"preset": stem})
        if run.prepared.skipped is not None:
            sys.exit(f"{stem}: outside the {variant} profile: {run.prepared.skipped}")
        out = run.output
        body = {k: v for k, v in out.to_dict().items() if k in KEEP}
        body["labels"] = encode_array(out.label_ids, fmt="npz_compressed", key="labels")
        body["scores"] = encode_array(out.scores, fmt="npz_compressed", key="scores")
        body["filter_keep"] = encode_array(out.filter_keep, fmt="npz_compressed", key="filter_keep")
        inp = run.prepared.inputs
        body["meta"] = {"variant": variant, "reference": "onnxruntime fp32 (deployed ONNX)",
                        "onnx_sha256": ref.weights.onnx_sha256, "num_unique_coors": inp.num_voxels,
                        "cropped_by_ego_box": int((~inp.keep_mask).sum()),
                        "generator": f"code/scripts/make_sample.py synthetic_scan({s['rows']} rows, {s['cols']} "
                                     f"columns, seed {SEED})"}
        path = samples / f"{stem}.{variant}.reference.json"
        path.write_text(json.dumps(body) + "\n")
        print(f"{stem}: N {inp.num_points} (raw {inp.num_points_raw}) M {inp.num_voxels} -> {path.name} "
              f"({path.stat().st_size / 1e6:.2f} MB)", flush=True)
    return 0


if __name__ == "__main__":
    sys.exit(main())