"""3D PointNet for vertex candidate classification + position refinement.

Architecture from S23DR 2025 winning solution (Skvrna & Neumann):
  - 7x Conv1D with channel attention
  - 3 heads: classification (is real vertex), position offset, score
"""
import numpy as np
import torch
import torch.nn as nn

VERTEX_CLASSES = [0, 1, 2]   # compressed class_id: apex=0, eave_end_point=1, flashing_end_point=2
CLUSTER_RADIUS = 0.5          # m — merge clusters sharing >50% points within this
PATCH_HALF     = 2.0          # m — cubic patch half-side around candidate centroid
SNAP_RADIUS    = 0.5          # m — GT vertex match threshold for labelling


# ---------------------------------------------------------------------------
# Model
# ---------------------------------------------------------------------------

class ChannelAttention(nn.Module):
    def __init__(self, channels, reduction=16):
        super().__init__()
        self.fc = nn.Sequential(
            nn.Linear(channels, channels // reduction),
            nn.ReLU(inplace=True),
            nn.Linear(channels // reduction, channels),
            nn.Sigmoid(),
        )

    def forward(self, x):
        # x: (B, C, N) — pool over N, gate channels
        g = x.mean(-1)          # (B, C)
        return x * self.fc(g).unsqueeze(-1)


class VertexPointNet(nn.Module):
    """
    Input:  (B, in_dim, N) point features
    Output: logits_cls (B,1), offset_3d (B,3), score (B,1)
    """
    def __init__(self, in_dim=11, dropout=0.3):
        super().__init__()
        dims = [in_dim, 64, 128, 256, 512, 1024, 1024, 2048]
        layers = []
        for d_in, d_out in zip(dims[:-1], dims[1:]):
            layers += [
                nn.Conv1d(d_in, d_out, 1),
                nn.BatchNorm1d(d_out),
                nn.LeakyReLU(0.2, inplace=True),
            ]
        self.encoder = nn.Sequential(*layers)
        self.attn = ChannelAttention(2048)

        self.shared = nn.Sequential(
            nn.Linear(2048 * 2, 1024),   # max + avg concat
            nn.GroupNorm(32, 1024),
            nn.ReLU(inplace=True),
            nn.Dropout(dropout),
            nn.Linear(1024, 512),
            nn.GroupNorm(16, 512),
            nn.ReLU(inplace=True),
            nn.Dropout(dropout),
            nn.Linear(512, 512),
            nn.GroupNorm(16, 512),
            nn.ReLU(inplace=True),
        )

        head = lambda out, act=None: nn.Sequential(
            nn.Linear(512, 256), nn.ReLU(inplace=True),
            nn.Linear(256, 128), nn.ReLU(inplace=True),
            nn.Linear(128, 64),  nn.ReLU(inplace=True),
            nn.Linear(64, out),
            *([act] if act else []),
        )
        self.head_cls   = head(1)
        self.head_score = head(1, nn.Softplus())
        self.head_off   = head(3)

    def forward(self, x):
        # x: (B, in_dim, N)
        f = self.encoder(x)           # (B, 2048, N)
        f = self.attn(f)
        g_max = f.max(-1).values      # (B, 2048)
        g_avg = f.mean(-1)            # (B, 2048)
        g = torch.cat([0.7 * g_max + 0.3 * g_avg,
                        g_max], dim=-1)               # (B, 4096)
        g = self.shared(g)                            # (B, 512)
        return self.head_cls(g), self.head_off(g), self.head_score(g)


class EdgePointNet(nn.Module):
    """Binary edge existence from cylindrical patch. Input: (B, 6, N).

    Uses max+avg pooling concat (4096-d global) + channel attention,
    matching VertexPointNet's aggregation strategy.
    """
    def __init__(self, in_dim=6, dropout=0.4):
        super().__init__()
        dims = [in_dim, 64, 128, 256, 512, 1024, 2048]
        layers = []
        for d_in, d_out in zip(dims[:-1], dims[1:]):
            layers += [
                nn.Conv1d(d_in, d_out, 1),
                nn.BatchNorm1d(d_out),
                nn.ReLU(inplace=True),
            ]
        self.encoder = nn.Sequential(*layers)
        self.attn    = ChannelAttention(2048)
        self.head = nn.Sequential(
            nn.Linear(2048 * 2, 1024), nn.GroupNorm(32, 1024), nn.ReLU(inplace=True), nn.Dropout(dropout),
            nn.Linear(1024, 512),      nn.GroupNorm(16, 512),  nn.ReLU(inplace=True), nn.Dropout(dropout),
            nn.Linear(512, 256),                               nn.ReLU(inplace=True), nn.Dropout(dropout),
            nn.Linear(256, 128),                               nn.ReLU(inplace=True),
            nn.Linear(128, 1),
        )

    def forward(self, x):
        f    = self.attn(self.encoder(x))              # (B, 2048, N)
        g    = torch.cat([f.max(-1).values,
                          f.mean(-1)], dim=-1)         # (B, 4096)
        return self.head(g)                            # (B, 1)


# ---------------------------------------------------------------------------
# Geometry helpers
# ---------------------------------------------------------------------------

def generate_vertex_candidates(xyz_world, class_id, source=None):
    """Stage 1: cluster COLMAP points of gestalt classes 1/2/3.

    Returns list of (centroid_3d, member_mask) where member_mask indexes xyz_world.
    """
    # Filter to COLMAP + vertex classes
    vmask = np.isin(class_id, VERTEX_CLASSES)
    if source is not None:
        vmask &= (source == 0)   # COLMAP only
    vpts = xyz_world[vmask]
    vidx = np.where(vmask)[0]
    if len(vpts) < 3:
        return []

    from scipy.spatial import cKDTree
    tree = cKDTree(vpts)
    pairs = list(tree.query_pairs(CLUSTER_RADIUS))

    # Union-Find clustering
    parent = list(range(len(vpts)))
    def find(x):
        while parent[x] != x:
            parent[x] = parent[parent[x]]
            x = parent[x]
        return x
    for a, b in pairs:
        ra, rb = find(a), find(b)
        if ra != rb:
            parent[ra] = rb

    from collections import defaultdict
    groups = defaultdict(list)
    for i in range(len(vpts)):
        groups[find(i)].append(i)

    candidates = []
    for members in groups.values():
        centroid = vpts[members].mean(0)
        # member_mask in original xyz_world space
        mask = vidx[members]
        candidates.append((centroid, mask))
    return candidates


def extract_vertex_patch(centroid, xyz_world, class_id, source, rgb=None,
                         half=PATCH_HALF, max_pts=512):
    """Extract 11D features for points in cubic patch around centroid."""
    offset = xyz_world - centroid
    inbox = np.all(np.abs(offset) < half, axis=1)
    idx = np.where(inbox)[0]
    if len(idx) == 0:
        return None

    # Cap at max_pts
    if len(idx) > max_pts:
        idx = idx[np.random.choice(len(idx), max_pts, replace=False)]

    off = offset[idx]                                 # (M, 3) — relative xyz
    cid = class_id[idx]                               # (M,)
    src = source[idx] if source is not None else np.zeros(len(idx), np.uint8)

    # Gestalt label as normalized "RGB" (3 channels, same as 2025 winner)
    NCLASSES = 30
    g_r = (cid % 10) / 9.0
    g_g = (cid // 10) / max(NCLASSES // 10, 1)
    g_b = np.clip(cid / NCLASSES, 0, 1)
    gestalt_rgb = np.stack([g_r, g_g, g_b], axis=-1).astype(np.float32)

    # RGB color normalized to [-1, 1]
    if rgb is not None:
        color = rgb[idx].astype(np.float32) / 127.5 - 1.0
    else:
        color = np.zeros((len(idx), 3), np.float32)

    # ADE20K house flag — placeholder (1 if COLMAP, heuristic)
    is_house = (src == 0).astype(np.float32).reshape(-1, 1)

    # In-cluster flag: 1 if original vertex class
    in_cluster = np.isin(cid, VERTEX_CLASSES).astype(np.float32).reshape(-1, 1)

    feats = np.concatenate([
        off.astype(np.float32),   # 3
        color,                     # 3
        is_house,                  # 1
        gestalt_rgb,               # 3
        in_cluster,                # 1
    ], axis=-1)                    # → 11D

    return feats.T  # (11, M) for Conv1d


def extract_edge_patch(v_a, v_b, xyz_world, class_id=None, rgb=None,
                       radius=1.0, extend=1.0, max_pts=512):
    """Extract 6D features for cylindrical patch between two vertices.

    Channel layout: xyz_offset(3) + gestalt_rgb(3) — same 6D as before but
    now the 'color' channels carry semantic class info instead of zeros.
    """
    mid = (v_a + v_b) / 2
    direction = v_b - v_a
    length = np.linalg.norm(direction)
    if length < 1e-6:
        return None
    d_hat = direction / length

    rel = xyz_world - mid
    t = rel @ d_hat
    dist_along = np.abs(t)
    dist_perp = np.linalg.norm(rel - np.outer(t, d_hat), axis=1)

    inbox = (dist_along < length / 2 + extend) & (dist_perp < radius)
    idx = np.where(inbox)[0]
    if len(idx) == 0:
        return None
    if len(idx) > max_pts:
        idx = idx[np.random.choice(len(idx), max_pts, replace=False)]

    off = (xyz_world[idx] - mid).astype(np.float32)

    # Prefer gestalt class encoding (semantic, 3D) over rgb (often zeros)
    if class_id is not None:
        NCLASSES = 30
        cid = class_id[idx].astype(np.float32)
        g_r = (cid % 10) / 9.0
        g_g = (cid // 10) / max(NCLASSES // 10, 1)
        g_b = np.clip(cid / NCLASSES, 0.0, 1.0)
        color = np.stack([g_r, g_g, g_b], axis=-1).astype(np.float32)
    elif rgb is not None:
        color = rgb[idx].astype(np.float32) / 127.5 - 1.0
    else:
        color = np.zeros((len(idx), 3), np.float32)

    feats = np.concatenate([off, color], axis=-1)  # (M, 6)
    return feats.T  # (6, M)


# ---------------------------------------------------------------------------
# Collate with variable-length point sets
# ---------------------------------------------------------------------------

def collate_patches(batch, n_pts=512):
    """Pad/subsample each patch to n_pts for batching."""
    feats_list, labels = [], []
    for feats, label in batch:
        M = feats.shape[1]
        if M >= n_pts:
            idx = np.random.choice(M, n_pts, replace=False)
        else:
            idx = np.concatenate([np.arange(M),
                                   np.random.choice(M, n_pts - M, replace=True)])
        feats_list.append(feats[:, idx])
        labels.append(label)
    return (torch.from_numpy(np.stack(feats_list)).float(),
            torch.tensor(labels, dtype=torch.float32))