"""Point-map faithfulness metrics for 3D completion (buffer-F1 suite).

Measures how faithfully a generated shape follows the *observed* evidence
(a partial point cloud / point map), NOT overall completion quality.

Everything operates on a shared canonical voxel grid (default 64^3 over
[-0.5, 0.5]^3, the TRELLIS convention). No registration is performed:
conditioned methods are born aligned; align unconditioned baselines before
calling.

Core metrics
  buffer-F1@r   voxel F1 with tolerance radius r (r=0,1,2 ~ user's 1/3/5 cubes)
  distances     mean / P95 / max (visible Hausdorff) from EDT, both directions
  FSV           free-space violation: predicted voxels in ray-carved empty space
  rim bands     F1@1 split by distance to the visibility border (near/mid/deep)

Usage:
  from metrics.faithfulness import evaluate_faithfulness, voxelize_points
  rep = evaluate_faithfulness(gt_vis_occ, pred_occ, free_occ)  # 64^3 bool grids

Self-test (no data needed):  python metrics/faithfulness.py
"""

from __future__ import annotations

import json
from dataclasses import dataclass, field, asdict

import numpy as np
from scipy import ndimage


# ---------------------------------------------------------------- grids

def canonicalize(points: np.ndarray, center: np.ndarray | None = None,
                 scale: float | None = None):
    """Map points (M,3) into [-0.5, 0.5]^3.

    If center/scale are None they are derived from the points' bbox
    ("realistic" mode). Pass the GT mesh's bbox center/extent for
    "oracle" mode so partial and GT share one frame.
    Returns (normalized points, center, scale).
    """
    pts = np.asarray(points, dtype=np.float64)
    if center is None:
        center = (pts.min(0) + pts.max(0)) / 2.0
    if scale is None:
        scale = (pts.max(0) - pts.min(0)).max()
    return (pts - center) / scale, center, scale


def voxelize_points(points: np.ndarray, n: int = 64) -> np.ndarray:
    """Points in [-0.5, 0.5]^3 -> (n,n,n) bool occupancy.

    Both ±0.5 boundaries are INCLUSIVE (a bbox-tight mesh has whole faces
    exactly at ±0.5; flooring them out would silently delete them).
    """
    pts = np.asarray(points, dtype=np.float64)
    atol = 1e-6
    inside = ((pts >= -0.5 - atol) & (pts <= 0.5 + atol)).all(axis=1)
    pts = np.clip(pts[inside], -0.5, np.nextafter(0.5, 0.0))
    idx = np.floor((pts + 0.5) * n).astype(np.int64)
    idx = np.minimum(idx, n - 1)
    occ = np.zeros((n, n, n), dtype=bool)
    occ[idx[:, 0], idx[:, 1], idx[:, 2]] = True
    return occ


def voxelize_mesh_surface(mesh, n: int = 64, samples: int = 500_000,
                          center=None, scale=None) -> np.ndarray:
    """Sample a trimesh surface and voxelize (surface shell, no solid fill —
    matches TRELLIS's triangle-only voxelization)."""
    pts = mesh.sample(samples)
    pts, _, _ = canonicalize(pts, center, scale)
    return voxelize_points(pts, n)


def carve_free_space(vis_occ: np.ndarray, cam_pos: np.ndarray,
                     margin: int = 2, step: float = 0.5) -> np.ndarray:
    """Mark voxels proven empty by the scan: cells traversed by camera rays
    *before* reaching an observed surface voxel (with a safety margin).

    vis_occ : (n,n,n) bool, observed surface voxels (canonical grid)
    cam_pos : (3,) camera position in canonical coords
    margin  : stop this many voxels before the surface hit
    Returns (n,n,n) bool free-space mask (never overlaps vis_occ).
    """
    n = vis_occ.shape[0]
    targets = np.argwhere(vis_occ) + 0.5                      # voxel centers
    cam = (np.asarray(cam_pos, np.float64) + 0.5) * n         # to grid coords
    d = targets - cam
    dist = np.linalg.norm(d, axis=1)
    free = np.zeros_like(vis_occ)
    n_steps = int(np.ceil(dist.max() / step)) if len(dist) else 0
    for i in range(n_steps):
        t = i * step
        live = dist - margin > t                              # not yet near hit
        if not live.any():
            break
        p = cam + d[live] / dist[live, None] * t
        ijk = np.floor(p).astype(np.int64)
        ok = ((ijk >= 0) & (ijk < n)).all(axis=1)
        ijk = ijk[ok]
        free[ijk[:, 0], ijk[:, 1], ijk[:, 2]] = True
    # grazing rays can carve a 1-voxel skin hugging the surface (and even
    # just behind it); a plausible reconstruction sitting 1 voxel inside GT
    # would then trip FSV. Keep free space Chebyshev->=1 away from surface.
    free &= ~ndimage.maximum_filter(vis_occ, size=3)
    return free


# ---------------------------------------------------------------- metrics

@dataclass
class FaithfulnessReport:
    f1: dict = field(default_factory=dict)         # r -> {precision, recall, f1}
    dist_pred_to_gt: dict = field(default_factory=dict)   # mean/p95/max (voxels)
    dist_gt_to_pred: dict = field(default_factory=dict)   # mean/p95/max
    fsv_rate: float = float("nan")                 # fraction of pred in free space
    fsv_max_depth: float = float("nan")            # deepest violation (voxels)
    rim_bands: dict = field(default_factory=dict)  # band -> {recall@1, n_voxels}
    n_gt_vis: int = 0
    n_pred: int = 0
    n_pred_observed: int = 0                       # pred voxels inside observed region

    def to_json(self) -> str:
        return json.dumps(asdict(self), indent=2, default=float)


def _dist_stats(d: np.ndarray) -> dict:
    if d.size == 0:
        return {"mean": float("nan"), "p95": float("nan"), "max": float("nan")}
    return {"mean": float(d.mean()), "p95": float(np.percentile(d, 95)),
            "max": float(d.max())}


def evaluate_faithfulness(gt_vis: np.ndarray, pred: np.ndarray,
                          free: np.ndarray | None = None,
                          gt_full: np.ndarray | None = None,
                          radii=(0, 1, 2),
                          rim_band_edges=(2, 4)) -> FaithfulnessReport:
    """All faithfulness metrics between observed voxels and a prediction.

    gt_vis : (n,n,n) bool — observed (visible) surface voxels, the reference
    pred   : (n,n,n) bool — generated occupancy in the SAME canonical grid
    free   : optional (n,n,n) bool — ray-carved proven-empty voxels.
             Enables FSV and restricts precision to the observed region
             (vis ∪ free); without it precision falls back to all pred voxels
             (pessimistic: the invented back half counts against precision).
    gt_full: optional (n,n,n) bool — FULL GT surface voxels. Enables rim
             bands: the border is the seam where the visible patch meets the
             hidden part of the surface (gt_full & ~gt_vis). Without it the
             border falls back to adjacency-to-unknown, which is only sound
             when `free` covers the whole camera frustum.
    radii  : tolerance radii in voxels (0,1,2 == 1^3, 3^3, 5^3 buffers)
    rim_band_edges : distances (voxels) splitting near/mid/deep from the
             visibility border.
    """
    rep = FaithfulnessReport()
    rep.n_gt_vis = int(gt_vis.sum())
    rep.n_pred = int(pred.sum())
    if rep.n_gt_vis == 0 or rep.n_pred == 0:
        return rep

    # EDT: distance of every voxel to nearest occupied voxel of the other set
    d_to_gt = ndimage.distance_transform_edt(~gt_vis)
    d_to_pred = ndimage.distance_transform_edt(~pred)

    # precision only over pred voxels where the scan asserts anything
    if free is not None:
        observed = gt_vis | free
        pred_obs = pred & observed
    else:
        pred_obs = pred
    rep.n_pred_observed = int(pred_obs.sum())

    d_of_preds = d_to_gt[pred_obs]        # per observed-region pred voxel
    d_of_gts = d_to_pred[gt_vis]          # per GT visible voxel
    rep.dist_pred_to_gt = _dist_stats(d_of_preds)
    rep.dist_gt_to_pred = _dist_stats(d_of_gts)   # max == visible Hausdorff

    for r in radii:
        # +1e-6: EDT of a diagonal neighbor is sqrt(2) etc.; buffer n^3 means
        # Chebyshev, so use the box distance via maximum filter instead of EDT
        # for exact n^3 semantics at small r.
        if r == 0:
            hit_p = gt_vis[pred_obs] if pred_obs.any() else np.array([])
            prec = float(hit_p.mean()) if hit_p.size else 0.0
            rec = float(pred[gt_vis].mean())
        else:
            gt_dil = ndimage.maximum_filter(gt_vis, size=2 * r + 1)
            pr_dil = ndimage.maximum_filter(pred, size=2 * r + 1)
            prec = float(gt_dil[pred_obs].mean()) if pred_obs.any() else 0.0
            rec = float(pr_dil[gt_vis].mean())
        f1 = 2 * prec * rec / (prec + rec) if prec + rec > 0 else 0.0
        rep.f1[int(r)] = {"precision": prec, "recall": rec, "f1": f1}

    # free-space violations
    if free is not None:
        viol = pred & free
        denom = int((pred & (gt_vis | free)).sum())
        rep.fsv_rate = float(viol.sum() / denom) if denom else 0.0
        if viol.any():
            # depth = how far inside proven-empty space, in voxels
            d_out = ndimage.distance_transform_edt(free)      # dist to non-free
            rep.fsv_max_depth = float(d_out[viol].max())
        else:
            rep.fsv_max_depth = 0.0

    # rim bands: border = the seam between the visible patch and the rest
    # of the surface. Preferred: gt_full gives the hidden surface directly.
    # Fallback (no gt_full): visible voxels adjacent to un-carved space —
    # only meaningful if `free` covers the whole frustum.
    rim = None
    if gt_full is not None:
        hidden = gt_full & ~gt_vis
        rim = gt_vis & ndimage.maximum_filter(hidden, size=3)
    elif free is not None:
        unknown = ~(gt_vis | free)
        rim = gt_vis & ndimage.maximum_filter(unknown, size=3)
    if rim is not None and rim.any():
        d_rim = ndimage.distance_transform_edt(~rim)
        dvals = d_rim[gt_vis]
        pr_dil1 = ndimage.maximum_filter(pred, size=3)
        hit1 = pr_dil1[gt_vis]
        lo, hi = rim_band_edges
        for name, m in (("near", dvals <= lo),
                        ("mid", (dvals > lo) & (dvals <= hi)),
                        ("deep", dvals > hi)):
            rep.rim_bands[name] = {
                "recall@1": float(hit1[m].mean()) if m.any() else float("nan"),
                "n_voxels": int(m.sum()),
            }
    return rep


# ---------------------------------------------------------------- self-test

def _sphere_shell(n=64, r=20, center=None, thickness=1.2):
    g = np.stack(np.meshgrid(*[np.arange(n)] * 3, indexing="ij"), -1) + 0.5
    c = np.array(center if center is not None else [n / 2] * 3)
    d = np.linalg.norm(g - c, axis=-1)
    return np.abs(d - r) < thickness


def _self_test():
    n = 64
    full = _sphere_shell(n)
    # front half visible (camera at -x)
    x = np.arange(n)[:, None, None] + 0.5
    vis = full & np.broadcast_to(x < n / 2, full.shape)
    cam = np.array([-1.5, 0.0, 0.0])  # canonical coords
    free = carve_free_space(vis, cam)

    print(f"vis={vis.sum()} free={free.sum()}")
    assert not (free & vis).any()

    # 1) perfect prediction -> F1@0 == 1 in visible region, FSV == 0
    rep = evaluate_faithfulness(vis, full, free)
    print("perfect:  F1@0=%.3f F1@1=%.3f FSV=%.4f Hausdorff=%.1f" % (
        rep.f1[0]["f1"], rep.f1[1]["f1"], rep.fsv_rate,
        rep.dist_gt_to_pred["max"]))
    assert rep.f1[0]["f1"] > 0.999 and rep.fsv_rate == 0.0

    # 2) shifted prediction -> F1@0 drops, F1@2 recovers, monotone in r
    shifted = np.roll(full, 2, axis=1)
    rep2 = evaluate_faithfulness(vis, shifted, free)
    print("shift+2:  F1@0=%.3f F1@1=%.3f F1@2=%.3f" % tuple(
        rep2.f1[r]["f1"] for r in (0, 1, 2)))
    assert rep2.f1[0]["f1"] < rep2.f1[1]["f1"] <= rep2.f1[2]["f1"]
    assert rep2.f1[0]["f1"] < rep.f1[0]["f1"]

    # 3) bulge toward camera -> free-space violation fires
    bulge = full.copy()
    bulge[8:12, 28:36, 28:36] = True
    rep3 = evaluate_faithfulness(vis, bulge, free)
    print("bulge:    FSV=%.4f max_depth=%.1f" % (rep3.fsv_rate, rep3.fsv_max_depth))
    assert rep3.fsv_rate > 0

    # 4) rim-broken prediction -> near band suffers, deep stays
    rimless = full.copy()
    rim_zone = np.abs(x - n / 2) < 4
    rimless &= ~np.broadcast_to(rim_zone, full.shape)
    rep4 = evaluate_faithfulness(vis, rimless, free, gt_full=full)
    b = rep4.rim_bands
    for name in ("near", "mid", "deep"):
        assert b[name]["n_voxels"] > 0, f"empty {name} band"
    print("rimcut:   near=%.3f mid=%.3f deep=%.3f" % (
        b["near"]["recall@1"], b["mid"]["recall@1"], b["deep"]["recall@1"]))
    assert b["near"]["recall@1"] < b["deep"]["recall@1"]

    print("ALL SELF-TESTS PASSED")


if __name__ == "__main__":
    _self_test()