scripts/run_batch_sqfit.py

"""
HY3D pipeline: sq_fit_v20 + mesh_mapper_v8 combined batch runner.

Input layout:
    <mesh_root>/<shard>/<uid>/full.ply

Output layout:
    <output_root>/<uid>/
        adaptive_map_v4/<uid>_adaptive_map_v4.npz
        sdf_vol/<uid>_sdf_vol.npz
        sq_fit_v20/{post_sq_*.ply, final_sq_*.ply, final_sq_model.ply, ...}
        curv_seg_v4/face_labels.npy            # segmentation labels (prerequisite for stage B)
        mesh_mapping_v8/{face_labels_v8.npy, mesh_mapped_v8.ply, report.json}
        full.ply                                # copy of the input mesh
        summary.json

Resume semantics:
    Stage A  sq_fit_v20         done marker: sq_fit_v20/final_sq_model.ply
    Stage B  mesh_mapping_v8    done marker: mesh_mapping_v8/face_labels_v8.npy
    - both done   -> skip the uid entirely
    - only A done -> run stage B only (plus the curv_seg_v4 prerequisite if needed)
    - none done   -> run A, then B
    - summary.json is rewritten at the end of each uid

Usage:
    python scripts/run_batch_sqfit.py --mesh_root <root> --output_root <out>
    python scripts/run_batch_sqfit.py --mesh_root <root> --output_root <out> --start_index 20 --end_index 50
    python scripts/run_batch_sqfit.py --mesh_root <root> --output_root <out> --uids_file list.json
    python scripts/run_batch_sqfit.py --mesh_root <root> --output_root <out> --no_isolate   # same process (debug)
"""
from __future__ import annotations

import argparse
import glob
import json
import os
import shutil
import sys
import time
import traceback
from pathlib import Path

import numpy as np

DEFAULT_SHARD = "00"

V20_PARAMS = dict(
    complex_levels=(16, 32),
    max_batches=50,
    max_cands=5,
    target_coverage=0.9,
    max_iter=200,
    cost_threshold=0.05,
    max_total_sqs=100,
)

V4_PARAMS = dict(
    cut_percentile=90.0,
    cut_abs_min_deg=6.0,
    cut_abs_max_deg=12.0,
    curv_weight=0.0,
    min_faces=60,
    smooth_iters=1,
    explode_scale=0.3,
)

V8_PARAMS = dict(
    vote_mode="count",
    vote_tau=0.025,
    orphan_face_dist=0.04,
    orphan_atom_frac=0.6,
    min_atom_size=3,
)


# ============================================================
# UID discovery
# ============================================================
def discover_uids(mesh_root: str, shard: str, uids_file: str | None = None) -> list[str]:
    if uids_file and os.path.exists(uids_file):
        if uids_file.endswith(".json"):
            data = json.load(open(uids_file))
            if isinstance(data, dict) and "uids" in data:
                return list(data["uids"])
            return list(data)
        return [ln.strip() for ln in open(uids_file) if ln.strip() and not ln.startswith("#")]

    shard_dir = os.path.join(mesh_root, shard)
    if not os.path.isdir(shard_dir):
        raise FileNotFoundError(shard_dir)
    uids = []
    for d in sorted(os.listdir(shard_dir)):
        full = os.path.join(shard_dir, d, "full.ply")
        if os.path.isfile(full):
            uids.append(d)
    return uids


# ============================================================
# Stage completion checks
# ============================================================
def stage_sq_done(uid_out: str) -> bool:
    return os.path.isfile(os.path.join(uid_out, "sq_fit_v20", "final_sq_model.ply"))


def stage_v8_done(uid_out: str) -> bool:
    return os.path.isfile(os.path.join(uid_out, "mesh_mapping_v8", "face_labels_v8.npy"))


def stage_v4_done(uid_out: str) -> bool:
    return os.path.isfile(os.path.join(uid_out, "curv_seg_v4", "face_labels.npy"))


# ============================================================
# Stage A: sq_fit_v20 (adaptive_map_v4 + sdf_vol + sq_fit_v20)
# ============================================================
def _normalize_mesh_o3d(mesh, half=0.5, margin=1e-6):
    import open3d as o3d
    v = np.asarray(mesh.vertices, dtype=np.float64).copy()
    f = np.asarray(mesh.triangles, dtype=np.int32).copy()
    if v.size == 0 or f.size == 0:
        raise ValueError("Empty mesh.")
    center = (v.min(0) + v.max(0)) * 0.5
    ext = float((v.max(0) - v.min(0)).max())
    if ext <= 0:
        raise ValueError(f"Invalid extent: {ext}")
    v = (v - center) * ((half - margin) * 2.0 / ext)
    return o3d.geometry.TriangleMesh(
        o3d.utility.Vector3dVector(v),
        o3d.utility.Vector3iVector(f),
    )


def run_stage_sqfit(uid: str, mesh_path: str, uid_out: str) -> dict:
    import open3d as o3d
    from hitops.adaptive.adaptive_block import build_block_maps_o3d
    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, sdf_volume_stats
    from hitops.sqfit.sq_fit import SDFAdaptiveFitter

    map_dir = os.path.join(uid_out, "adaptive_map_v4")
    sdf_dir = os.path.join(uid_out, "sdf_vol")
    sq_dir  = os.path.join(uid_out, "sq_fit_v20")
    for d in [map_dir, sdf_dir, sq_dir]:
        os.makedirs(d, exist_ok=True)

    stats = {"status": "ok"}
    t_start = time.time()

    mesh = o3d.io.read_triangle_mesh(mesh_path)
    if len(mesh.vertices) == 0:
        raise ValueError("empty mesh")
    stats["n_vertices"] = int(len(mesh.vertices))
    stats["n_triangles"] = int(len(mesh.triangles))
    mesh = _normalize_mesh_o3d(mesh)

    # Build the adaptive block map.
    t0 = time.time()
    prebuilt = {}
    for b in [16, 32, 64]:
        for attempt in range(3):
            try:
                r = build_block_maps_o3d(
                    mesh, B=b, R_max=32,
                    level_scheme="quantile", clip_percentile=99.5,
                )
                assert r["map"].shape == (b, b, b)
                prebuilt[b] = r["map"]
                break
            except (ValueError, AssertionError) as e:
                print(f"    [retry {attempt+1}/3] B={b}: {e}", flush=True)
        else:
            raise RuntimeError(f"build_block_maps_o3d B={b} failed")
    adaptive = build_adaptive_block_map(
        mesh=mesh, B_list=[16, 32, 64], R_max=32, prebuilt_maps=prebuilt,
    )
    stats["time_adaptive"] = round(time.time() - t0, 2)
    save_adaptive_map(adaptive, os.path.join(map_dir, f"{uid}_adaptive_map_v4.npz"))

    # Build the SDF volume.
    t0 = time.time()
    R = adaptive["level_map"].shape[0]
    sdf_vol, truncation = build_sdf_volume(mesh, resolution=R)
    save_sdf_volume(sdf_vol, truncation, os.path.join(sdf_dir, f"{uid}_sdf_vol.npz"))
    stats["time_sdf"] = round(time.time() - t0, 2)
    sdf_volume_stats(sdf_vol, truncation)

    # Fit superquadrics.
    t0 = time.time()
    fitter = SDFAdaptiveFitter(
        adaptive_map=adaptive, sdf_vol=sdf_vol, truncation=truncation,
        complex_levels=tuple(V20_PARAMS["complex_levels"]),
    )
    fitter.run(
        max_batches=V20_PARAMS["max_batches"],
        max_cands=V20_PARAMS["max_cands"],
        target_coverage=V20_PARAMS["target_coverage"],
        max_iter=V20_PARAMS["max_iter"],
        cost_threshold=V20_PARAMS["cost_threshold"],
        max_total_sqs=V20_PARAMS["max_total_sqs"],
    )
    fitter.save_results(sq_dir)
    stats["time_sqfit"] = round(time.time() - t0, 2)
    stats["n_sq"] = len(glob.glob(os.path.join(sq_dir, "post_sq_*.ply")))
    stats["time_total"] = round(time.time() - t_start, 2)
    return stats


# ============================================================
# Stage B prereq: curv_seg_v4 (produces face_labels.npy)
# ============================================================
def run_stage_v4(mesh_path: str, uid_out: str) -> dict:
    from hitops.segment.curvature_seg import segment_pipeline

    out = os.path.join(uid_out, "curv_seg_v4")
    os.makedirs(out, exist_ok=True)
    t0 = time.time()
    segment_pipeline(
        mesh_path=mesh_path,
        out_dir=out,
        **V4_PARAMS,
    )
    return {"status": "ok", "time_total": round(time.time() - t0, 2)}


# ============================================================
# Stage B: mesh_mapper_v8
# ============================================================
def run_stage_v8(mesh_path: str, uid_out: str) -> dict:
    import trimesh
    from hitops.mapping.mesh_mapper import _load_sq_voxel_groups, map_with_atoms, save_results

    fl_path = os.path.join(uid_out, "curv_seg_v4", "face_labels.npy")
    sq_dir  = os.path.join(uid_out, "sq_fit_v20")
    out_dir = os.path.join(uid_out, "mesh_mapping_v8")
    os.makedirs(out_dir, exist_ok=True)

    t0 = time.time()
    mesh = trimesh.load(mesh_path, force="mesh", process=False)
    face_labels_v4 = np.load(fl_path).astype(np.int32)
    sq_names, sq_pts = _load_sq_voxel_groups(sq_dir)
    if not sq_names:
        return {"status": "skipped", "reason": "empty sq_dir"}

    res = map_with_atoms(
        mesh, face_labels_v4, sq_pts,
        vote_mode=V8_PARAMS["vote_mode"],
        vote_tau=V8_PARAMS["vote_tau"],
        orphan_face_dist=V8_PARAMS["orphan_face_dist"],
        orphan_atom_frac=V8_PARAMS["orphan_atom_frac"],
        min_atom_size=V8_PARAMS["min_atom_size"],
    )
    save_results(mesh, res, sq_names, out_dir, save_per_sq=False)
    s = {"status": "ok", "t_total": round(time.time() - t0, 2), **res["stats"]}
    return s


# ============================================================
# Per-UID driver: runs only missing stages
# ============================================================
def process_one(uid: str, mesh_root: str, shard: str, output_root: str) -> dict:
    mesh_src = os.path.join(mesh_root, shard, uid, "full.ply")
    if not os.path.isfile(mesh_src):
        return {"uid": uid, "status": "failed", "error": f"mesh not found: {mesh_src}"}

    uid_out = os.path.join(output_root, uid)
    os.makedirs(uid_out, exist_ok=True)

    # Copy full.ply into the output directory to mirror the expected layout,
    # only when it is missing.
    mesh_dst = os.path.join(uid_out, "full.ply")
    if not os.path.isfile(mesh_dst):
        shutil.copy2(mesh_src, mesh_dst)

    result: dict = {"uid": uid, "mesh_path": mesh_src}
    # Run the stages against the input mesh rather than the copy; the bytes are
    # identical, and the pipeline code expects the input path convention.
    mesh_for_run = mesh_src

    # ---- Stage A: sq_fit_v20 ----
    if stage_sq_done(uid_out):
        n_sq = len(glob.glob(os.path.join(uid_out, "sq_fit_v20", "post_sq_*.ply")))
        result["sq_fit_v20"] = {"status": "skipped", "n_sq": n_sq}
    else:
        print(f"  [sq_fit_v20] running...", flush=True)
        try:
            result["sq_fit_v20"] = run_stage_sqfit(uid, mesh_for_run, uid_out)
        except Exception as e:
            result["sq_fit_v20"] = {"status": "failed", "error": str(e)}
            result["status"] = "failed"
            result["error"] = f"sq_fit_v20: {e}"
            return result

    # ---- Stage B: mesh_mapper_v8 ----
    if stage_v8_done(uid_out):
        result["mesh_mapping_v8"] = {"status": "skipped"}
    else:
        # Prerequisite: the curv_seg_v4 face labels.
        if not stage_v4_done(uid_out):
            print(f"  [curv_seg_v4] running (prereq for v8)...", flush=True)
            try:
                result["curv_seg_v4"] = run_stage_v4(mesh_for_run, uid_out)
            except Exception as e:
                result["curv_seg_v4"] = {"status": "failed", "error": str(e)}
                result["mesh_mapping_v8"] = {"status": "failed", "error": f"v4 prereq: {e}"}
                result["status"] = "failed"
                result["error"] = f"curv_seg_v4: {e}"
                return result
        else:
            result["curv_seg_v4"] = {"status": "skipped"}

        print(f"  [mesh_mapping_v8] running...", flush=True)
        try:
            result["mesh_mapping_v8"] = run_stage_v8(mesh_for_run, uid_out)
        except Exception as e:
            result["mesh_mapping_v8"] = {"status": "failed", "error": str(e)}
            result["status"] = "failed"
            result["error"] = f"mesh_mapping_v8: {e}"
            return result

    result["status"] = "ok"
    # Write the per-uid summary.json.
    _to_py = lambda o: (float(o) if isinstance(o, np.floating) else
                        int(o)   if isinstance(o, np.integer) else
                        o.tolist() if isinstance(o, np.ndarray) else o)
    with open(os.path.join(uid_out, "summary.json"), "w") as f:
        json.dump(result, f, indent=2, default=_to_py)
    return result


# ============================================================
# Subprocess isolation
# ============================================================
def _worker(uid, mesh_root, shard, output_root, q):
    try:
        r = process_one(uid, mesh_root, shard, output_root)
        q.put(r)
    except Exception as e:
        q.put({"uid": uid, "status": "failed", "error": str(e),
               "traceback": traceback.format_exc()})


def process_one_isolated(uid, mesh_root, shard, output_root, timeout_sec=1800):
    import multiprocessing as mp
    ctx = mp.get_context("spawn")
    q = ctx.Queue()
    p = ctx.Process(target=_worker, args=(uid, mesh_root, shard, output_root, q))
    p.start()
    p.join(timeout=timeout_sec)
    if p.is_alive():
        print(f"  [isolate] TIMEOUT {timeout_sec}s -> terminate", flush=True)
        p.terminate(); p.join(5)
        if p.is_alive():
            p.kill(); p.join()
        return {"uid": uid, "status": "failed", "error": f"timeout {timeout_sec}s"}
    res = None
    try:
        if not q.empty():
            res = q.get(timeout=2)
    except Exception:
        res = None
    if p.exitcode != 0:
        sig = f"signal {-p.exitcode}" if p.exitcode < 0 else f"exit {p.exitcode}"
        if res is not None:
            res.setdefault("error", f"worker died ({sig})")
            return res
        return {"uid": uid, "status": "failed", "error": f"worker died ({sig})"}
    if res is None:
        return {"uid": uid, "status": "failed", "error": "no result"}
    return res


# ============================================================
# Main
# ============================================================
def main():
    ap = argparse.ArgumentParser(description="HY3D sq_fit_v20 + mesh_mapper_v8 batch runner")
    ap.add_argument("--mesh_root",   required=True,
                    help="mesh root; meshes at <mesh_root>/<shard>/<uid>/full.ply")
    ap.add_argument("--shard",       default=DEFAULT_SHARD)
    ap.add_argument("--output_root", required=True)
    ap.add_argument("--uids_file",   default=None,
                    help="uid list .txt/.json; omit to scan the shard dir")
    ap.add_argument("--start_index", type=int, default=0,
                    help="index into uid list to start from (inclusive, default 0)")
    ap.add_argument("--end_index",   type=int, default=200,
                    help="index into uid list to stop at (exclusive, default 200)")
    ap.add_argument("--no_isolate",  action="store_true", help="run in same process (debug)")
    ap.add_argument("--timeout_sec", type=int, default=1800)
    args = ap.parse_args()

    uids_file = args.uids_file if args.uids_file else None
    all_uids = discover_uids(args.mesh_root, args.shard, uids_file)
    total = len(all_uids)
    s = max(0, args.start_index)
    e = min(total, args.end_index) if args.end_index > 0 else total
    if s >= e:
        raise SystemExit(f"[batch] empty slice: start_index={s} end_index={e} total={total}")
    uids = all_uids[s:e]
    print(f"[batch] {len(uids)} uid(s) to process  (slice [{s}:{e}) of {total})")
    print(f"[batch] uids_file   = {uids_file or '(scan shard dir)'}")
    print(f"[batch] mesh_root   = {args.mesh_root}")
    print(f"[batch] shard       = {args.shard}")
    print(f"[batch] output_root = {args.output_root}")
    os.makedirs(args.output_root, exist_ok=True)

    summary_path = os.path.join(args.output_root, "batch_summary.json")
    all_stats: list[dict] = []
    if os.path.exists(summary_path):
        try:
            all_stats = json.load(open(summary_path))
            done = {s["uid"] for s in all_stats if s.get("status") in ("ok", "skipped")}
            print(f"[batch] resume: {len(done)} uids already in batch_summary.json", flush=True)
        except Exception:
            all_stats = []

    for i, uid in enumerate(uids):
        print("\n" + "=" * 72, flush=True)
        print(f"[{i+1:3d}/{len(uids)}] {uid}", flush=True)
        print("=" * 72, flush=True)

        # Fast path: both stages are already complete.
        uid_out = os.path.join(args.output_root, uid)
        if stage_sq_done(uid_out) and stage_v8_done(uid_out):
            print("  -> both stages done, skip", flush=True)
            s = {"uid": uid, "status": "skipped", "reason": "both_done"}
        else:
            try:
                if args.no_isolate:
                    s = process_one(uid, args.mesh_root, args.shard, args.output_root)
                else:
                    s = process_one_isolated(
                        uid, args.mesh_root, args.shard, args.output_root,
                        timeout_sec=args.timeout_sec,
                    )
            except Exception as e:
                traceback.print_exc()
                s = {"uid": uid, "status": "failed", "error": str(e)}

        all_stats = [x for x in all_stats if x.get("uid") != uid]
        all_stats.append(s)
        _to_py = lambda o: (float(o) if isinstance(o, np.floating) else
                            int(o)   if isinstance(o, np.integer) else
                            o.tolist() if isinstance(o, np.ndarray) else o)
        with open(summary_path, "w") as f:
            json.dump(all_stats, f, indent=2, default=_to_py)

        st = s.get("status")
        if st == "ok":
            sq = s.get("sq_fit_v20", {})
            v8 = s.get("mesh_mapping_v8", {})
            print(f"  -> OK  sq:{sq.get('status')} v8:{v8.get('status')}", flush=True)
        elif st == "skipped":
            print(f"  -> SKIPPED ({s.get('reason','?')})", flush=True)
        else:
            print(f"  -> FAILED: {s.get('error','?')}", flush=True)

    ok = [s for s in all_stats if s.get("status") == "ok"]
    skipped = [s for s in all_stats if s.get("status") == "skipped"]
    failed = [s for s in all_stats if s.get("status") == "failed"]
    print(f"\n[done] ok={len(ok)} skipped={len(skipped)} failed={len(failed)} "
          f"| summary={summary_path}")


if __name__ == "__main__":
    main()