"""Vertex classifier and classifier-gated vertex augmentation.

Analog of edge_classifier.py for vertices. Contains:
  - VertexPointNetDualHead: PointNet with classification + regression heads
  - colmap_points_xyz_rgb / build_vertex_patch_6d: spherical patch builder
  - score_vertices_batched: per-sample batched inference
  - filter_hc_vertices: drop low-confidence handcrafted vertices

Only the classification head is used at deployment; the regression head is
present in the checkpoint but not used.
"""
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F


SPHERE_RADIUS = 1.0
MAX_PATCH_POINTS = 1024
MIN_PATCH_POINTS = 10
DEFAULT_THRESHOLD = 0.55
AUGMENT_THRESHOLD = 0.55
AUGMENT_DEDUP_RADIUS = 0.5         # matches the HSS vertex true-positive radius


class VertexPointNetDualHead(nn.Module):
    """PointNet backbone with dual classification + regression heads.

    Identical to training-time class. Regression head is loaded but not used
    at deployment (it didn't learn beyond predict-zero baseline).
    """

    def __init__(self, input_dim=6, max_points=1024):
        super().__init__()
        self.max_points = max_points
        self.conv1 = nn.Conv1d(input_dim, 64, 1)
        self.conv2 = nn.Conv1d(64, 128, 1)
        self.conv3 = nn.Conv1d(128, 256, 1)
        self.conv4 = nn.Conv1d(256, 512, 1)
        self.conv5 = nn.Conv1d(512, 1024, 1)
        self.conv6 = nn.Conv1d(1024, 2048, 1)
        self.bn1 = nn.BatchNorm1d(64)
        self.bn2 = nn.BatchNorm1d(128)
        self.bn3 = nn.BatchNorm1d(256)
        self.bn4 = nn.BatchNorm1d(512)
        self.bn5 = nn.BatchNorm1d(1024)
        self.bn6 = nn.BatchNorm1d(2048)
        self.fc1 = nn.Linear(2048, 1024)
        self.fc2 = nn.Linear(1024, 512)
        self.fc3 = nn.Linear(512, 256)
        self.d1 = nn.Dropout(0.3)
        self.d2 = nn.Dropout(0.4)
        self.d3 = nn.Dropout(0.3)
        self.cls_head = nn.Linear(256, 1)
        self.reg_head = nn.Linear(256, 3)

    def forward(self, x):
        x = F.relu(self.bn1(self.conv1(x)))
        x = F.relu(self.bn2(self.conv2(x)))
        x = F.relu(self.bn3(self.conv3(x)))
        x = F.relu(self.bn4(self.conv4(x)))
        x = F.relu(self.bn5(self.conv5(x)))
        x = F.relu(self.bn6(self.conv6(x)))
        g = torch.max(x, 2)[0]
        t = F.relu(self.fc1(g));   t = self.d1(t)
        t = F.relu(self.fc2(t));   t = self.d2(t)
        t = F.relu(self.fc3(t));   t = self.d3(t)
        return self.cls_head(t), self.reg_head(t)


def load_vertex_model(model_path, device=None):
    """Load a trained VertexPointNetDualHead checkpoint."""
    if device is None:
        device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    model = VertexPointNetDualHead(input_dim=6, max_points=MAX_PATCH_POINTS)
    ckpt = torch.load(model_path, map_location=device, weights_only=False)
    model.load_state_dict(ckpt['model_state_dict'])
    model.to(device).eval()
    return model


def colmap_points_xyz_rgb(colmap_rec):
    """Return (xyz, rgb_normalized_0_1) for all COLMAP points."""
    xyz_list, rgb_list = [], []
    for _, p3D in colmap_rec.points3D.items():
        xyz_list.append(p3D.xyz)
        rgb_list.append(p3D.color / 255.0)
    if not xyz_list:
        return np.empty((0, 3)), np.empty((0, 3))
    return np.array(xyz_list), np.array(rgb_list)


def build_vertex_patch_6d(v_xyz, colmap_xyz, colmap_rgb, radius=SPHERE_RADIUS):
    """Sphere of COLMAP points around vertex. 6D = [xyz_relative, rgb_signed]."""
    rel = colmap_xyz - v_xyz[np.newaxis, :]
    dist = np.linalg.norm(rel, axis=1)
    in_sphere = dist <= radius
    if int(in_sphere.sum()) <= MIN_PATCH_POINTS:
        return None
    pts_centered = rel[in_sphere]
    rgb_signed = colmap_rgb[in_sphere] * 2.0 - 1.0
    return np.hstack([pts_centered, rgb_signed])


def _pad_or_sample_patch(patch_6d, max_pts=MAX_PATCH_POINTS, rng=None):
    if rng is None:
        rng = np.random
    n = patch_6d.shape[0]
    if n >= max_pts:
        idx = rng.choice(n, max_pts, replace=False)
        return patch_6d[idx]
    out = np.zeros((max_pts, 6), dtype=np.float32)
    out[:n] = patch_6d
    return out


def score_vertices_batched(model, device, vertices, colmap_xyz, colmap_rgb,
                            rng=None, batch_size=64):
    """Score every vertex. Returns list of floats in [0,1] (None for failed patches)."""
    if rng is None:
        rng = np.random.RandomState(0)
    n = len(vertices)
    scores = [None] * n
    if n == 0 or len(colmap_xyz) == 0:
        return scores

    patches, indices = [], []
    for i in range(n):
        raw = build_vertex_patch_6d(vertices[i], colmap_xyz, colmap_rgb)
        if raw is None:
            continue
        patches.append(_pad_or_sample_patch(raw, rng=rng).astype(np.float32))
        indices.append(i)

    if not patches:
        return scores

    batch = np.stack(patches, axis=0).transpose(0, 2, 1)   # (N, 6, 1024)
    with torch.no_grad():
        for start in range(0, len(batch), batch_size):
            end = min(start + batch_size, len(batch))
            x = torch.from_numpy(batch[start:end]).to(device)
            cls_logit, _reg = model(x)
            probs = torch.sigmoid(cls_logit.squeeze(-1)).cpu().numpy().reshape(-1)
            for j, p in enumerate(probs):
                scores[indices[start + j]] = float(p)
    return scores


def filter_hc_vertices(vertices, scores, threshold=DEFAULT_THRESHOLD):
    """Keep vertices where score > threshold (or score is None -- couldn't build patch).

    Returns (filtered_vertices, keep_mask) where keep_mask is a bool ndarray.
    Vertices with score=None (too sparse to build a patch) are kept, since there
    is no confidence signal for them. The deployed policy uses
    augment_hybrid_with_filtered_hc_vertices below rather than this filter.
    """
    vertices = np.asarray(vertices)
    if len(vertices) == 0:
        return vertices, np.array([], dtype=bool)
    keep = np.array([(s is None) or (s > threshold) for s in scores], dtype=bool)
    return vertices[keep], keep


def augment_hybrid_with_filtered_hc_vertices(h_v, h_e, user_v, scores,
                                              threshold=AUGMENT_THRESHOLD,
                                              dedup_radius=AUGMENT_DEDUP_RADIUS):
    """Add high-scoring handcrafted vertices as orphan vertices (no edges).

    Deduplicates at the HSS vertex true-positive radius (0.5m): a candidate is
    skipped if within dedup_radius of any existing handcrafted vertex or an
    already-added orphan, so a single ground-truth vertex is not double-counted.

    Returns (combined_v, h_e_list). Edges are unchanged (orphans only).
    """
    h_v = np.asarray(h_v, dtype=np.float32)
    user_v = np.asarray(user_v, dtype=np.float32)
    h_e_list = [(int(a), int(b)) for a, b in h_e]
    if scores is None or len(user_v) == 0:
        return h_v, h_e_list

    kept_idx = [i for i, s in enumerate(scores) if s is not None and s > threshold]
    if not kept_idx:
        return h_v, h_e_list

    new_v_list = []
    for i in kept_idx:
        pos = user_v[i]
        if len(h_v) > 0:
            if float(np.linalg.norm(h_v - pos, axis=1).min()) <= dedup_radius:
                continue
        if new_v_list:
            stacked = np.stack(new_v_list)
            if float(np.linalg.norm(stacked - pos, axis=1).min()) <= dedup_radius:
                continue
        new_v_list.append(pos)

    if not new_v_list:
        return h_v, h_e_list

    if len(h_v):
        combined_v = np.concatenate([h_v, np.stack(new_v_list, axis=0)], axis=0)
    else:
        combined_v = np.stack(new_v_list, axis=0)
    return combined_v, h_e_list