#!/usr/bin/env python3
# SPDX-License-Identifier: Apache-2.0
"""Quickstart: the model card's Python snippet on the shipped sample, on one Blackhole p150.

    pip install -e .                          # once, from the repo root, on top of an environment that has ttnn (tt-metal)
    python examples/quickstart.py [input] [--variant ot128|qt128] [--out-dir examples/output]

Writes <out-dir>/quickstart.json (the same JSON as POST /predict) and <out-dir>/quickstart_bev.png: a bird's-eye view
of the output points coloured by the p150 classes, next to the stored fp32 CPU reference of the shipped sample and
the points where the two differ (for another input: the p150 panel only). The default input is the shipped
synthetic sample of the variant, found relative to this file (runs from any directory); an input given on the
command line is relative to the current directory.
"""
import argparse
import json
from pathlib import Path

REPO = Path(__file__).resolve().parents[1]
SAMPLES = REPO / "code" / "tt_frnet" / "samples"
ap = argparse.ArgumentParser()
ap.add_argument("input", nargs="?", default=None, help="sensor-frame cloud (default: the shipped synthetic sample)")
ap.add_argument("--variant", default="ot128", choices=["ot128", "qt128"])
ap.add_argument("--out-dir", default=str(REPO / "examples" / "output"))
ap.add_argument("--device-id", type=int, default=0)
args = ap.parse_args()
src = args.input or str(SAMPLES / f"synthetic_{args.variant}.npz")
out_dir = Path(args.out_dir)
out_dir.mkdir(parents=True, exist_ok=True)

# --- the model card snippet --------------------------------------------------------------------------------------
from tt_frnet import FRNet

# the sample's sensor -> base_link transform (calib/<stem>.json) feeds the node's ego crop box; for another cloud
# pass its own {"T_base_link_from_sensor": ...}, or {"ego_crop_box": None} to disable the box
stem = Path(src).name.split(".")[0]
calib_dir = REPO / "code" / "tt_frnet" / "calib"
calibration = {"preset": stem} if (calib_dir / f"{stem}.json").is_file() else {"ego_crop_box": None}

with FRNet.from_pretrained(device_id=args.device_id, variant=args.variant) as model:  # weights, traces captured
    out = model(src, calibration=calibration)                    # sensor-frame cloud + the ego crop transform
    palette = model.config.palette

print(out.class_counts())                                         # points per class (raw + interpolated)
# ------------------------------------------------------------------------------------------------------------------

(out_dir / "quickstart.json").write_text(json.dumps(out.to_dict(), indent=1))

import numpy as np  # noqa: E402

from tt_frnet import viz  # noqa: E402
from tt_frnet.ttaw.server.smoke import decode_array  # noqa: E402

ref_path = Path(src).with_name(f"{stem}.{args.variant}.reference.json")
title = f"FRNet {args.variant.upper()} on p150"
if ref_path.is_file():
    ref = json.loads(ref_path.read_text())
    ref_labels = np.asarray(decode_array(ref["labels"]).values, np.int64)
    agree = float(np.mean(ref_labels == out.label_ids)) if len(ref_labels) == len(out) else float("nan")
    img = viz.render_compare(out.points, out.label_ids, ref_labels, class_names=out.class_names, palette=palette,
                             titles=(title, "fp32 CPU reference (stored)"),
                             header=f"{Path(src).name}: {len(out):,} output points ({out.num_points_raw:,} "
                                    f"raw + Autoware's interpolated points), label agreement with the stored CPU "
                                    f"reference {agree * 100:.2f} %")
    print(f"label agreement with {ref_path.name}: {agree * 100:.3f} %")
else:
    img = viz.render_bev(out.points, out.label_ids, palette=palette, title=title)
img.save(out_dir / "quickstart_bev.png")
print(f"{len(out)} points -> {out_dir / 'quickstart.json'}, {out_dir / 'quickstart_bev.png'}  "
      f"timing_ms={ {k: round(v, 2) for k, v in out.timing_ms.items()} }")