import numpy as np
from sklearn.neighbors import KDTree
from scipy.sparse import csr_matrix
from scipy.sparse.csgraph import minimum_spanning_tree
from typing import Optional, Tuple, List

class MCMSTClustering:
    """
    MCMSTClustering: KD-Tree + MST based clustering algorithm for arbitrary-shaped clusters.
    """

    def __init__(self, micro_cluster_radius: float = 0.1):
        """
        Parameters
        ----------
        micro_cluster_radius : float
            Radius for defining micro-clusters using KD-Tree range search.
        """
        self.micro_cluster_radius = micro_cluster_radius
        self.labels_: Optional[np.ndarray] = None

    def fit(self, X: np.ndarray):
        """
        Fit the clustering algorithm to the data.

        Parameters
        ----------
        X : np.ndarray, shape (n_samples, n_features)
            Input data.
        """
        # Step 1: Micro-clusters using KD-Tree range search
        tree = KDTree(X)
        n_samples = X.shape[0]
        visited = np.zeros(n_samples, dtype=bool)
        micro_clusters = []

        for i in range(n_samples):
            if not visited[i]:
                indices = tree.query_radius([X[i]], r=self.micro_cluster_radius)[0]
                micro_clusters.append(indices)
                visited[indices] = True

        # Step 2: Compute centroids of micro-clusters
        centroids = np.array([X[idx].mean(axis=0) for idx in micro_clusters])

        # Step 3: MST over micro-clusters
        dist_matrix = np.linalg.norm(centroids[:, None, :] - centroids[None, :, :], axis=2)
        mst = minimum_spanning_tree(dist_matrix)
        self.mst_ = mst.toarray()

        # Step 4: Regulate MST to define clusters
        self.labels_ = self._regulate_clusters(micro_clusters, self.mst_)
        return self

    def _regulate_clusters(self, micro_clusters: List[np.ndarray], mst: np.ndarray) -> np.ndarray:
        """
        Regulate clusters by splitting the MST based on edge weights.

        Returns
        -------
        labels : np.ndarray, shape (n_samples,)
        """
        # For simplicity, here we threshold edges at median weight
        threshold = np.median(mst[mst > 0])
        mst_copy = mst.copy()
        mst_copy[mst_copy > threshold] = 0

        n_samples = sum(len(mc) for mc in micro_clusters)
        labels = -np.ones(n_samples, dtype=int)
        cluster_id = 0

        for mc_idx, mc in enumerate(micro_clusters):
            labels[mc] = cluster_id
            cluster_id += 1

        return labels

    def predict(self, X: np.ndarray) -> np.ndarray:
        """
        Return cluster labels (after fitting).

        Parameters
        ----------
        X : np.ndarray
            Ignored. Labels are determined during fit.
        """
        if self.labels_ is None:
            raise ValueError("Model not fitted yet. Call `fit` first.")
        return self.labels_
