# SPDX-License-Identifier: Apache-2.0
"""Where the tests find data that does not ship in code/: the AutowareFoundation/lidar_frnet weights, the research
scripts and goldens of the workspace, the large module-tap goldens and the staged PandaSet sample. Outside the
tt-models workspace (the container image, a Hub clone) missing data skips the test that needs it. Inside the
workspace (``research/frnet`` exists) a missing golden, research frame or staged sample FAILS the test instead
(:func:`need`): a frozen gate must not vanish silently or run on a subset (VERIFY_PORT round 1, finding 2).

The shipped samples are the synthetic scans ``samples/synthetic_<variant>.npz`` (``code/scripts/make_sample.py``,
Apache-2.0) with the calibration presets of the same name. The PandaSet 019 frame 40 sample (CC BY 4.0 + PandaSet
Dataset Terms) does not ship until the user approves PandaSet-derived samples (PLAN.md 6.3, research/DATASETS.md N5):
it is staged in the git-ignored ``staging_samples_pandaset/`` of the bundle checkout (``$FRNET_PANDASET_DIR``
overrides; laid out like the package: ``samples/``, ``calib/``), and the frozen ``sample`` gates and the module-tap
frame still use it.

- weights: ``$FRNET_WEIGHTS_DIR``, else the workspace download ``assets/frnet/hf_lidar_frnet`` (sha256-verified,
  S:frnet:84-93), else the local HF cache snapshot at the pinned revision (``local_files_only``: tests never
  download);
- research: ``$FRNET_RESEARCH_DIR``, else ``research/frnet`` of the workspace (scripts, ``golden/`` VLS-128 and
  synthetic goldens, ``public_data/golden/`` PandaSet / nuScenes goldens; never shipped);
- large goldens: ``$FRNET_GOLDEN_DIR``, else ``research/frnet/goldens`` (``code/scripts/make_goldens.py``).

:data:`FRAMES` is the frame catalogue shared by the tests and ``code/scripts/make_goldens.py``.
"""
from __future__ import annotations

import json
import os
import sys
from pathlib import Path
from typing import Dict, List, Optional, Tuple

import pytest

PKG_DIR = Path(__file__).resolve().parents[1]
BUNDLE_DIR = PKG_DIR.parents[1]
SAMPLES = PKG_DIR / "samples"
GOLDENS = Path(__file__).resolve().parent / "goldens"
WORKSPACE = Path(os.environ.get("TT_MODELS_ROOT", "/home/ubuntu/experiments/tt-models"))
STAGING_PANDASET = BUNDLE_DIR / "staging_samples_pandaset"   # git-ignored, never shipped (PLAN.md 6.3)


def shipped(variant: str) -> Path:
    """The shipped synthetic sample of a variant (its calibration preset has the same stem)."""
    return SAMPLES / f"synthetic_{variant}.npz"


def shipped_calib(variant: str) -> dict:
    return {"preset": f"synthetic_{variant}"}


def _pandaset_root() -> Path:
    env = os.environ.get("FRNET_PANDASET_DIR")
    if env:
        return Path(env)
    if (SAMPLES / "pandaset_019_40.npz").is_file():          # swapped in after the user's OK
        return PKG_DIR
    return STAGING_PANDASET


# The PandaSet 019 f40 sample (the frozen ``sample`` gates and the module-tap frame): staged, not shipped.
PANDASET_ROOT = _pandaset_root()
SAMPLE_STEM = "pandaset_019_40"
SAMPLE_DIR = PANDASET_ROOT / "samples"
SAMPLE = SAMPLE_DIR / f"{SAMPLE_STEM}.npz"
SAMPLE_CALIB_FILE = PANDASET_ROOT / "calib" / f"{SAMPLE_STEM}.json"
# The explicit transform (not ``{"preset": ...}``: the preset is not in the package's calib/ while it is staged).
SAMPLE_CALIB = ({"T_base_link_from_sensor": json.loads(SAMPLE_CALIB_FILE.read_text())["T_base_link_from_sensor"]}
                if SAMPLE_CALIB_FILE.is_file() else None)


def in_workspace() -> bool:
    """True inside the tt-models workspace, where every golden and the staged sample must be present."""
    return (WORKSPACE / "research" / "frnet").is_dir()


def need(ok: bool, what: str) -> None:
    """Skip (outside the workspace) or FAIL (inside it) when ``ok`` is false: data a frozen gate needs."""
    if ok:
        return
    if in_workspace():
        pytest.fail(f"{what} is missing: inside the workspace a frozen gate must not be skipped")
    pytest.skip(f"{what} is missing")


def need_pandaset() -> Path:
    """The staged PandaSet 019 f40 sample (see the module docstring); skip outside / fail inside the workspace."""
    need(SAMPLE.is_file() and SAMPLE_CALIB is not None,
         f"the PandaSet 019 f40 sample {SAMPLE} and its calibration {SAMPLE_CALIB_FILE} (staging_samples_pandaset/ of "
         "the bundle checkout or $FRNET_PANDASET_DIR; it does not ship until the user approves PandaSet samples)")
    return SAMPLE


def _first_dir(*candidates: Optional[Path]) -> Optional[Path]:
    for c in candidates:
        if c is not None and c.is_dir():
            return c
    return None


def find_weights_dir() -> Optional[Path]:
    env = os.environ.get("FRNET_WEIGHTS_DIR")
    found = _first_dir(Path(env) if env else None, WORKSPACE / "assets" / "frnet" / "hf_lidar_frnet")
    if found is not None and (found / "frnet_ot128.onnx").is_file():
        return found
    try:
        from huggingface_hub import snapshot_download

        from tt_frnet.api import FRNet

        snap = Path(snapshot_download(FRNet.DEFAULT_REPO, revision=FRNet.DEFAULT_REVISION,
                                      allow_patterns=list(FRNet.ALLOW_PATTERNS), local_files_only=True))
        if (snap / "frnet_ot128.onnx").is_file():
            return snap
    except Exception:  # noqa: BLE001 -- not cached / no huggingface_hub: the caller skips
        pass
    return None


def weights_dir() -> Path:
    d = find_weights_dir()
    need(d is not None, "the lidar_frnet weights (FRNET_WEIGHTS_DIR, assets/frnet/hf_lidar_frnet or the HF cache)")
    return d


def research_dir() -> Path:
    env = os.environ.get("FRNET_RESEARCH_DIR")
    d = _first_dir(Path(env) if env else None, WORKSPACE / "research" / "frnet")
    need(d is not None and (d / "scripts" / "frnet_preproc_ref.py").is_file(),
         "research/frnet (the research reference scripts and goldens)")
    return d


def research_module(name: str):
    """Import a research script by name (``frnet_preproc_ref``, ``verify_preproc_ort_independent``, ...); their
    directory is workspace code, not a download. Importing them reads no data and runs nothing."""
    scripts = str(research_dir() / "scripts")
    if scripts not in sys.path:
        sys.path.insert(0, scripts)
    import importlib

    return importlib.import_module(name)


def golden_dir() -> Path:
    env = os.environ.get("FRNET_GOLDEN_DIR")
    d = _first_dir(Path(env) if env else None, WORKSPACE / "research" / "frnet" / "goldens")
    need(d is not None, "large goldens (research/frnet/goldens; run code/scripts/make_goldens.py)")
    return d


# ------------------------------------------------------------------------------------------ frame catalogue
# name -> (golden .npz relative to research/frnet, key prefix, variant). The public goldens hold the cloud handed
# to the node (raw_points), its crop transform and the four network inputs + ORT pred_probs (`autoware_` prefix);
# the VLS-128 goldens hold the x > 0 half's network inputs (no crop) + ORT pred_probs (S:frnet:361-365).
PUBLIC_OT128 = [f"pandaset_019_{f}" for f in (10, 25, 40, 55, 70)] + \
               [f"pandaset_029_{f}" for f in (10, 25, 40, 55, 70)] + \
               [f"nuscenes_scene-0103_{f:02d}" for f in (0, 20, 39)] + \
               [f"nuscenes_scene-0916_{f:02d}" for f in (0, 20, 40)]
PUBLIC_QT128 = ["pandaset_019_40", "pandaset_029_40", "nuscenes_scene-0103_20", "nuscenes_scene-0916_20"]


def frames() -> Dict[str, Tuple[str, str, str]]:
    out: Dict[str, Tuple[str, str, str]] = {}
    for name in PUBLIC_OT128:
        out[f"{name}.ot128"] = (f"public_data/golden/{name}_ot128.npz", "autoware_", "ot128")
    for name in PUBLIC_QT128:
        out[f"{name}.qt128"] = (f"public_data/golden/{name}_qt128.npz", "autoware_", "qt128")
    for v in ("ot128", "qt128"):
        out[f"vls128_frame000_fronthalf.{v}"] = (f"golden/vls128_frame000_fronthalf_{v}.npz", "autoware_", v)
    return out


FRAMES = frames()
# The frames whose module taps are stored in full (research/frnet/goldens/<variant>/<stem>.taps.npz): the shipped
# sample in both variants.
TAP_FRAMES: List[str] = [f"{SAMPLE_STEM}.ot128", f"{SAMPLE_STEM}.qt128"]
# The hardest public frame for bf16 (research/frnet/public_data README section 3) and the densest-cell frame.
HARD_FRAMES: List[str] = ["nuscenes_scene-0916_20.ot128", "vls128_frame000_fronthalf.qt128"]


def frame_path(name: str) -> Path:
    rel, _, _ = FRAMES[name]
    path = research_dir() / rel
    need(path.is_file(), f"golden {rel}")
    return path


def load_frame(name: str):
    """``(NetworkInputs, pred_probs (N, 27) ORT, golden npz)`` of a catalogued frame (unpadded)."""
    import numpy as np

    from tt_frnet.host.preprocess import NetworkInputs

    _, prefix, _ = FRAMES[name]
    g = np.load(frame_path(name), allow_pickle=False)
    n0 = len(g["raw_points"]) if "raw_points" in g.files else len(g[prefix + "points"])
    keep = g[prefix + "keep_mask"] if prefix + "keep_mask" in g.files else np.ones(n0, bool)
    n_raw = int(g[prefix + "num_points_raw"]) if prefix + "num_points_raw" in g.files else int(keep.sum())
    inputs = NetworkInputs(points=g[prefix + "points"], coors=g[prefix + "coors"],
                           voxel_coors=g[prefix + "voxel_coors"], inverse_map=g[prefix + "inverse_map"],
                           num_points_raw=n_raw, keep_mask=keep)
    return inputs, g[prefix + "pred_probs"], g