"""Multi-view merger and graph optimizer for S23DR 2026."""

import numpy as np
from collections import defaultdict
from typing import Dict, List, Tuple, Optional
from scipy.spatial import KDTree


def merge_vertices_weighted(vert_edge_per_image, th=0.5, depth_confidence=None):
    all_verts, all_weights, connections_3d, types = [], [], [], []
    cur_start = 0
    for img_idx, data in vert_edge_per_image.items():
        if len(data) < 3: continue
        vertices_2d, connections, vertices_3d = data[:3]
        if len(vertices_3d) == 0: continue
        weight = depth_confidence.get(img_idx, 1.0) if depth_confidence else 1.0
        types.extend([v.get('type', 'apex') if isinstance(v, dict) else 'apex' for v in vertices_2d])
        all_verts.append(vertices_3d); all_weights.extend([weight] * len(vertices_3d))
        for (x, y) in connections: connections_3d.append((x + cur_start, y + cur_start))
        cur_start += len(vertices_3d)
    if not all_verts: return np.zeros((2, 3)), [(0, 1)]
    all_verts = np.concatenate(all_verts, axis=0); all_weights = np.array(all_weights)
    types = np.array([1 if t == 'apex' else 0 for t in types])
    tree = KDTree(all_verts); parent = list(range(len(all_verts)))
    def find(x):
        while parent[x] != x: parent[x] = parent[parent[x]]; x = parent[x]
        return x
    def union(a, b):
        ra, rb = find(a), find(b)
        if ra != rb: parent[ra] = rb
    for i, j in tree.query_pairs(r=th):
        if types[i] == types[j]: union(i, j)
    groups = defaultdict(list)
    for i in range(len(all_verts)): groups[find(i)].append(i)
    new_verts, old_to_new = [], {}
    for group_id, (root, members) in enumerate(groups.items()):
        w = all_weights[members]; pos = all_verts[members]
        new_verts.append(np.average(pos, weights=w, axis=0) if w.sum() > 0 else pos.mean(axis=0))
        for m in members: old_to_new[m] = group_id
    new_verts = np.array(new_verts)
    new_conns = set()
    for a, b in connections_3d:
        na, nb = old_to_new.get(a), old_to_new.get(b)
        if na is not None and nb is not None and na != nb: new_conns.add(tuple(sorted((na, nb))))
    return new_verts, list(new_conns)


def snap_to_planes(vertices, connections, n_planes=5, distance_threshold=0.3, min_inliers=4):
    if len(vertices) < 4: return vertices
    snapped = vertices.copy(); remaining = set(range(len(vertices)))
    for plane_iter in range(n_planes):
        if len(remaining) < min_inliers: break
        remaining_idx = list(remaining); pts = vertices[remaining_idx]
        best_plane, best_inliers = None, []
        rng = np.random.RandomState(42 + plane_iter)
        for _ in range(200):
            if len(pts) < 3: break
            sample_idx = rng.choice(len(pts), 3, replace=False)
            p1, p2, p3 = pts[sample_idx]
            normal = np.cross(p2 - p1, p3 - p1); norm = np.linalg.norm(normal)
            if norm < 1e-8: continue
            normal /= norm; d = -np.dot(normal, p1)
            inliers = np.where(np.abs(pts @ normal + d) < distance_threshold)[0]
            if len(inliers) > len(best_inliers): best_inliers = inliers; best_plane = (normal, d)
        if best_plane is None or len(best_inliers) < min_inliers: break
        normal, d = best_plane
        for local_idx in best_inliers:
            global_idx = remaining_idx[local_idx]; pt = vertices[global_idx]
            snapped[global_idx] = pt - (np.dot(normal, pt) + d) * normal; remaining.discard(global_idx)
    return snapped


def regularize_edges(vertices, connections, horizontal_threshold=0.15, vertical_threshold=0.15):
    reg_verts = vertices.copy()
    for a, b in connections:
        if a >= len(vertices) or b >= len(vertices): continue
        direction = reg_verts[b] - reg_verts[a]; length = np.linalg.norm(direction)
        if length < 1e-6: continue
        if abs(direction[2] / length) < np.sin(horizontal_threshold):
            avg_z = (reg_verts[a][2] + reg_verts[b][2]) / 2; reg_verts[a][2] = avg_z; reg_verts[b][2] = avg_z
    return reg_verts


def remove_redundant_edges(vertices, connections, min_edge_length=0.1, max_degree=8):
    edges = list(set(tuple(sorted(c)) for c in connections))
    filtered = [(a, b, np.linalg.norm(vertices[a] - vertices[b])) for a, b in edges if a < len(vertices) and b < len(vertices) and np.linalg.norm(vertices[a] - vertices[b]) >= min_edge_length]
    if not filtered: return [(0, 1)] if len(vertices) >= 2 else []
    filtered.sort(key=lambda x: x[2]); degree = defaultdict(int); final = []
    for a, b, l in filtered:
        if degree[a] < max_degree and degree[b] < max_degree: final.append((a, b)); degree[a] += 1; degree[b] += 1
    return final


def full_graph_optimization(vertices, connections, colmap_rec=None, do_plane_snap=True, do_regularize=True, do_prune=True):
    if len(vertices) < 2 or len(connections) < 1: return vertices, connections
    if do_prune: connections = remove_redundant_edges(vertices, connections)
    if do_plane_snap and len(vertices) >= 4: vertices = snap_to_planes(vertices, connections)
    if do_regularize: vertices = regularize_edges(vertices, connections)
    return vertices, connections