"""The settings sweep: one object, one encoder setting per run (the peak RAM is the process's). Encodes with every atom
(v5.4 tracked-fit through FCP_API, the low-memory selection of lowmem.py), cuts to the first atoms holding all but
1e-8 of the energy, at most 5 MB, and measures on 1M near-surface points (±2 % diag) and 1M uniform points in the
box + 10 %: wrong answers beyond 0.1–0.5 % of the diagonal, per query ring count NT, with the GPU time per point.
Appends to ../report/sweep.csv.

    env -u PYTHONPATH .venv/bin/python scripts/sweep.py <object> <layers> <N> <samples per ring> <min area frac>
"""
import csv
import os
import resource
import sys
import time
from pathlib import Path

import numpy as np

sys.path.insert(0, str(Path(__file__).parent))
from common import CACHE, REPORT, UP, cuda_sync, fcp, load_mesh, timeit  # noqa: E402
import lowmem  # noqa: E402, F401

import cupy as cp  # noqa: E402
import open3d as o3d  # noqa: E402

from fcpq import FCPQuery  # noqa: E402

PARTS = sys.modules["fcp._parts"]
ENC = sys.modules["fcp.encode"]
M = 64
EPS = (0.1, 0.2, 0.3, 0.4, 0.5)
LEFT = 1e-8  # energy left out by the cut ...
BUDGET = 5_000_000  # ... or at most this many bytes
NTS = (256, 512, 1024)


def cut(d):
    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 points(m):
    sc = o3d.t.geometry.RaycastingScene()
    sc.add_triangles(o3d.core.Tensor(np.asarray(m.vertices, np.float32)), o3d.core.Tensor(np.asarray(m.faces, np.uint32)))
    np.random.seed(0)
    rng = np.random.default_rng(0)
    lo, hi = m.bounds
    diag = np.linalg.norm(hi - lo)
    S, fi = m.sample(1_000_000, return_index=True, seed=0)  # trimesh ignores np.random.seed: seeded here (from 2026-10-30)
    sets = {"near": S + m.face_normals[fi] * rng.uniform(-0.02, 0.02, (len(S), 1)) * diag,
            "uniform": rng.uniform(lo - 0.1 * (hi - lo), hi + 0.1 * (hi - lo), (1_000_000, 3))}
    out = {}
    for k, P in sets.items():
        T = o3d.core.Tensor(P.astype(np.float32))
        out[k] = (P, sc.compute_occupancy(T, nsamples=3).numpy() > 0, np.abs(sc.compute_distance(T).numpy()) / diag * 100)
    return out


def main():
    name, layers, N, ns, ma = sys.argv[1], int(sys.argv[2]), int(sys.argv[3]), int(sys.argv[4]), float(sys.argv[5])
    PARTS.N_SAMPLES, PARTS.MIN_AREA_FRAC = ns, ma
    ENC.MAX_ATOMS = 1 << 26
    p = CACHE / f"sw_{name}_L{layers}_N{N}_s{ns}_ma{ma:g}.fcp"
    m = load_mesh(name)
    enc_s = float("nan")
    if p.exists():
        d = fcp.FCP.from_bytes(p.read_bytes())
    else:
        t0 = time.perf_counter()
        d = fcp.encode(m, atoms=1 << 26, layers=layers, N=N, M=M, up=UP, device="cuda", parts="tracked-fit")
        enc_s = time.perf_counter() - t0
        p.write_bytes(d.to_bytes())
    peak = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss / 1e6
    dc = cut(d)
    sets = points(m)
    rows = []
    for nt in NTS:
        t0 = time.perf_counter()
        q = FCPQuery(dc, nt=nt)
        q.gpu()
        cuda_sync()
        build = time.perf_counter() - t0
        row = dict(object=name, layers=layers, N=N, samples=ns, min_area=ma, M=M, nt=nt, parts=q.P, atoms_all=d.n_atoms,
                   bytes_all=d.nbytes, atoms=dc.n_atoms, bytes=dc.nbytes, encode_s=enc_s, encode_peak_gb=peak, build_s=build)
        for k, (P, gt, dist) in sets.items():
            Pg = cp.asarray(dc.pose.to_canonical(P).astype(np.float32))
            ins = q.inside_gpu(Pg, "table").get()
            row[f"{k}_ns"] = timeit(lambda: q.inside_gpu(Pg, "table"), runs=5, sync=cuda_sync) / len(P) * 1e9
            wrong = ins != gt
            row[f"{k}_wrong_pct"] = 100 * wrong.mean()
            for e in EPS:
                row[f"{k}_n{e}"] = int((wrong & (dist > e)).sum())
            row[f"{k}_max"] = float(dist[wrong].max()) if wrong.any() else 0.0
            del Pg
        rows.append(row)
        print(f"{name:9s} L{layers:4d} N{N:3d} s{ns:4d} ma{ma:<6g} NT{nt:5d}: {dc.bytes if False else dc.nbytes / 1e6:5.2f} MB "
              f"parts {q.P:4d} | near >" + " ".join(f"{e}:{row[f'near_n{e}']}" for e in EPS)
              + f" max {row['near_max']:.3f} | uniform >0.1:{row['uniform_n0.1']} | GPU {row['near_ns']:.1f}/{row['uniform_ns']:.1f} ns"
              f" | encode {enc_s:.0f} s {peak:.1f} GB", flush=True)
        del q
        cp.get_default_memory_pool().free_all_blocks()
    out = REPORT / "sweep.csv"
    new = not out.exists()
    with out.open("a", newline="") as f:
        w = csv.DictWriter(f, list(rows[0]))
        if new:
            w.writeheader()
        w.writerows(rows)


if __name__ == "__main__":
    main()