"""Edge filling from 2D gestalt evidence.

For each pair of predicted vertices (V_i, V_j) that is NOT currently an edge:
  1. Project both endpoints into every COLMAP view.
  2. Sample N points along the projected 2D segment.
  3. Count points falling on gestalt edge-class pixels (using the dilated mask).
  4. A view "supports" the candidate edge if support_frac >= min_pixel_frac.
  5. If at least min_views_support views agree, ADD the edge.

This is the inverse of edge_2d_filter.filter_edges_by_2d_support:
the filter regressed q5 because dropping edges based on a binary mask check
hurt recall. Adding edges is asymmetric: false positives waste a precision
slot, but false negatives are catastrophic (real edges missing). With strong
thresholds we should mostly add genuinely-missed edges.

Conservative defaults: 40% min support, 2+ views agreeing, max edge length
5m (most building edges are short). Pairs are scored by closest-first and
capped at max_pair_check to keep cost bounded.

Topology change only — adds new edges, never moves or removes vertices.
Falls back to (pv, pe) on any error.
"""

from __future__ import annotations

import numpy as np
import cv2


def fill_missing_edges_from_2d(
    pv,
    pe,
    sample,
    min_views_support: int = 3,
    min_pixel_frac: float = 0.60,
    max_edge_length_meters: float = 5.0,
    max_pair_check: int = 100,
    max_fills_abs: int = 6,
    max_fills_rel: float = 0.25,
    dilate_px: int = 4,
    sample_steps: int = 20,
):
    """Add edges between existing vertex pairs that have strong 2D edge support.

    Args:
        pv: (N, 3) vertices in world coordinates.
        pe: existing edge list. New edges are appended; existing edges
            are never removed.
        sample: raw dataset entry.
        min_views_support: minimum views with strong mask support to add edge.
        min_pixel_frac: fraction of sampled segment pixels that must lie on
            a gestalt edge-class pixel for a view to count as supporting.
        max_edge_length_meters: skip pairs whose 3D distance exceeds this.
        max_pair_check: hard cap on candidate pairs evaluated per sample
            (sorted by ascending 3D distance, so closest pairs go first).
        dilate_px: edge-mask dilation radius (same as edge_2d_filter).
        sample_steps: number of points sampled along each 2D segment.

    Returns:
        (pv, pe_extended). Vertex array unchanged; edges only grow.
        Falls back to inputs on any error.
    """
    try:
        from hoho2025.example_solutions import convert_entry_to_human_readable
        from mvs_utils import collect_views, project_world_to_image
        from edge_2d_filter import _build_edge_masks

        pv_arr = np.asarray(pv, dtype=np.float64)
        if pv_arr.ndim != 2 or pv_arr.shape[0] < 2:
            return pv, pe

        good = convert_entry_to_human_readable(sample)
        colmap_rec = good.get("colmap") or good.get("colmap_binary")
        if colmap_rec is None:
            return pv, pe

        views = collect_views(colmap_rec, good["image_ids"])
        if len(views) < min_views_support:
            return pv, pe

        view_masks = _build_edge_masks(good, views, dilate_px=dilate_px)
        if not view_masks:
            return pv, pe

        existing = set()
        for a, b in pe:
            a, b = int(a), int(b)
            lo, hi = (a, b) if a < b else (b, a)
            existing.add((lo, hi))

        N = pv_arr.shape[0]

        # Build candidate list: all pairs not already in `existing` and within
        # max_edge_length. Sort by ascending 3D distance and cap.
        candidates = []
        for i in range(N):
            for j in range(i + 1, N):
                if (i, j) in existing:
                    continue
                d = float(np.linalg.norm(pv_arr[i] - pv_arr[j]))
                if d > max_edge_length_meters or d < 1e-3:
                    continue
                candidates.append((d, i, j))
        candidates.sort()
        if len(candidates) > max_pair_check:
            candidates = candidates[:max_pair_check]

        if not candidates:
            return pv, pe

        # Cache view list once (avoid dict-iteration in hot loop)
        view_items = [(img_id, views[img_id]["P"], *view_masks[img_id])
                      for img_id in view_masks]

        # Score each candidate by (total support across views, # views supporting).
        # Then accept only those meeting the threshold, capped by ranking.
        scored = []
        for _, i, j in candidates:
            endpoints = np.stack([pv_arr[i], pv_arr[j]])

            supporting = 0
            total_support = 0.0
            for _img_id, P, mask_bool, H, W in view_items:
                uv, z = project_world_to_image(P, endpoints)
                if z[0] <= 0 or z[1] <= 0:
                    continue
                if not (
                    0 <= uv[0, 0] < W and 0 <= uv[0, 1] < H
                    and 0 <= uv[1, 0] < W and 0 <= uv[1, 1] < H
                ):
                    continue

                t = np.linspace(0.0, 1.0, sample_steps)
                xs = uv[0, 0] + t * (uv[1, 0] - uv[0, 0])
                ys = uv[0, 1] + t * (uv[1, 1] - uv[0, 1])
                xs_i = np.clip(xs.astype(np.int32), 0, W - 1)
                ys_i = np.clip(ys.astype(np.int32), 0, H - 1)

                frac = int(mask_bool[ys_i, xs_i].sum()) / float(sample_steps)
                if frac >= min_pixel_frac:
                    supporting += 1
                    total_support += frac

            if supporting >= min_views_support:
                # Score: prioritize multi-view agreement, break ties on total support
                scored.append((supporting, total_support, i, j))

        if not scored:
            return pv, pe

        # Cap additions: min(absolute cap, rel-fraction of existing edges)
        max_to_add = max(1, min(max_fills_abs,
                                int(max_fills_rel * max(len(pe), 1))))
        scored.sort(reverse=True)
        added = [(i, j) for _, _, i, j in scored[:max_to_add]]

        new_pe = list(pe) + added
        return pv, new_pe

    except Exception:
        return pv, pe