"""Snap + merge: snap baseline verts toward v9, append unmatched v9 verts.

Best params from sweep (170 val scenes):
  weight=0.80, radius=2.0m, greedy  ->  mean HSS 0.4151 (best of weight sweep 0.6-0.9)
"""
import numpy as np
from scipy.spatial.distance import cdist

SNAP_WEIGHT = 0.80
SNAP_DIST   = 2.0   # metres


def snap_midpoint_plus_unmatched(bl_v, bl_e, v9_v,
                                  weight=SNAP_WEIGHT, max_dist=SNAP_DIST):
    """For each baseline vert, snap weight fraction toward its nearest v9 vert
    if within max_dist. Append unmatched v9 verts as extra vertices.

    Returns (new_vertices, edges) — edges still index the original baseline vertices.
    """
    bl_v = np.asarray(bl_v, dtype=np.float64).copy()
    v9_v = np.asarray(v9_v, dtype=np.float64)
    bl_e = np.asarray(bl_e, dtype=np.int64)
    if len(v9_v) == 0:
        return bl_v, bl_e
    dists     = cdist(bl_v, v9_v)
    nearest   = dists.argmin(axis=1)
    near_dist = dists[np.arange(len(bl_v)), nearest]
    snap_mask = near_dist < max_dist
    bl_v[snap_mask] = ((1 - weight) * bl_v[snap_mask]
                       + weight * v9_v[nearest[snap_mask]])
    claimed   = set(nearest[snap_mask].tolist())
    unmatched = [v9_v[i] for i in range(len(v9_v)) if i not in claimed]
    if unmatched:
        bl_v = np.concatenate([bl_v, np.array(unmatched)], axis=0)
    return bl_v, bl_e