"""Build a position -> (material, layer_id) geometry oracle from steps parquet.

Backs the `dwarf build-geometry-oracle` subcommand. The oracle is consumed by
`giant rollout` to supply the material/layer conditioning at each step, since the
surrogate does not predict them.
"""

from __future__ import annotations

from pathlib import Path

from giant.data.loader import find_parquet_files
from giant.geometry import build_geometry_oracle


def run_build_geometry_oracle(
    data: Path,
    out: Path,
    method: str = "slab",
    k: int = 1,
    subsample: int = 500_000,
    escape_factor: float = 5.0,
    seed: int = 0,
    depth_axis: int = 2,
    n_bins: int = 2000,
) -> None:
    files = find_parquet_files(data)
    print(f"found {len(files)} parquet file(s); sampling up to {subsample:,} points")

    oracle = build_geometry_oracle(
        files,
        method=method,
        k=k,
        subsample=subsample,
        escape_factor=escape_factor,
        seed=seed,
        depth_axis=depth_axis,
        n_bins=n_bins,
    )

    print(
        f"method: {method}  reference points: {oracle.metadata['n_reference_points']:,}"
    )
    print("classes (material, layer_id):")
    for material, layer_id in oracle.classes:
        print(f"  {material:<12} layer_id={layer_id}")

    if method == "slab":
        z_lo, z_hi = oracle.metadata["z_range"]
        print(
            f"depth axis: {'xyz'[depth_axis]}  segments: {oracle.metadata['n_segments']}  "
            f"z range: [{z_lo:.3f}, {z_hi:.3f}]  radius_max: {oracle.metadata['radius_max']:.3f}"
        )
        print(
            f"median depth spacing: {oracle.metadata['median_z_spacing']:.3f}  "
            f"escape_threshold: {oracle.escape_threshold:.3f} (= {escape_factor}x spacing)"
        )
    else:
        print(
            f"median NN spacing: {oracle.metadata['median_nn_dist']:.3f}  "
            f"escape_threshold: {oracle.escape_threshold:.3f} "
            f"(= {escape_factor}x spacing)"
        )
    if oracle.escape_threshold <= 0.0:
        print(
            "warning: escape_threshold is 0 (reference points are coincident) — "
            "every rollout query would be flagged as escaped. Pass "
            "`giant rollout --escape-threshold <mm>` to override, or use data "
            "with distinct step positions."
        )

    out.parent.mkdir(parents=True, exist_ok=True)
    oracle.save(out)
    print(f"wrote oracle -> {out}")