"""Assign mesh faces to fitted superquadrics via curvature-atom-level voting
against per-SQ voxel point clouds.

The mesh is first over-segmented (by an upstream step) into small, curvature-
consistent "atoms" delimited by dihedral-angle boundaries. Each atom is treated
as an indivisible voting unit: every face within an atom votes for its nearest
SQ voxel point cloud, and the atom is assigned the SQ with the most votes. Each
face then inherits its atom's label. Because the atom boundaries are already
curvature-aware, the resulting seams follow those boundaries, avoiding erratic
seams across smooth surfaces while preserving connectivity.

Inputs:
  mesh_path        Original mesh.ply (same mesh used for SQ fitting, in its
                   original coordinate system).
  face_labels_npy  Per-face atom labels, shape (F,) int32, with atom ids in
                   [0..A-1] and -1 marking faces left unassigned by the
                   over-segmentation.
  sq_dir           Directory of fitted superquadrics containing post_sq_*.ply.
  output_dir       Output directory.

Outputs:
  face_labels_v8.npy           (F,) int32, the SQ index per face, -1 = orphan.
  mesh_mapped_v8.ply           Per-face colored mesh.
  per_sq_xx.ply                One submesh per assigned SQ.
  report.json                  Summary statistics.

Key parameters:
  vote_mode            "count" | "exp", the per-face vote weighting.
  vote_tau             Temperature used for "exp" voting.
  orphan_face_dist     A face with NN distance above this is excluded from voting.
  orphan_atom_frac     An atom with more than this fraction of orphan faces is
                       labeled -1.
  min_atom_size        Atoms smaller than this (in face count) are treated as
                       orphans.
"""

from __future__ import annotations

import os
import json
import glob
import time
import argparse
import colorsys
import numpy as np
import trimesh
from scipy.spatial import cKDTree


# =========================================================================
# Utilities
# =========================================================================
def _palette(n: int) -> np.ndarray:
    """Return n distinct RGBA colors from a golden-ratio HSV palette."""
    phi = 0.618033988749895
    out = np.empty((n, 4), dtype=np.uint8)
    for i in range(n):
        hue = (i * phi) % 1.0
        s   = 0.85 if (i % 2 == 0) else 0.65
        v   = 0.95 if (i % 2 == 0) else 0.78
        r, g, b = colorsys.hsv_to_rgb(hue, s, v)
        out[i] = (int(r * 255), int(g * 255), int(b * 255), 255)
    return out


def _normalize_mesh(verts: np.ndarray, half: float = 0.5,
                    margin: float = 1e-6) -> tuple[np.ndarray, float, np.ndarray]:
    """Normalize vertices to fit within [-0.5+margin, 0.5-margin]^3, matching
    the normalization used during SQ fitting. Returns the normalized vertices,
    the applied scale, and the center used."""
    center = (verts.min(0) + verts.max(0)) * 0.5
    ext    = float((verts.max(0) - verts.min(0)).max())
    if ext <= 0:
        raise ValueError(f"Invalid mesh extent: {ext}")
    scale = (half - margin) * 2.0 / ext
    verts_n = (verts - center) * scale
    return verts_n, scale, center


def _load_sq_voxel_groups(sq_dir: str) -> tuple[list[str], list[np.ndarray]]:
    """Load the canonical post_sq_<N>.ply voxel groups from the SQ fitting
    output directory, returning (group_names, voxel_points_list).

    Only post_sq_<number>.ply files are matched; visualization variants such as
    post_sq_tight_* and post_sq_mesh_* are excluded (otherwise multiple files
    for the same SQ would be double-counted and corrupt the mapping).
    """
    files = sorted(
        glob.glob(os.path.join(sq_dir, "post_sq_[0-9]*.ply")),
        key=lambda p: int(os.path.splitext(os.path.basename(p))[0].split("_")[-1]),
    )
    names: list[str] = []
    pts: list[np.ndarray] = []
    for f in files:
        m = trimesh.load(f, force="mesh", process=False)
        v = np.asarray(m.vertices)
        if len(v) == 0:
            continue
        # Voxel meshes contain many vertices (8 per cube x N voxels); subsample
        # to keep the KDTree small.
        if len(v) > 20000:
            sel = np.random.RandomState(0).choice(len(v), 20000, replace=False)
            v = v[sel]
        names.append(os.path.splitext(os.path.basename(f))[0])
        pts.append(v.astype(np.float32))
    return names, pts


def _build_global_tree(sq_pts: list[np.ndarray]) -> tuple[cKDTree, np.ndarray]:
    all_pts = np.vstack(sq_pts)
    all_gid = np.concatenate(
        [np.full(len(p), i, dtype=np.int32) for i, p in enumerate(sq_pts)]
    )
    return cKDTree(all_pts), all_gid


# =========================================================================
# Core: per-atom voting
# =========================================================================
def map_with_atoms(
    mesh: trimesh.Trimesh,
    face_labels_v4: np.ndarray,          # (F,) atom id, -1 = orphan
    sq_pts: list[np.ndarray],            # K SQ voxel point clouds
    *,
    vote_mode: str = "count",            # "count" | "exp"
    vote_tau: float = 0.025,             # temperature for "exp" voting
    orphan_face_dist: float = 0.04,      # per-face NN distance threshold
    orphan_atom_frac: float = 0.6,       # max orphan-face fraction within an atom
    min_atom_size: int = 3,              # atoms below this are demoted to orphan
    voxel_scale_hint: float | None = None,
) -> dict:
    """Assign each mesh face to an SQ via per-atom voting.

    Returns a dict with:
      face_label    (F,) int32   final SQ index, -1 = orphan
      atom_label    (A,) int32   per-atom label
      atom_conf     (A,) float64 winning vote share (winning vote / total)
      init_dist     (F,) float32 per-face nearest-neighbor distance
      stats         dict
    """
    # 0. Normalize to align with the SQ fitting coordinate system.
    verts_n, scale, center = _normalize_mesh(np.asarray(mesh.vertices, dtype=np.float64))
    faces = np.asarray(mesh.faces, dtype=np.int64)
    F = len(faces)

    face_centers = verts_n[faces].mean(axis=1)        # (F, 3) in normalized coords
    print(f"  [v8] mesh: F={F}  v4_atoms={int(face_labels_v4.max()) + 1}", flush=True)

    # 1. Global SQ KDTree.
    tree, voxel_gid = _build_global_tree(sq_pts)
    K = len(sq_pts)
    print(f"  [v8] {K} SQ groups, {len(voxel_gid)} total voxel pts", flush=True)

    # 2. Per-face nearest neighbor.
    init_dist, init_idx = tree.query(face_centers, k=1)
    init_sq = voxel_gid[init_idx].astype(np.int32)
    init_dist = init_dist.astype(np.float32)

    # 3. Atom-level voting.
    atom_ids = face_labels_v4.astype(np.int64)
    valid = atom_ids >= 0                              # exclude orphan faces
    A = int(atom_ids.max()) + 1 if valid.any() else 0

    if A == 0:
        print("  [v8] WARNING: no v4 atoms, falling back to per-face NN", flush=True)
        face_label = init_sq.copy()
        face_label[init_dist > orphan_face_dist] = -1
        return {
            "face_label": face_label,
            "atom_label": np.zeros(0, dtype=np.int32),
            "atom_conf":  np.zeros(0, dtype=np.float64),
            "init_dist":  init_dist,
            "scale":      scale,
            "center":     center,
            "stats": {"A": 0, "K": K, "F": int(F)},
        }

    # Compute per-face vote weights.
    is_face_orphan = init_dist > orphan_face_dist
    if vote_mode == "exp":
        w = np.exp(-init_dist / max(vote_tau, 1e-6)).astype(np.float64)
    elif vote_mode == "count":
        w = np.ones(F, dtype=np.float64)
    else:
        raise ValueError(f"unknown vote_mode {vote_mode}")
    w[~valid] = 0.0
    w[is_face_orphan] = 0.0  # orphan faces do not vote

    # Accumulate the (A, K) vote matrix.
    cnt = np.zeros((A, K), dtype=np.float64)
    if valid.any():
        idx_a = atom_ids[valid & ~is_face_orphan]
        idx_k = init_sq[valid & ~is_face_orphan]
        ww    = w[valid & ~is_face_orphan]
        np.add.at(cnt, (idx_a, idx_k), ww)

    total = cnt.sum(axis=1)                                       # (A,)
    atom_label = np.where(total > 0, cnt.argmax(axis=1), -1).astype(np.int32)
    atom_conf  = np.where(total > 0, cnt.max(axis=1) / np.clip(total, 1e-9, None), 0.0)

    # Atom-level orphan: too high a fraction of orphan faces.
    atom_size = np.zeros(A, dtype=np.int64)
    np.add.at(atom_size, atom_ids[valid], 1)
    orphan_in_atom = np.zeros(A, dtype=np.int64)
    np.add.at(orphan_in_atom, atom_ids[valid & is_face_orphan], 1)
    orphan_frac = orphan_in_atom / np.clip(atom_size, 1, None)
    too_orphan_mask = orphan_frac > orphan_atom_frac
    too_small_mask  = atom_size < min_atom_size
    atom_orphan = (atom_label < 0) | too_orphan_mask | too_small_mask
    n_atom_orphan = int(atom_orphan.sum())
    atom_label[atom_orphan] = -1

    # 4. Propagate atom labels to faces.
    face_label = np.full(F, -1, dtype=np.int32)
    face_label[valid] = atom_label[atom_ids[valid]]

    # 5. Orphan faces: assign individually by nearest neighbor if close enough.
    v4_orphan_mask = ~valid
    if v4_orphan_mask.any():
        good = v4_orphan_mask & (init_dist <= orphan_face_dist)
        face_label[good] = init_sq[good]

    # Force orphan faces back to -1 in case voting pulled them in.
    face_label[is_face_orphan & v4_orphan_mask] = -1

    n_hit = int((face_label >= 0).sum())
    sqs_used = int(np.unique(face_label[face_label >= 0]).size)
    print(f"  [v8] atom_orphan: {n_atom_orphan}/{A}  ({n_atom_orphan/max(A,1)*100:.1f}%)",
          flush=True)
    print(f"  [v8] face_label coverage: {n_hit}/{F}  ({n_hit/F*100:.2f}%)  | SQs hit: {sqs_used}/{K}",
          flush=True)
    print(f"  [v8] atom_conf median={np.median(atom_conf[~atom_orphan] if (~atom_orphan).any() else [0]):.3f}",
          flush=True)

    stats = {
        "F": int(F), "A": int(A), "K": int(K),
        "n_face_orphan_NN": int(is_face_orphan.sum()),
        "n_v4_orphan_face": int(v4_orphan_mask.sum()),
        "n_atom_orphan":    int(n_atom_orphan),
        "atom_size_min":    int(atom_size.min()) if A > 0 else 0,
        "atom_size_median": int(np.median(atom_size)) if A > 0 else 0,
        "atom_size_max":    int(atom_size.max()) if A > 0 else 0,
        "atom_conf_mean":   float(atom_conf[~atom_orphan].mean()) if (~atom_orphan).any() else 0.0,
        "atom_conf_p25":    float(np.percentile(atom_conf[~atom_orphan], 25)) if (~atom_orphan).any() else 0.0,
        "face_coverage":    float(n_hit / F),
        "sqs_used":         int(sqs_used),
        "vote_mode":        vote_mode,
        "vote_tau":         vote_tau,
        "orphan_face_dist": orphan_face_dist,
        "orphan_atom_frac": orphan_atom_frac,
        "min_atom_size":    min_atom_size,
    }
    return {
        "face_label": face_label,
        "atom_label": atom_label,
        "atom_conf":  atom_conf,
        "init_dist":  init_dist,
        "scale":      float(scale),
        "center":     center.astype(np.float64),
        "stats":      stats,
    }


# =========================================================================
# Save
# =========================================================================
def save_results(
    mesh: trimesh.Trimesh,
    result: dict,
    sq_names: list[str],
    output_dir: str,
    save_per_sq: bool = True,
):
    """Write the colored mesh, per-face labels, optional per-SQ submeshes, and
    a JSON report to output_dir."""
    os.makedirs(output_dir, exist_ok=True)
    face_label = result["face_label"]
    K = len(sq_names)
    palette = _palette(max(K, 1))
    orphan_color = np.array([90, 90, 90, 255], dtype=np.uint8)

    faces = np.asarray(mesh.faces, dtype=np.int64)
    verts = np.asarray(mesh.vertices, dtype=np.float64)

    # face colors
    face_colors = np.tile(orphan_color, (len(faces), 1))
    valid = face_label >= 0
    face_colors[valid] = palette[face_label[valid]]

    # Per-vertex color = mode of incident face labels (for eval-script compatibility).
    vertex_colors = np.tile(orphan_color, (len(verts), 1))
    # Each face adds its label to a count table for its 3 vertices, then argmax.
    if K > 0:
        # (V, K+1) table; the extra column counts orphans.
        bins = np.zeros((len(verts), K + 1), dtype=np.int32)
        # Orphan faces are counted in column K.
        eff = np.where(valid, face_label, K).astype(np.int64)
        for c in range(3):
            np.add.at(bins, (faces[:, c], eff), 1)
        winner = bins.argmax(axis=1)
        is_v_orphan = winner == K
        vertex_colors[~is_v_orphan] = palette[winner[~is_v_orphan]]

    vis = trimesh.Trimesh(vertices=verts, faces=faces, process=False)
    vis.visual.face_colors   = face_colors
    vis.visual.vertex_colors = vertex_colors
    vis_path = os.path.join(output_dir, "mesh_mapped_v8.ply")
    vis.export(vis_path)
    print(f"  -> {vis_path}", flush=True)

    np.save(os.path.join(output_dir, "face_labels_v8.npy"), face_label)

    if save_per_sq:
        n_hit = 0
        for k in range(K):
            mask = (face_label == k)
            if not mask.any():
                continue
            try:
                sub = mesh.submesh([np.where(mask)[0]],
                                   only_watertight=False, append=True)
            except Exception:
                sub = None
            if sub is None or len(sub.faces) == 0:
                continue
            sub.visual.face_colors = palette[k]
            sub.export(os.path.join(output_dir, f"per_sq_{k:03d}.ply"))
            n_hit += 1
        print(f"  -> {n_hit}/{K} per-SQ submeshes saved", flush=True)

    report = {
        "stats":    result["stats"],
        "sq_names": sq_names,
    }
    with open(os.path.join(output_dir, "report.json"), "w") as f:
        json.dump(report, f, indent=2, default=lambda x: float(x) if isinstance(x, np.floating) else int(x))


# =========================================================================
# CLI
# =========================================================================
def main():
    """Command-line entry point: load the mesh, atom labels, and SQ voxel
    groups, run per-atom voting, and write the results."""
    ap = argparse.ArgumentParser(description="mesh_mapper_v8: v4 atom -> SQ voting")
    ap.add_argument("--mesh_path",       type=str, required=True)
    ap.add_argument("--face_labels_npy", type=str, required=True)
    ap.add_argument("--sq_dir",          type=str, required=True,
                    help="dir with post_sq_*.ply files (sq_fit_v20 output)")
    ap.add_argument("--output_dir",      type=str, required=True)
    ap.add_argument("--vote_mode",       type=str, default="count",
                    choices=["count", "exp"])
    ap.add_argument("--vote_tau",        type=float, default=0.025)
    ap.add_argument("--orphan_face_dist", type=float, default=0.04)
    ap.add_argument("--orphan_atom_frac", type=float, default=0.6)
    ap.add_argument("--min_atom_size",   type=int, default=3)
    ap.add_argument("--no_per_sq",       action="store_true")
    args = ap.parse_args()

    t0 = time.time()
    mesh = trimesh.load(args.mesh_path, force="mesh", process=False)
    face_labels_v4 = np.load(args.face_labels_npy).astype(np.int32)
    sq_names, sq_pts = _load_sq_voxel_groups(args.sq_dir)
    if not sq_names:
        raise SystemExit(f"No post_sq_*.ply found in {args.sq_dir}")

    result = map_with_atoms(
        mesh, face_labels_v4, sq_pts,
        vote_mode=args.vote_mode, vote_tau=args.vote_tau,
        orphan_face_dist=args.orphan_face_dist,
        orphan_atom_frac=args.orphan_atom_frac,
        min_atom_size=args.min_atom_size,
    )
    save_results(mesh, result, sq_names, args.output_dir,
                 save_per_sq=not args.no_per_sq)

    print(f"\n[v8] done in {time.time() - t0:.1f}s")


if __name__ == "__main__":
    main()