"""Every working approach on every cloud: ../report/ablation.csv (resumes: rows already there are skipped).

Clouds: ../pointclouds/<name>_<tag>_<10000|2048>.ply, every tag of make_clouds.py (clean, mesh-noise, per-point noise,
variable density). Configurations:
  package          fcp.encode (surf slicer, tracked-fit, fit_tol="auto")
  pointtest        package slicer, cuts by the two-sided point test (pointfit.py)
  A                the slicer's rings refitted to their points (ringfit.py)
  A+pointtest
  hybrid           slicer topology repaired from the neighbouring layers, all fitted (hybrid.py)
  hybrid+pointtest
  hybrid-gated     the same, a carried set tried only where the slicer's layer is suspicious
  hybrid-gated+pointtest
  thinwall         thin walls (a parallel face within 1.5 spacings) added to the slicer's field as a band (thinwall.py)
  thinwall+pointtest
  eps              one scale eps = max(spacing, 4 sigma) for the slicer and the minimum loop, point test (eps.py);
                   sigma = the cloud's per-point noise, given (0 for clean, mesh-noise and density clouds)
  gwn              the GWN slicer (../../fcp_2810_gwn/scripts/gwn.py, defaults, sigma given): inside / outside from the
                   mollified generalized winding number; package cuts
  gwn+pointtest    the same, cuts by the point test
  gwn-tree         gwn with its two sums as tree codes (fastgwn.py; GPU): must match gwn
  local            the same idea per point: eps_i = max(local spacing_i, 4 sigma) in the slicer's kernels (surf_local.py,
                   CPU), point test
Reference: the mesh the cloud was sampled from (the noisy mesh for noise*, else the clean mesh), mesh route. Errors against the clean mesh (the truth).

    env -u PYTHONPATH PYTHONWARNINGS=ignore .venv/bin/python ~/code/data/meshes/fcp_2710_tests/scripts/ablation.py
"""
import csv
import sys
import time
from pathlib import Path

import numpy as np
import trimesh

sys.path.insert(0, str(Path(__file__).parent))
sys.path.insert(0, "/home/joao/code/fcp_devel/api")
import fcp  # noqa: E402
import hybrid as hy  # noqa: E402
import pointfit  # noqa: E402
import ringfit as rf  # noqa: E402
import thinwall as tw  # noqa: E402
import eps as epsmod  # noqa: E402
import surf_local  # noqa: E402
sys.path.insert(0, "/home/joao/code/data/meshes/fcp_2810_gwn/scripts")  # the GWN slicer (another agent's, used as is)
import gwn  # noqa: E402
from make_clouds import PN  # noqa: E402
from make_noisy import SIGMAS, label  # noqa: E402
from parts_report import OBJ, REPORT, ROOT, UP, stats  # noqa: E402

OUT = REPORT / "ablation.csv"
from make_clouds import TAGS  # noqa: E402  every cloud type
SIZES = (10000, 2048)
CONFIGS = {
    "package": lambda P, up: fcp.encode(P, device="cuda", up=up),
    "pointtest": lambda P, up: pointfit.encode_pointtest(P, up=up),
    "A": lambda P, up: rf.encode_ringfit(P, up=up),
    "A+pointtest": lambda P, up: rf.encode_ringfit(P, up=up, cuts="pointtest"),
    "hybrid": lambda P, up: hy.encode_hybrid(P, up=up),
    "hybrid+pointtest": lambda P, up: hy.encode_hybrid(P, up=up, cuts="pointtest"),
    "hybrid-gated": lambda P, up: hy.encode_hybrid(P, up=up, gate=True),
    "hybrid-gated+pointtest": lambda P, up: hy.encode_hybrid(P, up=up, cuts="pointtest", gate=True),
    "thinwall": lambda P, up: tw.encode_thinwall(P, up=up),
    "thinwall+pointtest": lambda P, up: tw.encode_thinwall(P, up=up, cuts="pointtest"),
    "eps": lambda P, up: epsmod.encode_eps(P, up=up, sigma=CTX["sigma"]),
    "local": lambda P, up: surf_local.encode_local(P, up=up, sigma=CTX["sigma"]),
    "gwn": lambda P, up: gwn.encode_gwn(P, up=up, sigma=CTX["sigma"]),
    "gwn+pointtest": lambda P, up: gwn.encode_gwn(P, up=up, sigma=CTX["sigma"], cuts="pointtest"),
    "gwn-tree": lambda P, up: _gwn_tree().encode_gwn(P, up=up, sigma=CTX["sigma"]),
}
_GT = {}
# gwn-tree is checked against gwn on a sample, not on every cloud: per object clean 10k and one other cloud type
_OTHER = ["pn0p6_10000", "pn1p2_10000", "dgrad_10000", "dstripe_10000", "clean_2048", "pn0p6_2048", "noise1_10000",
          "dgrad_2048", "pn1p2_2048", "noise0p6_10000"]
SAMPLED = {"gwn-tree": None}  # filled in main: (object, cloud tag, points) to run


def _sample(objects):
    return {(o, *c.rsplit("_", 1)) for i, o in enumerate(objects) for c in ("clean_10000", _OTHER[i % len(_OTHER)])}


def _gwn_tree():
    """A copy of gwn.py with the tree sums (fastgwn.py) on the GPU."""
    if not _GT:
        import gwn_speed

        _GT["m"] = gwn_speed.gpu_copy(True)
    return _GT["m"]
CTX = {"sigma": 0.0}  # the cloud's per-point noise (world units), given to "eps" (an encoder parameter: known noise)
FIELDS = ["object", "cloud", "points", "config", "parts", "holes", "short", "ref_parts", "mean", "p95", "empty", "seconds"]


def main():
    done = set()
    if OUT.exists():
        with OUT.open() as f:
            done = {(r["object"], r["cloud"], r["points"], r["config"]) for r in csv.DictReader(f)}
    new = not OUT.exists()
    with OUT.open("a", newline="") as f:
        w = csv.DictWriter(f, FIELDS)
        if new:
            w.writeheader()
        for name in OBJ:
            g = trimesh.load(OBJ[name], force="mesh")
            ref = np.asarray(g.sample(50_000, seed=1))
            up = UP.get(name, (0, 0, 1))
            for tag in TAGS:
                src = trimesh.load(ROOT / "meshes" / f"{name}_{tag}.ply", force="mesh") if tag.startswith("noise") else g
                ref_parts = None
                CTX["sigma"] = ({f"pn{label(v)}": v for v in PN}.get(tag, 0.0)) * float(np.linalg.norm(g.extents))
                for n in SIZES:
                    P = np.asarray(trimesh.load(ROOT / "pointclouds" / f"{name}_{tag}_{n}.ply").vertices)
                    for cfg, f_enc in CONFIGS.items():
                        if (name, tag, str(n), cfg) in done:
                            continue
                        if cfg in SAMPLED and (name, tag, str(n)) not in _sample(list(OBJ)):
                            continue
                        if ref_parts is None:
                            ref_parts = len(fcp.encode(src, device="cuda", up=up).parts)
                        t = time.perf_counter()
                        try:
                            d = f_enc(P, up)
                            sec = time.perf_counter() - t
                            m = fcp.decode(d, device="cuda")
                            mean, p95 = stats(d, ref, m) if m is not None else (float("nan"), float("nan"))
                            spans = np.array([h1 - h0 + 1 for _, h0, h1 in d.parts]) if d.parts else np.zeros(0)
                            row = dict(parts=len(d.parts), holes=sum(s < 0 for s, _, _ in d.parts), short=int((spans <= 3).sum()),
                                       mean=mean, p95=p95, empty=int(m is None), seconds=round(sec, 2))
                        except Exception as e:  # a prototype that fails on a cloud is a result too
                            row = dict(parts=-1, holes=-1, short=-1, mean=float("nan"), p95=float("nan"), empty=1,
                                       seconds=round(time.perf_counter() - t, 2))
                            print("FAILED", name, tag, n, cfg, repr(e), flush=True)
                        w.writerow(dict(object=name, cloud=tag, points=n, config=cfg, ref_parts=ref_parts, **row))
                        f.flush()
                    print(name, tag, n, flush=True)


if __name__ == "__main__":
    main()