"""Edge classifier and classifier-gated edge augmentation.

Contains:
  - ClassificationPointNet: PointNet binary classifier on 6D cylindrical patches
  - colmap_points_xyz_rgb / build_edge_patch_6d: patch builder
  - score_edges_batched: per-sample batched inference
  - augment_hybrid_with_filtered_hc: import high-confidence handcrafted edges
"""
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F


# Cylinder patch geometry (must match the training-time patch generation).
CYL_RADIUS = 0.5
CYL_EXT = 0.25            # extension at each end
MAX_PATCH_POINTS = 1024
AUGMENT_DEDUP_RADIUS = 0.3
AUGMENT_THRESHOLD = 0.55


class ClassificationPointNet(nn.Module):
    """PointNet binary classifier on 6D point cloud patches (xyz + rgb)."""

    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.fc1 = nn.Linear(2048, 1024)
        self.fc2 = nn.Linear(1024, 512)
        self.fc3 = nn.Linear(512, 256)
        self.fc4 = nn.Linear(256, 128)
        self.fc5 = nn.Linear(128, 64)
        self.fc6 = nn.Linear(64, 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.dropout1 = nn.Dropout(0.3)
        self.dropout2 = nn.Dropout(0.4)
        self.dropout3 = nn.Dropout(0.5)
        self.dropout4 = nn.Dropout(0.4)
        self.dropout5 = nn.Dropout(0.3)

    def forward(self, x):
        # x: (B, 6, max_points)
        x1 = F.relu(self.bn1(self.conv1(x)))
        x2 = F.relu(self.bn2(self.conv2(x1)))
        x3 = F.relu(self.bn3(self.conv3(x2)))
        x4 = F.relu(self.bn4(self.conv4(x3)))
        x5 = F.relu(self.bn5(self.conv5(x4)))
        x6 = F.relu(self.bn6(self.conv6(x5)))
        g = torch.max(x6, 2)[0]
        x = F.relu(self.fc1(g));   x = self.dropout1(x)
        x = F.relu(self.fc2(x));   x = self.dropout2(x)
        x = F.relu(self.fc3(x));   x = self.dropout3(x)
        x = F.relu(self.fc4(x));   x = self.dropout4(x)
        x = F.relu(self.fc5(x));   x = self.dropout5(x)
        return self.fc6(x)         # (B, 1) logits


def load_pnet_class(model_path, device=None):
    """Load a trained ClassificationPointNet checkpoint."""
    if device is None:
        device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    model = ClassificationPointNet(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_edge_patch_6d(u_xyz, v_xyz, colmap_xyz, colmap_rgb):
    """6D cylinder patch around edge (u, v). None if too sparse."""
    line = v_xyz - u_xyz
    L = float(np.linalg.norm(line))
    if L < 1e-6:
        return None
    direction = line / L
    ext_start = u_xyz - CYL_EXT * direction
    ext_L = L + 2 * CYL_EXT
    rel = colmap_xyz - ext_start[np.newaxis, :]
    proj = rel @ direction
    in_bounds = (proj >= 0) & (proj <= ext_L)
    closest = ext_start[np.newaxis, :] + proj[:, np.newaxis] * direction[np.newaxis, :]
    perp = np.linalg.norm(colmap_xyz - closest, axis=1)
    in_cyl = in_bounds & (perp <= CYL_RADIUS)
    if int(in_cyl.sum()) <= 10:
        return None
    midpoint = (u_xyz + v_xyz) / 2
    pts_centered = colmap_xyz[in_cyl] - midpoint
    rgb_signed = colmap_rgb[in_cyl] * 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_edges_batched(model, device, vertices, edges, colmap_xyz, colmap_rgb,
                         rng=None, batch_size=32):
    """Score every edge in `edges` (returns list of floats, or None per failed patch)."""
    if rng is None:
        rng = np.random.RandomState(0)
    n = len(edges)
    scores = [None] * n
    if n == 0 or len(vertices) == 0 or len(colmap_xyz) == 0:
        return scores

    patches, indices = [], []
    for i, (u, v) in enumerate(edges):
        u_xyz = vertices[int(u)]
        v_xyz = vertices[int(v)]
        raw = build_edge_patch_6d(u_xyz, v_xyz, 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)
            logits = model(x).squeeze(-1)
            probs = torch.sigmoid(logits).cpu().numpy().reshape(-1)
            for j, p in enumerate(probs):
                scores[indices[start + j]] = float(p)
    return scores


def augment_hybrid_with_filtered_hc(h_v, h_e, user_v, user_e, scores,
                                     thresh=AUGMENT_THRESHOLD,
                                     dedup_radius=AUGMENT_DEDUP_RADIUS):
    """Add handcrafted edges scoring above thresh into the hybrid (v, e)."""
    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_e) == 0 or len(user_v) == 0:
        return h_v, h_e_list

    kept_pairs = [
        (int(u), int(v)) for (u, v), s in zip(user_e, scores)
        if s is not None and s > thresh
    ]
    if not kept_pairs:
        return h_v, h_e_list

    needed_hc_idx = sorted({u for u, _ in kept_pairs} | {v for _, v in kept_pairs})
    new_v_list, user_to_combined = [], {}
    for u_idx in needed_hc_idx:
        pos = user_v[u_idx]
        if len(h_v) > 0:
            d = np.linalg.norm(h_v - pos, axis=1)
            best = int(np.argmin(d))
            if d[best] <= dedup_radius:
                user_to_combined[u_idx] = best
                continue
        user_to_combined[u_idx] = len(h_v) + len(new_v_list)
        new_v_list.append(pos)

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

    existing = {(min(a, b), max(a, b)) for a, b in h_e_list}
    new_edges = []
    for u, v in kept_pairs:
        a, b = user_to_combined[u], user_to_combined[v]
        if a == b:
            continue
        key = (min(a, b), max(a, b))
        if key in existing:
            continue
        existing.add(key)
        new_edges.append((a, b))
    return combined_v, h_e_list + new_edges