import torch
import math
import numpy as np

from os import sys, path
import torch.nn.functional as F
from torch.linalg import eigh
import time

def get_NGC_structure(points, raw_xyz, N2, N1, token_stage):
    B, N3, channels = points.shape
    reduce_factor = 2

    # first use nvg-style method clustering N3 to N1
    N1to3_binary, N1to3_lab2id, cur_points = clustering_points(points, B, N1, N3, reduce_factor)

    # now get N1 tokens structured
    ids = N1to3_lab2id[:, :, 0].unsqueeze(-1).expand(-1, -1, channels)
    cur_points = torch.gather(cur_points, 1, ids)
    N0to1_binary, N0to1_lab2id = build_tree(cur_points)

    time3 = time.perf_counter()

    # concat N0 -> N1 & N1 -> N3 binary structure
    N0to3_binarys = concat_N0to3_binary(N0to1_binary, N1to3_lab2id, N1to3_binary, N1)

    # get N3 -> N2 token masks
    if N2 == N3:
        structure_gt02_bin, points_new, raw_xyz_new = N0to3_binarys, points, raw_xyz
    else:
        structure_gt02_bin, points_new, raw_xyz_new = get_3to2(N0to3_binarys, N2, points, raw_xyz)

    # get N2 final decimal structure
    structure_gt02_dec = get_structure_dec(structure_gt02_bin, N2, token_stage)

    return structure_gt02_bin, structure_gt02_dec, points_new, raw_xyz_new

def get_structure_dec(structure_gt02, N2, token_stage):
    B = len(structure_gt02)
    _, len_max = structure_gt02[0].shape
    len_start = int(math.log(N2, 2)) + 1
    structure_gt02_dec = []
    for b in range(B):
        structure_gt02_b = structure_gt02[b]
        structure_gt02_b_dec = []
        len_end = len_start
        label = None
        for i in range(1, len_max):
            labels = binary_prefix_to_label(structure_gt02_b[:, :i])  # (N,)

            # first store the decimal infrom
            structure_gt02_b_dec.append(labels)
            if i < len_start:
                continue
            # judge if get the end length
            unique_labels, inverse = torch.unique(labels, return_inverse=True)
            M = unique_labels.numel()
            len_end = i
            if M == N2:
                break
        if token_stage > len_end:
            buffer = token_stage - len_end
            for _ in range(buffer):
                labels = torch.arange(N2, device=structure_gt02[0].device, dtype=torch.long)
                structure_gt02_b_dec.append(labels)
        structure_gt02_b_dec = torch.stack(structure_gt02_b_dec, dim=0).permute(1, 0)
        structure_gt02_dec.append(structure_gt02_b_dec)

    return structure_gt02_dec

def get_3to2(N0to3_binarys, N2, points, raw_xyz):
    B = len(N0to3_binarys)
    device = N0to3_binarys[0].device
    structure_gt02 = []
    points_new = []
    raw_xyz_new = []
    for b in range(B):
        binary_b03 = N0to3_binarys[b]
        N3, length = binary_b03.shape
        labels = binary_prefix_to_label(binary_b03)
        token_mask = select_diverse_tokens(labels, N2)
        assert N2 == token_mask.sum()

        binary_b02 = binary_b03[token_mask] # N2 len1+len2 
        points_new_b = points[b, :, :][token_mask] # B N3 -> N2 C 
        raw_xyz_new_b = raw_xyz[b, :, :][token_mask] # B N3 -> N2 C 
        points_new.append(points_new_b)
        raw_xyz_new.append(raw_xyz_new_b)
        structure_gt02.append(binary_b02) 

    points_new = torch.stack(points_new, dim=0)
    raw_xyz_new = torch.stack(raw_xyz_new, dim=0)
    return structure_gt02, points_new, raw_xyz_new

def select_diverse_tokens(labels, N):
    """
    labels: (8192,) int64
    return: token_mask (8192,) bool, sum == N
    """
    device = labels.device

    unique_labels, inverse = torch.unique(labels, return_inverse=True)
    M = unique_labels.numel()
    assert M >= N

    # === Step 1: ancestor-based grouping ===
    shift = max(0, int(torch.floor(torch.log2(torch.tensor(M / N))).item()))
    group_id = unique_labels >> shift
    unique_groups, group_inv = torch.unique(group_id, return_inverse=True)

    # === Step 2: select N labels, one per group first ===
    perm = torch.randperm(M, device=device)

    selected_label = torch.zeros(M, dtype=torch.bool, device=device)
    used_group = torch.zeros(unique_groups.numel(), dtype=torch.bool, device=device)

    cnt = 0
    for idx in perm:
        g = group_inv[idx]
        if not used_group[g]:
            selected_label[idx] = True
            used_group[g] = True
            cnt += 1
            if cnt == N:
                break

    # === Step 3: map label → first token ===
    token_mask = torch.zeros(labels.size(0), dtype=torch.bool, device=device)

    # inverse: token -> label index
    for lab_idx in selected_label.nonzero(as_tuple=True)[0]:
        t = (inverse == lab_idx).nonzero(as_tuple=True)[0][0]
        token_mask[t] = True

    return token_mask

def binary_prefix_to_label(prefix):
    k = prefix.size(1)
    weights = (2 ** torch.arange(k - 1, -1, -1, device=prefix.device))
    labels = (prefix * weights).sum(dim=1)
    return labels


def concat_N0to3_binary(N0to1_binary, N1to3_lab2id, N1to3_binary, N1):
    B, N3, length2 = N1to3_binary.shape
    device = N1to3_binary.device
    K = N3 // N1
    N0to3_binarys = []
    # end_indxs = []
    for b in range(B):
        binary_b01 = N0to1_binary[b] # N1 len1
        length1 = binary_b01.shape[-1]
        binary_b03 = - torch.ones(N3, length1 + length2, device=device) # N3 len1+len2
        binary_b03 = binary_b03.long()
        for i in range(N1):
            binary_b03[N1to3_lab2id[b, i, :], :length1] = binary_b01[i, :]
        mask = (binary_b03 == -1)
        has_neg1 = mask.any(dim=1)       # (N,) bool
        first_idx = torch.where(
            has_neg1,
            mask.int().argmax(dim=1),    
            torch.full((binary_b03.size(0),),
                length1,        
                device=binary_b03.device,
                dtype=torch.long)
        )
        offset = torch.arange(length2, device=device).unsqueeze(0)  # (1, len2)
        write_idx = first_idx.unsqueeze(1) + offset  # (N, len2)
        binary_b03.scatter_(1, write_idx, N1to3_binary[b, :, :])  
        binary_b03 = (binary_b03 > 0).long() # turns -1 into 0
        N0to3_binarys.append(binary_b03)

    return N0to3_binarys

def clustering_points(points, B, N1, N3, reduce_factor):
    B, N3, channels = points.shape
    device = points.device

    iter_nums = int(math.log(N3, reduce_factor) - math.log(N1, reduce_factor))
    N1to3_decimal = torch.arange(N3, device=device).unsqueeze(0).unsqueeze(0).repeat(B, iter_nums+1, 1) 
    cur_points = points.clone()
    id2lab = torch.arange(N3, device=device).unsqueeze(0).repeat(B, 1) 
    lab2id = torch.arange(N3, device=device).unsqueeze(0).repeat(B, 1).unsqueeze(-1)  
    for i in range(1, iter_nums + 2):
        N1to3_decimal[:, iter_nums+1-i, :] = id2lab
        if i == iter_nums + 1:
            break
        ids = lab2id[:, :, 0].unsqueeze(-1).expand(-1, -1, channels)
        cur_points = torch.gather(cur_points, 1, ids)
        lab2id, id2lab = pair_points(cur_points, i, lab2id, reduce_factor)
        cur_points = pr2vis(lab2id, points, channels)
    
    N1to3_decimal = reorganize(N1to3_decimal, int(math.log(N1, reduce_factor)))
    # turn decimal to binary
    N1to3_binary = N1to3_decimal[:, 1:, :] & 1
    N1to3_binary = N1to3_binary.permute(0, 2, 1) # B L len2 
    return N1to3_binary, lab2id, cur_points

def build_tree(tokens, min_ratio=0.25, k=8):
    """
    tokens: (B, L, C)
    return:
        id2lab: List of B tensors, each size (L, depth)
        lab2id: List of B lists: each layer is list of tensors of token ids
    """
    B, L, C = tokens.shape
    device = tokens.device
    # all_id2lab = []
    all_lab2id = []
    all_id2lab_binary = []

    for b in range(B):
        feats = tokens[b]     # (L, C)
        # id2lab_b = []         
        lab2id_b = []         
        id2lab_b_binary = []
        # id2lab_b.append(torch.zeros(L, dtype=torch.long, device=device))
        id2lab_b_binary.append(torch.zeros(L, dtype=torch.long, device=device))
        lab2id_b.append([torch.arange(L, device=device)])
        frontier = [torch.arange(L, device=device)]   
        depth = 1
        while True: 
            next_frontier = []
            # id2lab_level = - torch.ones(L, dtype=torch.long, device=device)
            lab2id_level = []
            id2lab_level_binary = - torch.ones(L, dtype=torch.long, device=device)
            
            for group_ids in frontier:
                if group_ids.numel() == 1:
                    ids = group_ids
                    # label_counter = id2lab_b[depth - 1][ids]
                    # id2lab_level[ids] =  2 * label_counter
                    id2lab_level_binary[ids] = -1
                    lab2id_level.append(ids)
                elif group_ids.numel() == 2:
                    # label_counter = id2lab_b[depth - 1][group_ids[0]]
                    ids0 = group_ids[0]
                    ids1 = group_ids[1]
                    # id2lab_level[ids0] =  2 * label_counter
                    lab2id_level.append(ids0)
                    id2lab_level_binary[ids0] = 0
                    # id2lab_level[ids1] = 2 * label_counter + 1
                    lab2id_level.append(ids1)
                    id2lab_level_binary[ids1] = 1
                    next_frontier.append(ids0)
                    next_frontier.append(ids1)      
                else:
                    # label_counter = id2lab_b[depth - 1][group_ids[0]] #
                    subfeat = feats[group_ids]      
                    A = build_knn_graph(subfeat, k=k)
                    part0, part1 = constrained_bipartition(A, min_ratio=min_ratio)

                    ids0 = group_ids[part0]
                    ids1 = group_ids[part1]
                    # id2lab_level[ids0] =  2 * label_counter
                    id2lab_level_binary[ids0] = 0
                    lab2id_level.append(ids0)
                    # id2lab_level[ids1] = 2 * label_counter + 1
                    id2lab_level_binary[ids1] = 1
                    lab2id_level.append(ids1)
                    next_frontier.append(ids0)
                    next_frontier.append(ids1)

            # id2lab_b.append(id2lab_level)
            lab2id_b.append(lab2id_level)
            id2lab_b_binary.append(id2lab_level_binary)

            if len(next_frontier) == 0:
                break
            frontier = next_frontier
            depth += 1

        all_id2lab_binary.append(torch.stack(id2lab_b_binary, dim=1))
        # all_id2lab.append(torch.stack(id2lab_b, dim=1))  # (L, depth)
        all_lab2id.append(lab2id_b)

    return all_id2lab_binary, all_lab2id

def pairwise_dist(x):
    # x: (N, C)
    # returns (N, N)
    diff = x.unsqueeze(1) - x.unsqueeze(0)
    return (diff * diff).sum(-1)

def build_knn_graph(feature, k=8):
    # feature: (N, C)
    N = feature.size(0)
    if N <= 1:
        return torch.zeros(N, N, device=feature.device)
    k_eff = min(k, N - 1)
    dist = pairwise_dist(feature)             # (N, N)
    knn_idx = dist.topk(k_eff + 1, dim=1, largest=False).indices[:, 1:]  

    N = feature.size(0)
    A = torch.zeros(N, N, device=feature.device)
    for i in range(N):
        A[i, knn_idx[i]] = 1

    A = torch.maximum(A, A.T)
    return A

def spectral_bipartition(A):
    # A: (N, N)
    N = A.size(0)
    D = A.sum(dim=1)
    L = torch.diag(D) - A
    vals, vecs = eigh(L)
    fv = vecs[:, 1]   # Fiedler vector
    part0 = (fv > 0)
    part1 = ~part0

    return part0, part1

def constrained_bipartition(A, min_ratio=0.25):
    part0, part1 = spectral_bipartition(A)
    N = A.size(0)

    min_size = int(N * min_ratio)

    if part0.sum() < min_size or part1.sum() < min_size:
        vals, vecs = eigh(torch.diag(A.sum(1)) - A)
        fv = vecs[:, 1]
        sorted_ids = torch.argsort(fv)

        part0 = torch.zeros(N, dtype=torch.bool, device=A.device)
        part1 = torch.zeros(N, dtype=torch.bool, device=A.device)

        part0[sorted_ids[:min_size]] = True
        part1[sorted_ids[min_size:]] = True

    return part0, part1

def reorganize(cluster_tensor, start_stage):
    reorganized = torch.zeros_like(cluster_tensor)
    reorganized[:, :1, :] = cluster_tensor[:, :1, :]
    B, L, N = cluster_tensor.shape
    for level in range(1, L):
        parent2token = reorganized[:, level - 1]
        child2token = cluster_tensor[:, level]

        children = torch.empty(B, 2**(level + start_stage), device=cluster_tensor.device, dtype=cluster_tensor.dtype)
        children.scatter_(1, child2token, parent2token)

        current_parent2child = torch.sort(children, dim=1, stable=True)[1]

        # Simplified inner loop
        new_indices = torch.empty_like(current_parent2child)
        new_indices.scatter_(1, current_parent2child, torch.arange(current_parent2child.size(1), device=cluster_tensor.device).unsqueeze(0).repeat(B, 1))
        reorganized[:, level] = torch.gather(new_indices, 1, child2token)

    return reorganized

def pair_points(points_batch, iter_num, lab2id, reduce_factor):

    batch_size = points_batch.size(0)
    num_points = points_batch.size(1)
    device = points_batch.device

    inf_value = 1e8

    distance_matrix = torch.cdist(points_batch, points_batch, p=2)
    distance_matrix += torch.eye(num_points, device=device).unsqueeze(0).expand(batch_size, -1, -1) * inf_value

    cluster_num = reduce_factor ** iter_num
    new_lab2id = torch.empty((batch_size, num_points // reduce_factor, cluster_num), dtype=torch.long, device=device)
    new_id2lab = torch.empty((batch_size, num_points * cluster_num // reduce_factor), dtype=torch.long, device=device)

    for match_id in range(num_points // reduce_factor):
        row_indices = torch.arange(batch_size, device=device)

        _, min_idx = torch.min(distance_matrix.view(batch_size, -1), dim=1)
        i = min_idx // num_points
        j = min_idx % num_points

        ids_i = lab2id[row_indices, i]
        ids_j = lab2id[row_indices, j]

        new_lables = torch.ones(batch_size, cluster_num // reduce_factor, dtype=torch.long, device=device) * match_id
        new_id2lab.scatter_(1, ids_i, new_lables)
        new_id2lab.scatter_(1, ids_j, new_lables)


        new_lab2id[:, match_id, :(cluster_num // reduce_factor)] = ids_i
        new_lab2id[:, match_id, (cluster_num // reduce_factor):] = ids_j


        distance_matrix[row_indices, i, :] = inf_value
        distance_matrix[row_indices, :, i] = inf_value
        distance_matrix[row_indices, j, :] = inf_value
        distance_matrix[row_indices, :, j] = inf_value

    return new_lab2id, new_id2lab

def pr2vis(lab2id, colors, channels):
    # pair results -> colors
    bs, label_num, cluster_num = lab2id.shape
    device = lab2id.device

    new_colors = torch.zeros(bs, label_num * cluster_num, channels, device=device).to(colors.dtype)

    for i in range(label_num):
        ids = lab2id[:, i, :].unsqueeze(-1).expand(-1, -1, channels)
        target_color = torch.gather(colors, 1, ids)
        target_color = torch.mean(target_color, 1)
        new_colors.scatter_(1, ids, target_color.unsqueeze(1).repeat(1, cluster_num, 1))

    return new_colors