"""Objects, descriptors (cached), timing. Run every script with the study's venv:

    cd ~/code/data/meshes/fcp_2910_collision && env -u PYTHONPATH .venv/bin/python scripts/<script>.py
"""
import os
import sys
import time
from pathlib import Path

import numpy as np
import trimesh

sys.path.insert(0, os.environ.get("FCP_API", "/home/joao/code/fcp_devel/api"))
import fcp  # noqa: E402

ROOT = Path(__file__).resolve().parent.parent
REPORT = ROOT / "report"
IMG = REPORT / "img"
CACHE = ROOT / "cache"
D = Path("/home/joao/code/data/meshes")
OBJ = {  # watertight meshes of the earlier studies (fcp_2810_partition)
    "mug": D / "data/meshes_clean/gso__Room_Essentials_Mug_White_Yellow.ply",
    "bowl": D / "data/meshes_clean/gso__Threshold_Bead_Cereal_Bowl_White.ply",
    "dino": D / "data/meshes_clean/gso__Schleich_Therizinosaurus_ln9cruulPqc.ply",
    "shoe": D / "data/meshes_clean/gso__Timberland_Mens_Classic_2Eye_Boat_Shoe.ply",
    "horse": D / "data/meshes_clean/gso__Breyer_Horse_Of_The_Year_2015.ply",
    "dress": Path("/home/joao/code/fousd-3d/outputs/heavy/MANNEQUIN_SUMMER_DRESS_-_ORANGE_envelope_solid_target.ply"),
}
COMPLEX = {  # concave / interleaved / perforated: real (the manifest's highest genus) and synthetic (make_shapes.py)
    "pan_holes": D / "data/meshes_clean/objectfolder__73.ply",  # perforated pan, genus 334
    "knot": ROOT / "meshes/knot.ply",
    "spring": ROOT / "meshes/spring.ply",
    "sponge": ROOT / "meshes/sponge.ply",
    "gear": ROOT / "meshes/gear.ply",
}
_G = D / "data/meshes_clean/gso__"
HOUSEHOLD = {  # 10 household objects with clean watertight meshes (Google Scanned Objects), the CPU study
    "mug": _G.with_name("gso__Room_Essentials_Mug_White_Yellow.ply"),
    "mug_ace": _G.with_name("gso__ACE_Coffee_Mug_Kristen_16_oz_cup.ply"),
    "bowl": _G.with_name("gso__Threshold_Bead_Cereal_Bowl_White.ply"),
    "bowl_serving": _G.with_name("gso__Threshold_Porcelain_Serving_Bowl_Coupe_White.ply"),
    "bowl_dog": _G.with_name("gso__Top_Paw_Dog_Bowl_Blue_Paw_Bone_Ceramic_25_fl_oz_total.ply"),
    "flower_pot": _G.with_name("gso__Central_Garden_Flower_Pot_Goo_425.ply"),
    "shoe": _G.with_name("gso__Timberland_Mens_Classic_2Eye_Boat_Shoe.ply"),
    "sneaker": _G.with_name("gso__Reebok_ZIGSTORM.ply"),
    "jogger": _G.with_name("gso__Reebok_CLASSIC_JOGGER.ply"),
    "boot": _G.with_name("gso__Tory_Burch_Kiernan_Riding_Boot.ply"),
}
ALL = {**HOUSEHOLD, **OBJ, **COMPLEX,  # dropped from the study (v2, 2026-10-30), kept loadable for the earlier scripts
       "elephant": D / "data/meshes_clean/gso__Elephant.ply", "dog": D / "data/meshes_clean/gso__Dog.ply"}
SET = os.environ.get("FCP_SET", "standard")  # "complex": the objects above; outputs get the suffix TAG
TAG = ("" if SET == "standard" else f"_{SET}") + os.environ.get("FCP_TAGX", "")  # FCP_TAGX: e.g. "_L512" for a variant run
if SET == "complex":
    OBJ = COMPLEX
if os.environ.get("FCP_OBJECTS"):
    OBJ = {k: v for k, v in OBJ.items() if k in os.environ["FCP_OBJECTS"].split(",")}
UP = (0, 0, 1)
# v2 (2026-10-30): one setting for every object, chosen by scripts/sweep.py (report section "Settings"); encoded with
# the package at FCP_API (v5.4, parts="tracked-fit"), every atom, then cut to all but CUT_LEFT of the energy, at most
# CUT_BYTES. The v1 runs (1024 layers, N 64, 12288 atoms, NT 256) are archived in report/v1/.
LAYERS = int(os.environ.get("FCP_LAYERS", 1024))
N = int(os.environ.get("FCP_N", 128))
M = int(os.environ.get("FCP_M", 64))
SAMPLES = int(os.environ.get("FCP_SAMPLES", 2048))  # arc-length samples per ring (the package's _parts.N_SAMPLES: 512)
MIN_AREA = float(os.environ.get("FCP_MINAREA", 0))  # the package's _parts.MIN_AREA_FRAC (0.002: small islands dropped)
NT = int(os.environ.get("FCP_NT", 512))  # ring vertices in the query (y-bins: NT / 4)
RULE = os.environ.get("FCP_RULE", "nearest")  # the query's z rule: "nearest" layer centre (v2), or "table" (centres agree, else blend)
CUT_LEFT, CUT_BYTES = 1e-8, 5_000_000
ATOMS = None  # v1: an atom budget; v2 cuts by energy and size


def load_mesh(name):
    m = trimesh.load(ALL[name], force="mesh")
    return m


def cut(d, left=CUT_LEFT, budget=CUT_BYTES):
    """The first atoms holding all but `left` of the energy (the atoms are orthonormal: energy = sum of squared values),
    at most `budget` bytes."""
    e = np.cumsum(np.asarray(d.atoms[1], np.float64) ** 2)
    k = min(int(np.searchsorted(e, (1 - left) * e[-1])) + 1, d.n_atoms)
    if k < d.n_atoms:
        d = d.prefix(k)
    if d.nbytes <= budget:
        return d
    lo, hi = 0, d.n_atoms
    while hi - lo > 1:
        mid = (lo + hi) // 2
        lo, hi = (mid, hi) if d.prefix(mid).nbytes <= budget else (lo, mid)
    return d.prefix(lo)


def descriptor(name, layers=None, n=None, samples=None, min_area=None):
    """The object's descriptor at the study's setting: encoded once on CUDA with every atom (cached as bytes, the
    same files as sweep.py), returned cut (`cut`)."""
    layers, n = layers or LAYERS, n or N
    samples = samples or SAMPLES
    min_area = MIN_AREA if min_area is None else min_area
    CACHE.mkdir(exist_ok=True)
    p = CACHE / f"{os.environ.get('FCP_VTAG', '')}sw_{name}_L{layers}_N{n}_s{samples}_ma{min_area:g}.fcp"  # FCP_VTAG: another package version
    if not p.exists():
        import lowmem  # noqa: F401  (the encoder's selection without its padded basis array)

        parts, enc = sys.modules["fcp._parts"], sys.modules["fcp.encode"]
        parts.N_SAMPLES, parts.MIN_AREA_FRAC, enc.MAX_ATOMS = samples, min_area, 1 << 26
        d = fcp.encode(load_mesh(name), atoms=1 << 26, layers=layers, N=n, M=M, up=UP, device="cuda", parts="tracked-fit")
        p.write_bytes(d.to_bytes())
    return cut(fcp.FCP.from_bytes(p.read_bytes()))


def timeit(fn, runs=5, sync=None, warm=1):
    """Median seconds of fn() over runs (after warm-up); sync() called before stopping the clock (CUDA)."""
    for _ in range(warm):
        fn()
        if sync:
            sync()
    ts = []
    for _ in range(runs):
        t0 = time.perf_counter()
        fn()
        if sync:
            sync()
        ts.append(time.perf_counter() - t0)
    return float(np.median(ts))


def cuda_sync():
    import cupy as cp

    cp.cuda.runtime.deviceSynchronize()