import warnings
import numpy as np
import torch
import math
from tqdm import trange
from sklearn.base import BaseEstimator, ClusterMixin
from sklearn.utils import check_array, check_random_state
from scipy.sparse.csgraph import connected_components
from scipy.sparse import csr_matrix


class KBC(BaseEstimator, ClusterMixin):
    """
    KBC: Isolation-Kernel + Binary Connected-component Clustering

    Parameters
    ----------
    k : int
        Number of clusters.
    tau : float
        Threshold for binarising the subgraph (relative to max kernel value).
    psi : int
        Number of centroids per isolation-tree.
    t : int, default=100
        Number of isolation trees.
    subsample_size : int, default=10_000
        Size of subset used to build the kernel matrix.
    batch_size : int, default=10_000
        GPU batch size for IK transform.
    random_state : int or None, default=None
        Random seed for reproducibility.
    post_processing : bool, default=False
        Refine labels with k-means iterations after graph clustering.

    Attributes
    ----------
    labels_ : ndarray of shape (n_samples,)
        Cluster labels (0-based).
    """
    def __init__(self, k: int, tau: float, psi: int, t: int = 100,
                 subsample_size: int = 10_000, batch_size: int = 10_000,
                 random_state=None, post_processing: bool = True, device: str = 'auto'):
        self.k = k
        self.tau = tau
        self.psi = psi
        self.t = t
        self.subsample_size = subsample_size
        self.batch_size = batch_size
        self.random_state = random_state
        self.post_processing = post_processing
        self.device = device

    def fit(self, X, y=None):
        """
        Perform clustering.

        Parameters
        ----------
        X : array-like of shape (n_samples, n_features)
            Training data.
        y : Ignored
            Not used, present for sklearn API consistency.

        Returns
        -------
        self : object
        """
        X = check_array(X, dtype=np.float32, order='C')
        n, d = X.shape

        if self.device == 'auto':
            self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

        psi = max(2, min(self.psi, n))
        ik = _IsolationKernel(t=self.t, psi=psi, device=self.device)
        ik.fit(X)

        # 
        self.ndata_indices = ik.transform_indices(X, batch_size=self.batch_size)  # (n, t), uint8

        rng = check_random_state(self.random_state)
        s = min(self.subsample_size, n)
        sID = rng.choice(n, s, replace=False)

        # 
        sub_indices = torch.tensor(self.ndata_indices[sID], dtype=torch.uint8, device=self.device)  # (s, t)

        # 
        # K = torch.eq(sub_indices[:, None, :], sub_indices[None, :, :]).sum(dim=2)  # (s,s)
        K = self._compute_intersection_tiled(sub_indices, sub_indices, tile_size=1000)  # (s,s)
        tau_s = self.tau * float(K.max().item())
        Kg = (K >= tau_s)

        n_comp, labels_sub = connected_components(
            csr_matrix(Kg.cpu().numpy()), directed=False
        )
        labels_sub = torch.from_numpy(labels_sub).to(self.device, dtype=torch.int32)

        if n_comp < self.k:
            raise ValueError(
                f"KBC found only {n_comp} connected components, "
                f"but k={self.k} was requested. Try increasing tau (current={self.tau:.4f})."
            )

        # 
        comp_size = np.bincount(labels_sub.cpu().numpy(), minlength=n_comp)
        topk_idx = np.argsort(comp_size)[-self.k:][::-1]

        sub2full = torch.full((s,), -1, dtype=torch.int32, device=self.device)
        for new_id, old_id in enumerate(topk_idx):
            sub2full[labels_sub == old_id] = new_id

        # 
        if (sub2full == -1).any():
            centers = torch.stack([
                sub_indices[sub2full == cid].float().mean(dim=0) for cid in range(self.k)
            ]).to(torch.uint8)  # (k, t)

            unassigned = (sub2full == -1)
            unassigned_idx = torch.where(unassigned)[0]

            chunk = 10000
            for i in trange(0, len(unassigned_idx), chunk, desc="Assigning unassigned in subset", leave=False):
                idx = unassigned_idx[i:i+chunk]
                sim = torch.eq(sub_indices[idx][:, None, :], centers[None, :, :]).sum(dim=2)  # (chunk, k)
                sub2full[idx] = sim.argmax(dim=1).to(torch.int32)

        # 
        full_labels = torch.full((n,), -1, dtype=torch.int32, device=self.device)
        full_labels[sID] = sub2full

        # 
        if self.post_processing:
            indices_tensor = torch.tensor(self.ndata_indices, dtype=torch.uint8, device=self.device)
            full_labels = self._refine_gpu(indices_tensor, full_labels, self.k)

        self.labels_ = full_labels.cpu().numpy()
        return self

    def fit_predict(self, X, y=None):
        return self.fit(X, y).labels_

    @torch.no_grad()
    def _refine_gpu(self, indices, init_labels, k, max_iter=50, tol=0.001):
        """
        Refine labels using k-means-style updates in the Isolation Kernel space.

        Parameters
        ----------
        ndata : torch.Tensor, shape (n_samples, n_features_ik)
            Data in Isolation Kernel space (dense float32).
        labels : torch.Tensor, shape (n_samples,)
            Current cluster labels (0-based).
        k : int
            Number of clusters.
        max_iter : int
            Maximum number of refinement iterations.
        tol : float
            Convergence threshold (fraction of points that can change).

        Returns
        -------
        new_labels : torch.Tensor, shape (n_samples,)
            Refined labels.
        """
        labels = init_labels.clone()
        n = indices.shape[0]
        th = max(1, int(math.ceil(tol * n)))

        for _ in range(max_iter):
            # 
            centers = torch.stack([
                torch.mode(indices[labels == c], dim=0).values for c in range(k)
            ]).to(torch.uint8)  # (k, t)

            # 
            # sim = torch.eq(indices[:, None, :], centers[None, :, :]).sum(dim=2)  # (n, k)
            sim = self._compute_intersection_tiled(indices, centers, tile_size=1000) 
            new_labels = sim.argmax(dim=1).to(torch.int32)

            changed = (new_labels != labels).sum().item()
            if changed < th:
                break
            labels = new_labels
        return labels

    def _compute_intersection_tiled(self, indices_1, indices_2, tile_size=1000):
        """
       Calculate the intersection matrix K[i,j] = |indices_1[i] ∩ indices_2[j]|
        Avoid explicitly creating an N×M×L tensor by using block partitioning.
        """
        N, _ = indices_1.shape  
        M, _ = indices_2.shape  # 
        
        
        K = torch.zeros(N, M, device=indices_1.device, dtype=torch.long)
        
        n_tiles_1 = (N + tile_size - 1) // tile_size  # indices_1
        n_tiles_2 = (M + tile_size - 1) // tile_size  # indices_2
        
        for i in range(n_tiles_1):
            i_start = i * tile_size
            i_end = min((i + 1) * tile_size, N)
            block_i = indices_1[i_start:i_end]  # (tile_i, L)
            
            for j in range(n_tiles_2):  # n_tiles_2
                j_start = j * tile_size
                j_end = min((j + 1) * tile_size, M)  # 
                block_j = indices_2[j_start:j_end]  # (tile_j, L)
                
                #  = tile_i × tile_j × L
                tile = torch.eq(
                    block_i[:, None, :],  # (tile_i, 1, L)
                    block_j[None, :, :]   # (1, tile_j, L)
                ).sum(dim=2)  # (tile_i, tile_j)
                
                K[i_start:i_end, j_start:j_end] = tile
        
        return K


# ============================  IsolationKernel ============================
class _IsolationKernel:
    """GPU-accelerated Isolation Kernel returning partition indices (n, t)"""

    def __init__(self, t: int, psi: int, device: torch.device):
        self.t = t
        self.psi = psi
        self.device = device
        self._center_index_list = None  # (t, psi)
        self._radius_list = None       # (t, psi) on GPU
        self.X = None

    def fit(self, X: np.ndarray):
        self.X = X
        n_samples = X.shape[0]
        psi = min(self.psi, n_samples)
        if psi != self.psi:
            self.psi = psi
            warnings.warn(f"psi reduced to {psi} (n_samples limit)")

        rng = np.random.default_rng()
        self._center_index_list = np.vstack([
            rng.choice(n_samples, size=psi, replace=False) for _ in range(self.t)
        ])  # (t, psi)

        #
        self._radius_list = torch.zeros((self.t, psi), dtype=torch.float32, device=self.device)
        centers_all = X[self._center_index_list]  # (t, psi, d)

        for i in range(self.t):
            sample = torch.tensor(centers_all[i], dtype=torch.float32, device=self.device)
            dist = torch.cdist(sample, sample, p=2)
            dist.fill_diagonal_(float('inf'))
            self._radius_list[i] = dist.min(dim=0).values
        return self

    @torch.no_grad()
    def transform_indices(self, X: np.ndarray, batch_size: int = 10000) -> np.ndarray:
        """
        Returns a uint8 array of (n_samples, t), where each column is the partition number (0~psi-1) in the t-th tree where the point falls.
        """
        if X.ndim == 1:
            X = X.reshape(1, -1)
        n = X.shape[0]
        indices = np.zeros((n, self.t), dtype=np.uint8)

        X_cuda = torch.tensor(X, dtype=torch.float32, device=self.device)

        for start in trange(0, n, batch_size, desc="IK Indexing", leave=False):
            end = min(start + batch_size, n)
            batch = X_cuda[start:end]

            for i in range(self.t):
                sl = slice(i * self.psi, (i + 1) * self.psi)
                centers_i = torch.tensor(self.X[self._center_index_list[i]], 
                                       dtype=torch.float32, device=self.device)  # (psi, d)

                dist = torch.cdist(batch, centers_i, p=2)          # (batch, psi)
                min_dist, nn_idx = dist.min(dim=1)                # (batch,)

                # Determine if it falls within the radius of the nearest neighbor voronoi cell.
                radius = self._radius_list[i][nn_idx]              # (batch,)
                accept = min_dist <= radius

                # Only assign values ​​to the points that are accepted (theoretically, almost all points are accepted, as defined by IK).
                indices[start:end, i][accept.cpu().numpy()] = nn_idx[accept].cpu().numpy().astype(np.uint8)

        return indices














