import numpy as np
from sklearn.neighbors import KDTree
from scipy.spatial.distance import pdist, squareform
from typing import Optional, List, Tuple
import warnings
warnings.filterwarnings("ignore")

class MCMSTClustering:
    """
    MCMSTClustering: KD-Tree + MST based clustering algorithm for arbitrary-shaped clusters.
    
    Parameters
    ----------
    N : int, default=5
        Minimum number of data points to define a Micro Cluster.
    r : float, default=0.05
        Radius of the Micro Cluster.
    n_micro : int, default=5
        Minimum number of Micro Clusters to define a Macro Cluster.
    """
    
    def __init__(self, N: int = 5, r: float = 0.05, n_micro: int = 5):
        self.N = N
        self.r = r
        self.n_micro = n_micro
        self.labels_: Optional[np.ndarray] = None
        self._MCs: Optional[np.ndarray] = None
        self._processed_data: Optional[np.ndarray] = None
        self._MacroClusters: 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.
        """
        # Initialize data structures
        self._XX = X.copy()
        self._d = X.shape[1]
        self._MC_Num = 0
        self._MacroC_Num = 0
        
        # Create processed data array: [id, MC_id, MacroC_id, features...]
        self._processed_data = np.hstack((
            np.arange(X.shape[0]).reshape(-1, 1) + 1,
            np.zeros((X.shape[0], 2)),
            X
        ))
        
        self._MCs = np.empty((0, self._d + 3), dtype=float)
        self._MacroClusters = np.empty((0, 4), dtype=object)
        
        # Main algorithm steps
        self._DefineMC()
        self._RegulateClusters()
        self._DefMacros()
        self._assignMacroN()
        
        return self
        
    def _DefineMC(self):
        """Define Micro Clusters using KD-Tree range search."""
        while True:
            # Get unassigned data points
            X = self._processed_data[self._processed_data[:, 1] == 0, :]
            num_MCs_before = self._MC_Num
            
            if X.shape[0] >= self.N:
                tree = KDTree(X[:, 3:])
                
                for i in range(X.shape[0]):
                    if X[i, 2] == 0:  # Not yet in any macro cluster
                        # Range search
                        indices = tree.query_radius(X[i:i+1, 3:], self.r)[0]
                        points = X[indices, 3:]
                        
                        if len(points) >= self.N:
                            # Create new micro cluster
                            center = np.mean(points, axis=0)
                            self._MC_Num += 1
                            
                            # Add to MCs array: [MC_id, num_points, MacroC_id, center...]
                            new_mc = np.hstack((
                                np.array([self._MC_Num, len(points), 0]),
                                center
                            ))
                            self._MCs = np.vstack((self._MCs, new_mc))
                            
                            # Assign points to this micro cluster
                            for point in points:
                                idx = np.where((self._processed_data[:, 3:] == point).all(axis=1))[0]
                                if len(idx) > 0:
                                    self._processed_data[idx[0], 1] = self._MC_Num
                            break
            
            # Stop if no new micro clusters were created
            if num_MCs_before == self._MC_Num:
                break
    
    def _MST(self, mc_id: int) -> List[Tuple[int, int]]:
        """
        Compute Minimum Spanning Tree for micro clusters using Prim's algorithm.
        Edges longer than 2*r are removed (incomplete graph).
        """
        # Get micro clusters not assigned to any macro cluster
        P = self._MCs[self._MCs[:, 2] == 0, :]
        if len(P) == 0:
            return []
            
        indices = P[:, 0]
        # Compute distance matrix between micro cluster centers
        distances = squareform(pdist(P[:, 3:]))
        
        # Remove edges longer than 2*r (create incomplete graph)
        distances[distances > 2 * self.r] = 0
        
        N = len(P)
        selected = np.zeros(N, dtype=bool)
        
        # Find starting node (the one with given mc_id)
        start_idx = np.where(indices == mc_id)[0]
        if len(start_idx) == 0:
            return []
        selected[start_idx[0]] = True
        
        edges = []
        num_edges = 0
        
        while num_edges < N - 1:
            min_dist = np.inf
            a = -1
            b = -1
            
            for m in range(N):
                if selected[m]:
                    for n in range(N):
                        if not selected[n] and distances[m, n] > 0 and distances[m, n] < min_dist:
                            min_dist = distances[m, n]
                            a = m
                            b = n
            
            if a == -1 or b == -1:  # No more connections
                break
                
            edges.append((int(indices[a]), int(indices[b])))
            selected[b] = True
            num_edges += 1
            
        return edges
    
    def _DefMacros(self):
        """Define Macro Clusters from micro clusters using MST."""
        for mc in self._MCs:
            if mc[2] == 0:  # Micro cluster not assigned to any macro cluster
                edges = self._MST(mc[0])
                
                if len(edges) > 0:
                    # Flatten edges to get all connected micro clusters
                    all_edges = np.array(edges).flatten()
                    unique_mcs = np.unique(all_edges)
                    
                    # Count total points in these micro clusters
                    total_points = 0
                    for mc_id in unique_mcs:
                        mc_data = self._MCs[self._MCs[:, 0] == mc_id]
                        if len(mc_data) > 0:
                            total_points += mc_data[0, 1]
                    
                    # Check if enough points or enough micro clusters
                    if (total_points >= self.n_micro * self.N or 
                        len(unique_mcs) >= self.n_micro):
                        
                        self._MacroC_Num += 1
                        
                        # Create macro cluster record
                        macro_cluster = np.array([
                            self._MacroC_Num,
                            len(unique_mcs),
                            edges,
                            None  # Color placeholder
                        ], dtype=object)
                        
                        self._MacroClusters = np.vstack((self._MacroClusters, macro_cluster))
                        
                        # Assign micro clusters to this macro cluster
                        for mc_id in unique_mcs:
                            idx = np.where(self._MCs[:, 0] == mc_id)[0]
                            if len(idx) > 0:
                                self._MCs[idx[0], 2] = self._MacroC_Num
        
        # Handle large micro clusters that didn't get connected
        for mc in self._MCs:
            if mc[2] == 0 and mc[1] >= self.N * self.n_micro:
                self._MacroC_Num += 1
                idx = np.where(self._MCs[:, 0] == mc[0])[0]
                if len(idx) > 0:
                    self._MCs[idx[0], 2] = self._MacroC_Num
                    
    def _RegulateClusters(self):
        """Regulate clusters: assign orphan points to nearest micro cluster within 2*r."""
        # Get unassigned points
        unassigned = self._processed_data[self._processed_data[:, 1] == 0, :]
        
        if len(unassigned) > 0 and len(self._MCs) > 0:
            tree = KDTree(self._MCs[:, 3:])
            
            for point in unassigned:
                distance, idx = tree.query(point[3:].reshape(1, -1))
                if distance[0] <= 2 * self.r:
                    mc_id = self._MCs[idx[0], 0]
                    point_id = point[0]
                    self._processed_data[self._processed_data[:, 0] == point_id, 1] = mc_id
    
    def _assignMacroN(self):
        """Assign macro cluster labels to all data points."""
        for mc in self._MCs:
            mc_id = mc[0]
            macro_id = mc[2]
            # Assign macro cluster id to all points in this micro cluster
            self._processed_data[self._processed_data[:, 1] == mc_id, 2] = macro_id
        
        self.labels_ = self._processed_data[:, 2].astype(int)
    
    def predict(self, X: np.ndarray) -> np.ndarray:
        """
        Predict cluster labels for new data.
        
        Parameters
        ----------
        X : np.ndarray, shape (n_samples, n_features)
            New data to predict.
            
        Returns
        -------
        labels : np.ndarray, shape (n_samples,)
            Cluster labels for each data point.
        """
        if self._MCs is None:
            raise ValueError("Model not fitted yet. Call `fit` first.")
        
        if len(self._MCs) == 0:
            return np.zeros(X.shape[0], dtype=int)
        
        tree = KDTree(self._MCs[:, 3:])
        labels = np.zeros(X.shape[0], dtype=int)
        
        for i in range(X.shape[0]):
            distance, idx = tree.query(X[i:i+1])
            if distance[0] <= self.r:
                mc_id = self._MCs[idx[0], 0]
                # Find which macro cluster this micro cluster belongs to
                mc_data = self._MCs[self._MCs[:, 0] == mc_id]
                if len(mc_data) > 0:
                    labels[i] = int(mc_data[0, 2])
        
        return labels
    
    def fit_predict(self, X: np.ndarray) -> np.ndarray:
        """
        Fit the model and return cluster labels.
        
        Parameters
        ----------
        X : np.ndarray, shape (n_samples, n_features)
            Input data.
            
        Returns
        -------
        labels : np.ndarray, shape (n_samples,)
            Cluster labels for each data point.
        """
        self.fit(X)
        return self.labels_