import os
import sys
import time
import numpy as np
from tqdm import tqdm
import matplotlib.pyplot as plt
from scipy.spatial.distance import squareform
from scipy.cluster import hierarchy

def _check_umap_available():
    try:
        import umap
    except ImportError:
        raise ImportError(
            "UMAP is required for this functionality. "
            "Install it with: pip install umap-learn"
        )
    return umap

def _check_phate_available():
    try:
        import phate
    except ImportError:
        raise ImportError(
            "PHATE is required for this functionality. "
            "Install it with: pip install phate"
        )
    return phate

def _check_tsne_available():
    try:
        from openTSNE import TSNE
    except ImportError:
        raise ImportError(
            "t-SNE requires the 'openTSNE' package. "
            "Install it with: pip install openTSNE"
        )
    return TSNE

def _check_sklearn_available():
    try:
        from sklearn.cluster import KMeans, DBSCAN
        from sklearn.decomposition import PCA
    except ImportError:
        raise ImportError(
            "scikit-learn is required for this functionality. "
            "Install it with: pip install scikit-learn"
        )
    return KMeans, DBSCAN, PCA

class Clusterer:
    """
    Clusterer: A unified interface for embedding and clustering attribution maps
    
    This implementation provides implementations of common embedding 
    and clustering methods for attribution maps:
    
    Embedding Methods:
    - UMAP (requires umap-learn)
    - PHATE (requires phate)
    - t-SNE (requires openTSNE)
    - PCA (requires scikit-learn)
    - Diffusion Maps (not yet implemented)

    Clustering Methods:
    - Hierarchical (GPU-optimized available)
    - K-means (requires scikit-learn)
    - DBSCAN (requires scikit-learn)
    
    Requirements:
    - numpy
    - scipy
    - scikit-learn (for PCA, K-means, DBSCAN)
    
    Optional Requirements:
    - tensorflow (for GPU-accelerated hierarchical clustering)
    - umap-learn (for UMAP)
    - phate (for PHATE)
    - openTSNE (for t-SNE)

    Additional Requirements:
    - scikit-learn (for clustering)
    - matplotlib (for visualization)
    
    Example usage:
        # Initialize clusterer with attribution maps
        clusterer = Clusterer(
            maps,
            method='umap',
            n_components=2
        )
        
        # Compute embedding
        embedding = clusterer.embed()
        
        # For K-means or DBSCAN:
        clusters = clusterer.cluster(embedding, method='kmeans', n_clusters=10)
        
        # For hierarchical clustering:
        linkage = clusterer.cluster(method='hierarchical')
        # Then get cluster labels using different criteria:
        labels = clusterer.get_cluster_labels(linkage, criterion='distance', max_distance=8)
        # or
        labels, cut_level = clusterer.get_cluster_labels(linkage, criterion='maxclust', n_clusters=100)
    """
    
    SUPPORTED_METHODS = {'umap', 'phate', 'tsne', 'pca', 'diffmap'}
    SUPPORTED_CLUSTERERS = {'hierarchical', 'kmeans', 'dbscan'}

    def __init__(self, attribution_maps, method='umap', n_components=2, gpu=True):
        """Initialize the Clusterer.
        
        Args:
            attribution_maps: numpy array of shape (N, L, A) containing attribution maps
            method: Embedding method (default: 'umap')
            n_components: Number of dimensions for embedding (default: 2)
            gpu: Whether to use GPU acceleration when available (default: True)
        """
        if method not in self.SUPPORTED_METHODS:
            raise ValueError(f"Method must be one of {self.SUPPORTED_METHODS}")
        
        self.embedding = None
        self.cluster_labels = None
        self.maps = attribution_maps
        self.method = method
        self.n_components = n_components
        self.gpu = gpu
        
        # Reshape maps if needed
        if len(self.maps.shape) == 3:
            N, L, A = self.maps.shape
            self.maps = self.maps.reshape((N, L*A))
            
    def embed(self, **kwargs):
        """Compute embedding using specified method.
        
        Args:
            **kwargs: Method-specific parameters
            
        Returns:
            numpy.ndarray: Embedded coordinates
        """
        t0 = time.time()
        
        if self.method == 'umap':
            embedding = self._embed_umap(**kwargs)
        elif self.method == 'phate':
            embedding = self._embed_phate(**kwargs)
        elif self.method == 'tsne':
            embedding = self._embed_tsne(**kwargs)
        elif self.method == 'pca':
            embedding = self._embed_pca(**kwargs)
            
        print(f'Embedding time: {time.time() - t0:.2f}s')
        return embedding
    
    def _embed_umap(self, **kwargs):
        """Compute UMAP embedding with optional GPU acceleration."""
        umap = _check_umap_available()
        
        if self.gpu:
            try:
                from cuml.manifold import UMAP
                print("Using GPU-accelerated UMAP")
            except ImportError:
                print("cuml not available. Falling back to CPU implementation.")
                print("For GPU acceleration, install cuml via conda:")
                print("conda install -c rapidsai -c conda-forge -c nvidia cuml cuda-version=11.8")
                UMAP = umap.UMAP
        else:
            UMAP = umap.UMAP
            print("Using CPU UMAP")
        
        fit = UMAP(n_components=self.n_components, **kwargs)
        return fit.fit_transform(self.maps)
    
    def _embed_phate(self, **kwargs):
        """Compute PHATE embedding."""
        phate = _check_phate_available()
        phate_op = phate.PHATE(n_components=self.n_components, **kwargs)
        return phate_op.fit_transform(self.maps)
    
    def _embed_tsne(self, perplexity=30, n_jobs=8, random_state=42, **kwargs):
        """Compute t-SNE embedding."""
        TSNE = _check_tsne_available()
        tsne = TSNE(
            perplexity=perplexity,
            metric="euclidean",
            n_jobs=n_jobs,
            random_state=random_state,
            **kwargs
        )
        return tsne.fit_transform(self.maps)
    
    def _embed_pca(self, **kwargs):
        """Compute PCA embedding."""
        _, _, PCA = _check_sklearn_available()
        return PCA(n_components=self.n_components, **kwargs).fit_transform(self.maps)
    
    def _embed_diffusion_maps(self, epsilon=None, batch_size=10000, dist_fname='distances.dat', **kwargs):
        """Compute Diffusion Maps embedding."""
        '''D = self._compute_distance_matrix(batch_size, dist_fname)
        import diffusion_maps
        vals, embedding = diffusion_maps.op(D, epsilon, **kwargs)
        return embedding[:, :self.n_components]'''
        raise NotImplementedError(
        "Diffusion Maps embedding is not yet implemented. "
        "This feature will be available in a future release."
    )
    
    
    def _compute_distance_matrix(self, batch_size=10000, dist_fname='distances.dat', return_flat=False):
        """Compute pairwise distance matrix with GPU acceleration if available.
        
        Args:
            batch_size: Batch size for GPU computation
            dist_fname: Temporary file for distance matrix memmap
            return_flat: If True, returns flattened distance matrix for hierarchical clustering
        
        Returns:
            numpy.ndarray: Distance matrix (flattened if return_flat=True)
        """
        t1 = time.time()
        
        if self.gpu:
            try:
                import tensorflow as tf
                D = self._compute_distances_gpu(batch_size, dist_fname)
                D = np.array(D).astype(np.float32)
                np.fill_diagonal(D, 0)
                D = (D + D.T) / 2  # matrix may be antisymmetric due to precision errors
                if return_flat:
                    D = squareform(D)
                os.remove(dist_fname)  # Clean up temporary file
            except ImportError:
                print("TensorFlow not available. Falling back to CPU implementation.")
                D_flat = self._compute_distances_cpu()
                D = squareform(D_flat) if not return_flat else D_flat
        else:
            D_flat = self._compute_distances_cpu()
            D = squareform(D_flat) if not return_flat else D_flat
                
        print(f'Distances time: {time.time() - t1:.2f}s')
        return D
    
    def cluster(self, embedding=None, method='kmeans', n_clusters=10, **kwargs):
        """Cluster the embedded data.
        
        Args:
            embedding: Optional pre-computed embedding. If None, uses stored embedding
            method: Clustering method ('kmeans', 'dbscan', or 'hierarchical')
            n_clusters: Number of clusters for kmeans
            **kwargs: Additional clustering parameters
                For DBSCAN:
                    eps: Maximum distance between samples (default: 0.01)
                    min_samples: Minimum samples per cluster (default: 10)
                For KMeans:
                    random_state: Random seed (default: 0)
                    n_init: Number of initializations (default: 10)
                For Hierarchical:
                    batch_size: Batch size for GPU computation (default: 10000)
                    link_method: Linkage method (default: 'ward')
                    dist_fname: Temporary file for distance matrix
                    store_distances: Whether to return distances (default: False)
        
        Returns:
            For kmeans/dbscan:
                numpy.ndarray: Cluster labels for each sample
            For hierarchical:
                scipy.cluster.hierarchy.linkage: Linkage matrix for hierarchical clustering
                (use get_cluster_labels() to obtain cluster assignments)
            If store_distances=True with hierarchical:
                tuple: (linkage_matrix, distance_matrix)
        """
        if method not in self.SUPPORTED_CLUSTERERS:
            raise ValueError(f"Method must be one of {self.SUPPORTED_CLUSTERERS}")
            
        # Special case for hierarchical clustering - works directly on maps
        if method == 'hierarchical':
            return self._cluster_hierarchical(**kwargs)
            
        # For other methods, need embedding
        if embedding is None:
            if self.embedding is None:
                raise ValueError("No embedding provided or computed. Run embed() first.")
            embedding = self.embedding

        if method in ['kmeans', 'dbscan']:
            KMeans, DBSCAN, _ = _check_sklearn_available()
            
            if method == 'kmeans':
                clusterer = KMeans(
                    n_clusters=n_clusters,
                    init='k-means++',
                    random_state=kwargs.get('random_state', 0),
                    n_init=kwargs.get('n_init', 10)
                )
            else:  # dbscan
                clusterer = DBSCAN(
                    eps=kwargs.get('eps', 0.01),
                    min_samples=kwargs.get('min_samples', 10)
                )
            
        self.cluster_labels = clusterer.fit_predict(embedding)
        return self.cluster_labels
    
    def _cluster_hierarchical(self, batch_size=10000, link_method='ward', 
                            dist_fname='distances.dat', store_distances=False):
        """Perform hierarchical clustering with optional GPU acceleration."""
        D_flat = self._compute_distance_matrix(batch_size, dist_fname, return_flat=True)
        
        print('Computing hierarchical clusters...')
        t1 = time.time()
        linkage = hierarchy.linkage(D_flat, method=link_method, metric='euclidean')
        print(f'Linkage time: {time.time() - t1:.2f}s')
        
        if store_distances:
            return linkage, D_flat
        return linkage

    def _compute_distances_gpu(self, batch_size, dist_fname):
        """Compute pairwise distances on GPU using TensorFlow with memory mapping."""
        import tensorflow as tf
        
        A = tf.cast(self.maps, tf.float32)
        B = tf.cast(self.maps, tf.float32)
        num_A = A.shape[0]
        num_B = B.shape[0]
        
        # Create memory-mapped file for distance matrix
        distance_matrix = np.memmap(dist_fname, dtype=np.float32, mode='w+', 
                                shape=(num_A, num_B))
        
        def pairwise_distance(A_batch, B_batch):
            v_A = tf.expand_dims(tf.reduce_sum(tf.square(A_batch), 1), 1)
            v_B = tf.expand_dims(tf.reduce_sum(tf.square(B_batch), 1), 1)
            p1 = tf.reshape(v_A, (-1, 1))
            p2 = tf.reshape(v_B, (1, -1))
            dist_squared = tf.add(p1, p2) - 2 * tf.matmul(A_batch, B_batch, transpose_b=True)
            dist_squared = tf.maximum(dist_squared, 0.0)  # ensure non-negative values
            return tf.sqrt(dist_squared)
        
        for i in tqdm(range(0, num_A, batch_size), desc='Distance batch'):
            A_batch = A[i:i + batch_size]
            for j in range(0, num_B, batch_size):
                B_batch = B[j:j + batch_size]
                distances = pairwise_distance(A_batch, B_batch).numpy()
                distance_matrix[i:i + batch_size, j:j + batch_size] = distances
        
        distance_matrix.flush()  # Ensure all data is written to disk
        return distance_matrix

    def _compute_distances_cpu(self):
        """Compute pairwise distances on CPU."""
        nS = self.maps.shape[0]
        D_upper = np.zeros(shape=(nS, nS))
        
        for i in tqdm(range(nS), desc='Computing distances'):
            for j in range(i + 1, nS):
                D_upper[i,j] = np.linalg.norm(self.maps[i,:] - self.maps[j,:])
        
        D = D_upper + D_upper.T - np.diag(np.diag(D_upper))  # Match original code
        return squareform(D)
    
    def normalize(self, embedding, to_sum=False, copy=True):
        """Normalize embedding vectors to [0,1] range.
        
        Args:
            embedding: Array of shape (n_samples, n_dimensions)
            to_sum: If True, normalize to sum=1. If False, normalize to range [0,1]
            copy: If True, operate on a copy of the data
        
        Returns:
            numpy.ndarray: Normalized embedding
        """
        d = embedding if not copy else np.copy(embedding)
        d -= np.min(d, axis=0)
        d /= (np.sum(d, axis=0) if to_sum else np.ptp(d, axis=0))  # normalize to [0,1]
        return d
    
    def plot_embedding(self, embedding, labels=None, dims=[0,1], 
                    normalize=False, cmap='jet', s=2.5, alpha=1.0, 
                    linewidth=0.1, colorbar_label=None, sort_order=None, 
                    ref_index=None, legend_loc='upper left', save_path=None, dpi=200):
        """Plot embedding and optionally color by labels/values.
        
        Args:
            embedding: Array of shape (n_samples, n_dimensions)
            labels: Values for coloring points. Can be:
                - numpy array of shape (N,) or (N,1)
                - pandas Series/DataFrame column (e.g., mave['DNN'])
                - None (points will be single color)
            dims: Which dimensions to plot [dim1, dim2]
            normalize: Whether to normalize embedding to [0,1] range
            cmap: Colormap for points (e.g., 'viridis', 'jet', 'tab10')
                - Use 'viridis'/'jet' for continuous values
                - Use 'tab10'/'Set3' for discrete clusters
            s: Point size (default: 2.5)
            alpha: Point transparency (default: 1.0)
            linewidth: Width of point edges (default: 0.1)
            colorbar_label: Label for colorbar (if None, no colorbar shown)
            sort_order: Order to plot points ('ascending', 'descending', or None)
                - Useful for ensuring important points are plotted on top
                - Points are sorted based on their label values
                - Works with both numpy arrays and pandas Series/DataFrames
            ref_index: Index of reference/wild-type sequence to highlight (default: None)
                - Will be shown as a black star on the plot
            legend_loc: Location of legend for reference sequence ('best', 'top left', 'upper right', etc.)
            save_path: Path to save figure (if None, displays plot)
            dpi: DPI for saved figure (default: 200)
        
        Example usage:
            # Basic plot with reference sequence
            clusterer.plot_embedding(embedding, ref_index=0)
            
            # Color by DNN predictions with colorbar and reference
            clusterer.plot_embedding(
                embedding,
                labels=mave['DNN'],  # or y_mut numpy array
                colorbar_label='DNN prediction',
                sort_order='descending',  # high predictions on top
                ref_index=ref_idx
            )
        """
        try:
            plt.close()
        except:
            pass

        # Normalize embedding if requested  
        if normalize:
            embedding = self.normalize(embedding)
        
        # Ensure embedding has enough dimensions
        if embedding.shape[1] <= max(dims):
            raise ValueError(f"Embedding has {embedding.shape[1]} dimensions, but dims={dims} was requested")
        
        # Convert labels to numpy array if needed and flatten
        if labels is not None:
            if hasattr(labels, 'values'):  # pandas Series/DataFrame
                labels = labels.values
            labels = labels.flatten()  # Handle both (N,1) and (N,) arrays
        
        # Sort points if requested
        if labels is not None and sort_order is not None:
            if sort_order not in ['ascending', 'descending']:
                raise ValueError("sort_order must be 'ascending', 'descending', or None")
            sort_idx = np.argsort(labels)
            if sort_order == 'descending':
                sort_idx = sort_idx[::-1]
        else:
            sort_idx = slice(None)  # equivalent to ':'
        
        scatter = plt.scatter(
            embedding[:, dims[0]][sort_idx], 
            embedding[:, dims[1]][sort_idx], 
            c=None if labels is None else labels[sort_idx],
            cmap=cmap,
            s=s,
            alpha=alpha,
            linewidth=linewidth,
            edgecolors='k',
            zorder=0
        )
        
        # Plot reference sequence if specified
        if ref_index is not None:
            plt.scatter(
                embedding[ref_index, dims[0]], 
                embedding[ref_index, dims[1]], 
                marker='*', 
                c='k', 
                s=200,  # make star larger than regular points
                label='Reference',
                zorder=100
            )
            plt.legend(markerscale=1, fontsize='small', frameon=False, loc=legend_loc)

        # Create colorbar using the main scatter object
        if colorbar_label is not None and labels is not None:
            plt.colorbar(scatter, label=colorbar_label)
        
        plt.gca().spines['top'].set_visible(False)
        plt.gca().spines['right'].set_visible(False)
        plt.xlabel(f'Ψ{dims[0]+1}')
        plt.ylabel(f'Ψ{dims[1]+1}')
        plt.gca().set_aspect('equal')
        plt.tight_layout()
        
        if save_path:
            plt.savefig(save_path, facecolor='w', dpi=dpi, bbox_inches='tight')
            plt.close()
        else:
            plt.show()

    def plot_histogram(self, embedding, dims=[0,1], bins=101, cmap='viridis', 
                    colorbar_label='Count', save_path=None, dpi=200):
        """Plot 2D histogram of embedding points.
        
        Args:
            embedding: Array of shape (n_samples, n_dimensions)
            dims: Which dimensions to plot [dim1, dim2]
            bins: Number of bins for histogram (default: 101)
            cmap: Colormap for histogram (default: 'viridis')
            colorbar_label: Label for colorbar (if None, shows 'Count')
            save_path: Path to save figure (if None, displays plot)
            dpi: DPI for saved figure (default: 200)
        """
        try:
            plt.close()
        except:
            pass
        
        # Ensure embedding has enough dimensions
        if embedding.shape[1] <= max(dims):
            raise ValueError(f"Embedding has {embedding.shape[1]} dimensions, but dims={dims} was requested")
        
        # Always normalize the embedding
        embedding = self.normalize(embedding)
        
        # Create 2D histogram
        H, edges = np.histogramdd(
            np.vstack((embedding[:, dims[0]], embedding[:, dims[1]])).T,
            bins=bins,
            range=((0,1), (0,1)),
            density=False,
            weights=None
        )
        
        # Plot histogram
        mesh = plt.pcolormesh(H.T, cmap=cmap)
        
        if colorbar_label is not None:
            plt.colorbar(mesh, label=colorbar_label)
        else:
            plt.colorbar(mesh, label='Count')
        
        plt.gca().spines['top'].set_visible(False)
        plt.gca().spines['right'].set_visible(False)
        plt.xlabel(f'Ψ{dims[0]+1}')
        plt.ylabel(f'Ψ{dims[1]+1}')
        plt.gca().set_aspect('equal')
        plt.tight_layout()
        
        if save_path:
            plt.savefig(save_path, facecolor='w', dpi=dpi, bbox_inches='tight')
            plt.close()
        else:
            plt.show()

    def plot_dendrogram(self, linkage, figsize=(15, 10), leaf_rotation=90, 
                        leaf_font_size=8, cut_level=None, save_path=None, dpi=200):
        """Plot dendrogram from hierarchical clustering linkage matrix.
        
        Args:
            linkage: Hierarchical clustering linkage matrix
            figsize: Figure size (width, height)
            leaf_rotation: Rotation of leaf labels
            leaf_font_size: Font size for leaf labels
            cut_level: Optional height at which to draw horizontal cut line
            save_path: Path to save figure (if None, displays plot)
            dpi: DPI for saved figure (default: 200)
        """
        sys.setrecursionlimit(100000)  # Fix for large dendrograms
        
        plt.figure(figsize=figsize)
        plt.title('Hierarchical Clustering Dendrogram')
        plt.xlabel('Sample index')
        plt.ylabel('Distance')
        
        with plt.rc_context({'lines.linewidth': 2}):
            hierarchy.dendrogram(
                linkage,
                leaf_rotation=leaf_rotation,
                leaf_font_size=leaf_font_size,
            )

        if cut_level is not None:
            plt.axhline(y=cut_level, color='r', linestyle='--', label=f'Cut level: {cut_level:.3f}')
            plt.legend(frameon=False, loc='best')
        
        plt.xticks([])
        plt.gca().spines['top'].set_visible(False)
        plt.gca().spines['right'].set_visible(False)
        
        if save_path:
            plt.savefig(save_path, facecolor='w', dpi=dpi, bbox_inches='tight')
            plt.close()
        else:
            plt.show()

    def get_cluster_labels(self, linkage, criterion='maxclust', max_distance=10, n_clusters=200):
        """Get cluster labels from hierarchical clustering linkage matrix.
        
        Args:
            linkage: Hierarchical clustering linkage matrix
            criterion: How to form flat clusters ('distance' or 'maxclust')
                'distance': Cut tree at specified height 
                'maxclust': Produce specified number of clusters
            max_distance: Maximum cophenetic distance within clusters 
                        (only used if criterion='distance')
            n_clusters: Desired number of clusters to produce
                    (only used if criterion='maxclust')
        
        Returns:
            numpy.ndarray: Cluster labels (zero-indexed)
            float: Cut level (max_distance if criterion='distance', or computed level if criterion='maxclust')
        """
        if criterion == 'distance':
            clusters = hierarchy.fcluster(linkage, max_distance, criterion='distance')
            clusters = clusters - 1  # Zero-index clusters
            self.cluster_labels = clusters  # Store labels in the instance
            return clusters, max_distance
            
        elif criterion == 'maxclust':
            clusters = hierarchy.fcluster(linkage, n_clusters, criterion='maxclust')
            clusters = clusters - 1  # Zero-index clusters
            
            # Find the cut level that gives the desired number of clusters
            sorted_heights = np.sort(linkage[:, 2])[::-1]  # Heights from highest to lowest
            for height in sorted_heights:
                temp_clusters = hierarchy.fcluster(linkage, height, criterion='distance')
                if len(np.unique(temp_clusters)) == n_clusters:
                    max_d = height
                    break
                    
            print(f"Cut level for {n_clusters} clusters: {max_d:.3f}")
            self.cluster_labels = clusters  # Store labels in the instance
            return clusters, max_d
            
        else:
            raise ValueError("criterion must be either 'distance' or 'maxclust'")