from __future__ import annotations

import numpy as np


def merge_vertices_iterative(vertices: np.ndarray, edges: np.ndarray,
                              start: float = 0.15, end: float = 0.6,
                              n_iters: int = 5):
    """Iterative merge: start with tight threshold, gradually widen.

    Avoids the worst transitive chaining effects of a single wide threshold.
    Each pass merges only the closest pairs first, establishing stable cluster
    centers before wider merges pull in more distant endpoints.

    +0.004 HSS / +0.007 F1 over single-pass merge(0.4) on 1024 val samples.
    """
    pv, pe = vertices, edges
    for t in np.linspace(start, end, n_iters):
        pv, pe = merge_vertices(pv, pe, t)
    return pv, pe


def merge_vertices(vertices: np.ndarray, edges: np.ndarray, thresh: float):
    verts = np.asarray(vertices, dtype=np.float32)
    edges = np.asarray(edges, dtype=np.int64)
    if verts.size == 0 or edges.size == 0:
        return verts, edges

    n = verts.shape[0]
    parent = np.arange(n, dtype=np.int64)

    def find(i):
        while parent[i] != i:
            parent[i] = parent[parent[i]]
            i = parent[i]
        return i

    def union(i, j):
        ri = find(i)
        rj = find(j)
        if ri != rj:
            parent[rj] = ri

    for i in range(n):
        vi = verts[i]
        for j in range(i + 1, n):
            if np.linalg.norm(vi - verts[j]) <= thresh:
                union(i, j)

    clusters = {}
    for i in range(n):
        root = find(i)
        clusters.setdefault(root, []).append(i)

    new_vertices = []
    mapping = {}
    for new_idx, idxs in enumerate(clusters.values()):
        pts = verts[idxs]
        center = pts.mean(axis=0)
        new_vertices.append(center)
        for i in idxs:
            mapping[i] = new_idx

    new_edges = []
    seen = set()
    for a, b in edges:
        na = mapping.get(int(a), int(a))
        nb = mapping.get(int(b), int(b))
        if na == nb:
            continue
        key = (na, nb) if na <= nb else (nb, na)
        if key in seen:
            continue
        seen.add(key)
        new_edges.append([na, nb])

    return np.asarray(new_vertices, dtype=np.float32), np.asarray(new_edges, dtype=np.int64)