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)
        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)
            new_labels = sim.argmax(dim=1).to(torch.int32)

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


# ============================  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
















# """
# KBC clustering — Kernel-Bounded Clustering: Achieving the Objective of Spectral Clustering without Eigendecomposition in Artificial Intelligence Journal (AIJ) 2025.
# Author: Hang Zhang
# """

# 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

# # Suppress warnings for cleaner output
# warnings.filterwarnings('ignore')


# 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: torch.device = 'None'):
#         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 == 'None':
#             self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

#         # Ensure psi is at least 2 to avoid degenerate radius
#         psi = max(2, min(self.psi, n))
#         ik = _IsolationKernel(t=self.t, psi=psi, device=self.device)
#         ik.fit(X)
#         self.ndata = ik.transform(X, batch_size=self.batch_size)  # (n, t*psi), int8

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

#         # Move subset to GPU for kernel computation
#         sub = torch.tensor(self.ndata[sID], dtype=torch.float32, device=self.device)  # (s, t*psi)

#         # Compute kernel on subset
#         K = sub @ sub.T  # (s, s)
#         tau_s = self.tau * float(K.max().item())
#         Kg = (K >= tau_s).int()

#         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})."
#             )

#         # Select top-k largest components
#         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

#         # Assign unassigned points in subset to nearest centroid (by IK similarity)
#         centers = torch.stack([
#             sub[sub2full == cid].mean(dim=0) for cid in range(self.k)
#         ])  # (k, t*psi)

#         unassigned = (sub2full == -1)
#         # if unassigned.any():
#         #     similarity = sub[unassigned] @ centers.T  # (n_unassigned, k)
#         #     sub2full[unassigned] = similarity.argmax(dim=1).to(torch.int32)

#         if unassigned.any():
#             chunk_size = 10000
#             unassigned_indices = torch.where(unassigned)[0]
#             for i in trange(0, len(unassigned_indices), chunk_size, desc="Assigning data to the distribution", leave=False):
#                 idx_chunk = unassigned_indices[i:i + chunk_size]
#                 sub_chunk = sub[idx_chunk] 
#                 sim_chunk = sub_chunk @ centers.T   
#                 best_center = sim_chunk.argmax(dim=1).to(torch.int32)
#                 sub2full[idx_chunk] = best_center


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

#         # Optional refinement via k-means in IK space
#         if self.post_processing:
#             ndata_tensor = torch.tensor(self.ndata, dtype=torch.float32, device=self.device)
#             full_labels = self._refine_gpu(ndata_tensor, full_labels, self.k)

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

#     def fit_predict(self, X, y=None):
#         """Return cluster labels."""
#         return self.fit(X, y).labels_

#     @torch.no_grad()
#     def _refine_gpu(self, ndata, labels, k, max_iter=100, tol=0.01):
#         """
#         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.
#         """
#         n = ndata.shape[0]
#         th = max(1, int(math.ceil(tol * n)))
#         for _ in range(max_iter):
#             # Recompute centroids
#             centers = torch.stack([
#                 ndata[labels == c].mean(dim=0) for c in range(k)
#             ])  # (k, d_ik)

#             # Assign to most similar centroid (max inner product)
#             similarity = ndata @ centers.T  # (n, k)
#             new_labels = similarity.argmax(dim=1).to(labels.dtype)  # 0-based

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


# # ===================================================================
# # Isolation-Kernel GPU implementation (internal use only) @Yi-xiao Ma
# # ===================================================================
# class _IsolationKernel:
#     """Minimal GPU Isolation-Kernel transformer."""

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

#     def fit(self, X: np.ndarray):
#         self.X = X
#         n_samples = X.shape[0]

#         if self._psi > n_samples:
#             self._psi = n_samples
#             warnings.warn(f"psi is set to {n_samples} "
#                           "as it is greater than the number of data points.")

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

#         self._center_list = np.zeros((self._t * self._psi, X.shape[1]), dtype=X.dtype)
#         self._radius_list = torch.zeros((self._t, self._psi), dtype=torch.float32, device=self.device)

#         for i in range(self._t):
#             sample = X[self._center_index_list[i]]
#             self._center_list[i * self._psi:(i + 1) * self._psi] = sample

#             sample_cuda = torch.tensor(sample, dtype=torch.float32, device=self.device)
#             s2s = torch.cdist(sample_cuda, sample_cuda, p=2)
#             s2s.fill_diagonal_(float('inf'))
#             self._radius_list[i] = torch.min(s2s, dim=0).values
#         return self

#     def transform(self, X: np.ndarray, batch_size: int = 10000):
#         if X.ndim == 1:
#             X = X.reshape(1, -1)

#         n_samples = X.shape[0]
#         output = torch.zeros((n_samples, self._psi * self._t), device='cpu', dtype=torch.int4)

#         for start in trange(0, n_samples, batch_size, desc="Transforming data to Kernel space", leave=False):
#             end = min(start + batch_size, n_samples)
#             batch_cuda = torch.tensor(X[start:end], dtype=torch.float32, device=self.device)

#             for i in range(self._t):
#                 sl = slice(i * self._psi, (i + 1) * self._psi)
#                 sample_cuda = torch.tensor(self._center_list[sl], dtype=torch.float32, device=self.device)

#                 p2s = torch.cdist(batch_cuda, sample_cuda, p=2)
#                 p2ns_idx = torch.argmin(p2s, dim=1)
#                 p2ns = p2s[torch.arange(p2ns_idx.size(0), device=self.device), p2ns_idx]
#                 accept = p2ns <= self._radius_list[i, p2ns_idx]

#                 col_idx = p2ns_idx[accept] + i * self._psi
#                 output[start:end][accept, col_idx] = 1

#         # Always return 2D array for consistency
#         return output.numpy()