scripts/run_pipeline.py

"""
Single-mesh pipeline (SQ fitting only):
  1. load + normalize mesh
  2. build adaptive block map (or load existing)
  3. compute SDF volume (or load existing)
  4. run SDFAdaptiveFitter
  5. save results

Example:
  python scripts/run_pipeline.py \
    --mesh_path /path/to/watertight_mesh.ply \
    --output_root /path/to/output \
    --complex_levels 16,32 \
    --max_batches 50 --max_cands 5 --target_coverage 0.9 --max_iter 200
"""

import os
import sys
import argparse
import json
import numpy as np
import open3d as o3d
import trimesh

from hitops.adaptive.adaptive_res_block import build_adaptive_block_map, save_adaptive_map
from hitops.sdf.build_adaptive_sdf import (
    build_sdf_volume, save_sdf_volume, load_sdf_volume, sdf_volume_stats,
)
from hitops.sqfit.sq_fit import SDFAdaptiveFitter


# -----------------------------------------------------------------------
# Utility functions
# -----------------------------------------------------------------------

def normalize_to_unit_cube_safe(mesh, half=0.5, margin=1e-6):
    v = np.asarray(mesh.vertices).copy()
    f = np.asarray(mesh.triangles).copy()
    if v.size == 0 or f.size == 0:
        raise ValueError("Empty mesh.")
    if not np.isfinite(v).all():
        raise ValueError("Mesh has NaN/Inf vertices.")
    center = (v.min(0) + v.max(0)) * 0.5
    max_ext = float((v.max(0) - v.min(0)).max())
    if max_ext <= 0:
        raise ValueError(f"Invalid extent: {max_ext}")
    v = (v - center) * ((half - margin) * 2.0 / max_ext)
    out = o3d.geometry.TriangleMesh()
    out.vertices  = o3d.utility.Vector3dVector(v)
    out.triangles = o3d.utility.Vector3iVector(f)
    return out


def _load_adaptive_input(path: str) -> dict:
    """Load an adaptive map from an npz file or an npy dict."""
    loaded = np.load(path, allow_pickle=True)
    if isinstance(loaded, np.lib.npyio.NpzFile):
        if "level_map" in loaded and "block_size_map" in loaded:
            return {
                "level_map":      loaded["level_map"].astype(np.int32),
                "block_size_map": loaded["block_size_map"].astype(np.int32),
            }
        raise ValueError(f"NPZ missing required keys: {path}")
    obj = loaded.item()
    if "level_map" in obj and "block_size_map" in obj:
        return {
            "level_map":      np.asarray(obj["level_map"], dtype=np.int32),
            "block_size_map": np.asarray(obj["block_size_map"], dtype=np.int32),
        }
    if "map" in obj:
        lm = np.asarray(obj["map"], dtype=np.int32)
        print("[Warn] Legacy 'map' key, using fallback block_size_map=32.")
        return {"level_map": lm, "block_size_map": np.where(lm > 0, 32, 0).astype(np.int32)}
    raise ValueError(f"Cannot parse input file: {path}")


# -----------------------------------------------------------------------
# Main pipeline
# -----------------------------------------------------------------------

def run_pipeline(
    mesh_path:        str,
    output_root:      str,
    b_list:           list,
    r_max:            int,
    level_scheme:     str,
    clip_percentile:  float,
    complex_levels:   list,
    max_batches:      int,
    max_cands:        int,
    target_coverage:  float,
    max_iter:         int,
    cost_threshold:   float,
    max_total_sqs:    int,
    map_npz_path:     str = None,  # existing adaptive map
    sdf_npz_path:     str = None,  # existing SDF volume
):
    if not os.path.exists(mesh_path):
        raise FileNotFoundError(f"Mesh not found: {mesh_path}")

    mesh_name = os.path.splitext(os.path.basename(mesh_path))[0]
    out_dir   = os.path.join(output_root, mesh_name)
    map_dir   = os.path.join(out_dir, "adaptive_map")
    sdf_dir   = os.path.join(out_dir, "sdf_vol")
    sq_dir    = os.path.join(out_dir, "sq_fit_v20")
    for d in [map_dir, sdf_dir, sq_dir]:
        os.makedirs(d, exist_ok=True)

    # ---- Step 1: load mesh ----
    print(f"\n[1/4] Loading mesh: {mesh_path}")
    mesh = o3d.io.read_triangle_mesh(mesh_path)
    if len(mesh.vertices) == 0:
        raise ValueError("Empty mesh.")
    print(f"      Vertices: {len(mesh.vertices)}, Triangles: {len(mesh.triangles)}")

    # ---- Step 2: normalize ----
    print("[2/4] Normalizing to [-0.5, 0.5]^3")
    mesh = normalize_to_unit_cube_safe(mesh)

    # ---- Step 3: Adaptive block map ----
    if map_npz_path and os.path.exists(map_npz_path):
        print(f"[3a/4] Loading pre-built adaptive map: {map_npz_path}")
        adaptive_map = _load_adaptive_input(map_npz_path)
        saved_map    = map_npz_path
    else:
        print("[3a/4] Building adaptive block map")
        adaptive_map = build_adaptive_block_map(
            mesh=mesh, B_list=b_list, R_max=r_max,
            level_scheme=level_scheme, clip_percentile=clip_percentile,
        )
        saved_map = os.path.join(map_dir, f"{mesh_name}_adaptive_map.npz")
        save_adaptive_map(adaptive_map, saved_map)

    # ---- Step 3b: SDF volume ----
    R = adaptive_map["level_map"].shape[0]
    if sdf_npz_path and os.path.exists(sdf_npz_path):
        print(f"[3b/4] Loading pre-built SDF volume: {sdf_npz_path}")
        sdf_vol, truncation = load_sdf_volume(sdf_npz_path)
        saved_sdf = sdf_npz_path
    else:
        print(f"[3b/4] Computing SDF volume at {R}³ resolution...")
        sdf_vol, truncation = build_sdf_volume(mesh, resolution=R)
        saved_sdf = os.path.join(sdf_dir, f"{mesh_name}_sdf_vol.npz")
        save_sdf_volume(sdf_vol, truncation, saved_sdf)

    print("  SDF stats:")
    sdf_volume_stats(sdf_vol, truncation)

    # Report statistics on complex (high-resolution) regions
    lmap      = adaptive_map["level_map"]
    total_occ = int(np.sum(lmap > 0))
    comp_cnt  = int(np.sum(np.isin(lmap, complex_levels)))
    print(f"\n  complex_levels={complex_levels}:  {comp_cnt}/{total_occ}  ({comp_cnt/max(total_occ,1)*100:.1f}%)")

    # ---- Step 4: v20 SQ fitting ----
    print(f"\n[4/4] Running v20 SDFAdaptiveFitter (cost_threshold={cost_threshold})")
    fitter = SDFAdaptiveFitter(
        adaptive_map=adaptive_map,
        sdf_vol=sdf_vol,
        truncation=truncation,
        complex_levels=tuple(complex_levels),
    )
    fitter.run(
        max_batches=max_batches,
        max_cands=max_cands,
        target_coverage=target_coverage,
        max_iter=max_iter,
        cost_threshold=cost_threshold,
        max_total_sqs=max_total_sqs,
    )
    fitter.save_results(sq_dir)

    # ---- Summary ----
    summary = {
        "mesh_path":    mesh_path,
        "mesh_name":    mesh_name,
        "adaptive_map": saved_map,
        "sdf_vol":      saved_sdf,
        "sq_dir":       sq_dir,
        "n_sq":         len(fitter.all_sqs_params),
        "complex_levels": list(complex_levels),
        "complex_voxel_ratio": f"{comp_cnt/max(total_occ,1)*100:.1f}%",
        "params": {
            "B_list": b_list, "R_max": r_max,
            "level_scheme": level_scheme, "clip_percentile": clip_percentile,
            "complex_levels": list(complex_levels),
            "max_batches": max_batches, "max_cands": max_cands,
            "target_coverage": target_coverage, "max_iter": max_iter,
            "cost_threshold": cost_threshold,
        },
    }
    summary_path = os.path.join(out_dir, "run_summary_v20.json")
    with open(summary_path, "w") as f:
        json.dump(summary, f, indent=2)

    print(f"\n=== Done ===")
    print(f"Adaptive map : {saved_map}")
    print(f"SDF volume   : {saved_sdf}")
    print(f"SQ outputs   : {sq_dir}  ({len(fitter.all_sqs_params)} SQs)")
    print(f"Summary      : {summary_path}")


# -----------------------------------------------------------------------
# CLI
# -----------------------------------------------------------------------

def parse_int_list(s):
    return [int(x.strip()) for x in s.split(",") if x.strip()]


def main():
    parser = argparse.ArgumentParser(description="v20 adaptive SDF + SQ fitting pipeline")
    parser.add_argument("--mesh_path",        type=str, required=True)
    parser.add_argument("--output_root",      type=str, required=True,
                        help="output root; each mesh is written to <output_root>/<mesh_name>/")
    parser.add_argument("--b_list",           type=str, default="16,32,64")
    parser.add_argument("--r_max",            type=int, default=32)
    parser.add_argument("--level_scheme",     type=str, default="quantile",
                        choices=["quantile", "fixed"])
    parser.add_argument("--clip_percentile",  type=float, default=99.5)
    parser.add_argument("--complex_levels",   type=str, default="16,32")
    parser.add_argument("--max_batches",      type=int,   default=50)
    parser.add_argument("--max_cands",        type=int,   default=5)
    parser.add_argument("--target_coverage",  type=float, default=0.9)
    parser.add_argument("--max_iter",         type=int,   default=200)
    parser.add_argument("--cost_threshold",   type=float, default=0.05,
                        help="max SDF fitting error; SQs above this are discarded")
    parser.add_argument("--max_total_sqs",    type=int,   default=50)
    parser.add_argument("--map_npz_path",     type=str,   default=None,
                        help="path to an existing adaptive map .npz; skips map construction")
    parser.add_argument("--sdf_npz_path",     type=str,   default=None,
                        help="path to an existing SDF volume .npz; skips SDF computation")

    args = parser.parse_args()
    run_pipeline(
        mesh_path=args.mesh_path,
        output_root=args.output_root,
        b_list=parse_int_list(args.b_list),
        r_max=args.r_max,
        level_scheme=args.level_scheme,
        clip_percentile=args.clip_percentile,
        complex_levels=parse_int_list(args.complex_levels),
        max_batches=args.max_batches,
        max_cands=args.max_cands,
        target_coverage=args.target_coverage,
        max_iter=args.max_iter,
        cost_threshold=args.cost_threshold,
        max_total_sqs=args.max_total_sqs,
        map_npz_path=args.map_npz_path,
        sdf_npz_path=args.sdf_npz_path,
    )


if __name__ == "__main__":
    main()