"""Build the object catalog from NVIDIA SimReady Central assets on the Hugging Face Hub.

For every asset in the dataset that a robot arm could handle (mass <= MAX_KG), writes
  objects/<id>.glb      the decimated visual geometry, metres, z up
and objects/catalog.js, the list app.js searches: name, class, description, mass,
bounding-box size and a link to the asset's thumbnail and page on the Hub.

Mass, semantic class and description come from each asset's payloads/base.usda; the size
and the geometry come from composing the asset's root USD with usd-core.

  uv pip install usd-core huggingface_hub trimesh fast-simplification numpy
  python tools/build_objects.py
  for f in objects/*.glb; do npx @gltf-transform/cli optimize $f $f --compress meshopt --simplify false \\
      --texture-compress false; done
"""
import json
import re
import sys
from collections import Counter, defaultdict
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path

import fast_simplification
import numpy as np
import trimesh
from huggingface_hub import HfApi, hf_hub_download
from pxr import Usd, UsdGeom

REPO = "nvidia/simready-assets"
MAX_KG = 35.0                 # heaviest object the heaviest catalog robot can lift
MAX_GEOMETRY = 40e6           # skip huge scans
FACES = 2500
OUT = Path(__file__).resolve().parent.parent / "objects"
HUB = f"https://huggingface.co/datasets/{REPO}"


def assets():
    by = defaultdict(dict)
    for f in HfApi().list_repo_tree(REPO, repo_type="dataset", recursive=True):
        if hasattr(f, "size"):
            by[f.path.split("/")[0]][f.path] = f.size
    out = []
    for name, files in sorted(by.items()):
        roots = [p for p in files if p.count("/") == 2 and p.split("/")[1] == "simready_usd" and p.endswith(".usd")]
        base = [p for p in files if p.endswith("payloads/base.usda")]
        if not roots or not base:
            continue
        root = roots[0]
        thumb = root.replace("simready_usd/", "simready_usd/.thumbs/256x256/") + "_transparent.png"
        geo = sum(s for p, s in files.items() if p.endswith(".usd") and "/payloads/" in p)
        out.append({"id": name, "root": root, "base": base[0], "thumb": thumb if thumb in files else None, "geo": geo,
                    "files": [p for p in files if "/simready_usd/" in p and "/Textures/" not in p and "/source_assets/" not in p
                              and "/.thumbs/" not in p and not p.endswith((".png", ".jpg", ".mdl"))]})
    return out


def metadata(a):
    text = Path(hf_hub_download(REPO, a["base"], repo_type="dataset")).read_text(errors="ignore")
    mass = sum(float(m) for m in re.findall(r"pxr:usd:physics_mass = ([0-9.eE+-]+)", text))
    if not mass:
        m = re.search(r"physics:mass = ([0-9.eE+-]+)", text)
        mass = float(m.group(1)) if m else None
    cls = re.search(r'semantics:labels:wikidata_class = \["([^"]+)"', text)
    doc = re.search(r'omni:simready:documentation = "([^"]*)"', text)
    kind = re.search(r'string type = "([^"]+)"', text)
    return {"mass": mass, "cls": cls.group(1) if cls else None, "doc": doc.group(1) if doc else "", "type": kind.group(1) if kind else None}


def pretty(name):
    words = re.sub(r"_[a-z]\d\d$", "", name).replace("_", " ")
    words = re.sub(r"([a-z])([A-Z])", r"\1 \2", words).lower()
    return words[:1].upper() + words[1:]


def geometry(a):
    for p in a["files"]:
        hf_hub_download(REPO, p, repo_type="dataset")
    root = hf_hub_download(REPO, a["root"], repo_type="dataset")
    stage = Usd.Stage.Open(root)
    up = UsdGeom.GetStageUpAxis(stage)
    scale = UsdGeom.GetStageMetersPerUnit(stage) or 1.0
    cache = UsdGeom.XformCache()
    parts = []
    for prim in Usd.PrimRange.Stage(stage, Usd.TraverseInstanceProxies()):
        if not prim.IsA(UsdGeom.Mesh) or UsdGeom.Imageable(prim).ComputePurpose() not in (UsdGeom.Tokens.default_, UsdGeom.Tokens.render):
            continue
        mesh = UsdGeom.Mesh(prim)
        pts, counts, idx = mesh.GetPointsAttr().Get(), mesh.GetFaceVertexCountsAttr().Get(), mesh.GetFaceVertexIndicesAttr().Get()
        if not pts or not counts:
            continue
        pts = np.asarray(pts, float)
        M = np.asarray(cache.GetLocalToWorldTransform(prim), float)        # row-vector convention
        pts = (np.c_[pts, np.ones(len(pts))] @ M)[:, :3]
        faces, k = [], 0
        idx = np.asarray(idx)
        for c in counts:
            faces += [[idx[k], idx[k + i], idx[k + i + 1]] for i in range(1, c - 1)]
            k += c
        color = mesh.GetDisplayColorAttr().Get()
        parts.append((pts, np.asarray(faces, int), tuple(color[0]) if color else None))
    if not parts:
        return None
    meshes = []
    total = sum(len(f) for _, f, _ in parts)
    for pts, faces, color in parts:
        target = max(60, int(FACES * len(faces) / total))
        if len(faces) > target:
            pts, faces = fast_simplification.simplify(pts.astype(np.float32), faces.astype(np.int32), 1 - target / len(faces))
        m = trimesh.Trimesh(pts * scale, faces, process=True)
        rgba = (np.r_[color, 1] if color else np.array([0.78, 0.74, 0.68, 1])) * 255
        m.visual = trimesh.visual.TextureVisuals(material=trimesh.visual.material.PBRMaterial(baseColorFactor=rgba.astype(np.uint8), roughnessFactor=0.7))
        meshes.append(m)
    scene = trimesh.Scene(meshes)
    if up == "Y":
        scene.apply_transform(trimesh.transformations.rotation_matrix(np.pi / 2, [1, 0, 0]))
    lo, hi = scene.bounds
    scene.apply_translation([-(lo[0] + hi[0]) / 2, -(lo[1] + hi[1]) / 2, -lo[2]])    # centred, resting on z = 0
    return scene


def build(a):
    try:
        meta = metadata(a)
    except Exception as e:
        return a["id"], None, f"metadata: {e!r}"
    if meta["mass"] is None or meta["mass"] > MAX_KG or a["geo"] > MAX_GEOMETRY:
        return a["id"], None, "skipped"
    try:
        scene = geometry(a)
    except Exception as e:
        return a["id"], None, f"geometry: {e!r}"
    if scene is None:
        return a["id"], None, "no mesh"
    dims = scene.extents
    if max(dims) > 3:                      # mis-scaled asset
        return a["id"], None, f"size {dims}"
    (OUT / f"{a['id']}.glb").write_bytes(trimesh.exchange.gltf.export_glb(scene))
    variant = re.search(r"_([a-z]\d\d)$", a["id"])
    name = (meta["cls"][:1].upper() + meta["cls"][1:]) if meta["cls"] else pretty(a["id"])
    item = {"id": a["id"], "name": name + (f" ({variant.group(1)})" if variant else ""), "cls": meta["cls"], "doc": meta["doc"][:160], "kg": round(meta["mass"], 4),
            "size": [round(float(x), 4) for x in dims], "thumb": a["thumb"]}
    return a["id"], item, "ok"


def dedupe(items):
    """Several folders share a class (three "Bottle (a01)"): name those after their folder instead."""
    count = Counter(i["name"] for i in items)
    for i in items:
        if count[i["name"]] > 1:
            variant = re.search(r"_([a-z]\d\d)$", i["id"])
            i["name"] = pretty(i["id"]) + (f" ({variant.group(1)})" if variant else "")


def main():
    OUT.mkdir(exist_ok=True)
    todo = assets()
    if len(sys.argv) > 1:
        todo = [a for a in todo if a["id"] in sys.argv[1:]]
    print(f"{len(todo)} assets", flush=True)
    items, stats = [], defaultdict(int)
    with ThreadPoolExecutor(12) as pool:
        for i, (aid, item, status) in enumerate(pool.map(build, todo)):
            stats[status.split(":")[0]] += 1
            if item:
                items.append(item)
            elif status != "skipped":
                print(aid, status, flush=True)
            if i % 100 == 0:
                print(i, dict(stats), flush=True)
    dedupe(items)
    items.sort(key=lambda x: x["name"])
    (OUT / "catalog.js").write_text(f"// Generated by tools/build_objects.py from {HUB} (CC-BY-4.0).\n"
                                    f"const OBJECTS_SOURCE = {json.dumps(HUB)};\nconst OBJECTS = "
                                    + json.dumps(items, separators=(",", ":")) + ";\n")
    print(len(items), "objects", dict(stats))


if __name__ == "__main__":
    main()