"""Geometry metrics vs a GT mesh (LOCKED spec).

Both meshes are already canonical in [-0.5,0.5]^3; NO ICP / NO alignment.
100k surface points sampled per mesh (same count both sides) via
pytorch3d.ops.sample_points_from_meshes. NN distances via knn_points.

Reported:
  * Chamfer L1  (headline): mean over both directions of the sqrt Euclidean
    nearest-neighbour distance. Units = fraction of the unit-cube edge.
  * Chamfer L2 : pytorch3d.loss.chamfer_distance default (mean squared).
  * F-score@{0.01,0.02,0.05}: P = frac pred within tau of GT,
    R = frac GT within tau of pred, F = 2PR/(P+R). Headline F@0.02.
  * Normal Consistency: mean abs cosine of normals at mutual NN (both dirs).
  * Volume IoU@128^3: solid voxel occupancy on a shared 128^3 grid, I/U.
"""

from __future__ import annotations
# FINAL: = metrics/appeval/geometry.py + extent-relative CD/F keys, optional vol IoU. Original keys unchanged.

import numpy as np
import torch
import trimesh

from pytorch3d.structures import Meshes
from pytorch3d.ops import sample_points_from_meshes, knn_points
from pytorch3d.loss import chamfer_distance

N_SAMPLES = 100_000
TAUS = (0.01, 0.02, 0.05)
VOX_RES = 128


def _to_p3d(mesh: trimesh.Trimesh, device):
    v = torch.as_tensor(np.asarray(mesh.vertices), dtype=torch.float32, device=device)
    f = torch.as_tensor(np.asarray(mesh.faces), dtype=torch.int64, device=device)
    return Meshes(verts=[v], faces=[f])


def _sample(mesh, device, n=N_SAMPLES, seed=0):
    torch.manual_seed(seed)
    m = _to_p3d(mesh, device)
    pts, nrm = sample_points_from_meshes(m, n, return_normals=True)
    return pts, nrm  # (1,n,3), (1,n,3)


def _voxelize_solid(mesh: trimesh.Trimesh, res=VOX_RES):
    """Solid occupancy on a fixed res^3 grid over [-0.5,0.5]^3 (bool array)."""
    pitch = 1.0 / res
    grid = np.zeros((res, res, res), dtype=bool)
    try:
        vox = mesh.voxelized(pitch=pitch)
        try:
            vox = vox.fill()
        except Exception:
            pass  # non-watertight -> hollow occupancy, still comparable
        pts = np.asarray(vox.points)  # world-space centres of occupied cells
    except Exception:
        # last-resort: dense surface samples
        pts, _ = trimesh.sample.sample_surface(mesh, 2_000_000, seed=0)
        pts = np.asarray(pts)
    idx = np.floor((pts + 0.5) * res).astype(np.int64)
    idx = np.clip(idx, 0, res - 1)
    grid[idx[:, 0], idx[:, 1], idx[:, 2]] = True
    return grid


def geometry_metrics(pred_mesh: trimesh.Trimesh, gt_mesh: trimesh.Trimesh,
                     device="cuda", seed=0, with_vol=True):
    """Return a dict of all geometry metrics. NO alignment applied."""
    pp, pn = _sample(pred_mesh, device, seed=seed)
    gp, gn = _sample(gt_mesh, device, seed=seed)

    # NN both directions (knn returns squared dists)
    p2g = knn_points(pp, gp, K=1)
    g2p = knn_points(gp, pp, K=1)
    d_p2g = p2g.dists[..., 0].clamp_min(0).sqrt()[0]   # pred->gt  (n,)
    d_g2p = g2p.dists[..., 0].clamp_min(0).sqrt()[0]   # gt->pred  (n,)

    cd_l1 = 0.5 * (d_p2g.mean() + d_g2p.mean())
    cd_l2, _ = chamfer_distance(pp, gp)                # default mean squared

    fscore = {}
    for tau in TAUS:
        precision = (d_p2g < tau).float().mean()
        recall = (d_g2p < tau).float().mean()
        f = (2 * precision * recall / (precision + recall + 1e-12))
        fscore[tau] = {"precision": float(precision), "recall": float(recall),
                       "f": float(f)}

    # normal consistency at mutual NN, both directions
    idx_p2g = p2g.idx[..., 0][0]      # for each pred pt, nearest gt idx
    idx_g2p = g2p.idx[..., 0][0]
    gn0, pn0 = gn[0], pn[0]
    cos_p = (pn0 * gn0[idx_p2g]).sum(-1).abs()
    cos_g = (gn0 * pn0[idx_g2p]).sum(-1).abs()
    nc = 0.5 * (cos_p.mean() + cos_g.mean())

    # volume IoU on shared grid
    iou = float("nan")
    if with_vol:   # optional (slow; meaningless for non-watertight GT, see memory uv-flip note)
        gp_vox = _voxelize_solid(pred_mesh)
        gg_vox = _voxelize_solid(gt_mesh)
        inter = np.logical_and(gp_vox, gg_vox).sum()
        union = np.logical_or(gp_vox, gg_vox).sum()
        iou = float(inter) / float(union) if union > 0 else 0.0

    # ADDED: extent-relative metrics so FB150 (GT max extent 0.9) is comparable with
    # Toys/Omni (1.0): distances divided by the GT max bbox extent (tau scaled by it).
    ext = float((np.asarray(gt_mesh.vertices).max(0) - np.asarray(gt_mesh.vertices).min(0)).max())
    rel = {}
    for tau in TAUS:
        t = tau * ext
        pr_, rc_ = (d_p2g < t).float().mean(), (d_g2p < t).float().mean()
        rel[tau] = float(2 * pr_ * rc_ / (pr_ + rc_ + 1e-12))

    return {
        "gt_extent": ext,
        "cd_l1_rel": float(cd_l1) / ext,
        "f01_rel": rel[0.01], "f02_rel": rel[0.02], "f05_rel": rel[0.05],
        "cd_l1": float(cd_l1),
        "cd_l2": float(cd_l2),
        "fscore": {str(t): fscore[t] for t in TAUS},
        "f01": fscore[0.01]["f"],
        "f02": fscore[0.02]["f"],   # headline
        "f05": fscore[0.05]["f"],
        "normal_consistency": float(nc),
        "vol_iou": iou,
    }