from .cuda_kernels import TritonOperations
import numpy as np


class DistanceCalculator:
    """Handles distance matrix computation for clustering."""
    
    def __init__(self, use_triton=True, use_cuda=False):
        self.use_triton = use_triton
        self.use_cuda = use_cuda
        
        if use_triton:
            try:
                import triton
                self.triton_ops = TritonOperations()
                print("Using Triton kernels for distance computation")
            except Exception as e:
                print(f"Triton not available ({e}), falling back to CUDA")
                self.use_triton = False
        
        assert not use_cuda, "CUDA support is not implemented yet."
        
        # if not self.use_triton and use_cuda:
        #     try:
        #         from .cuda_kernels import CudaOperations
        #         self.cuda_ops = CudaOperations()
        #         print("Using CUDA kernels for distance computation")
        #     except Exception as e:
        #         print(f"CUDA not available ({e}), falling back to CPU")
        #         self.use_cuda = False
    
    def compute_distance_matrix(self, queries, method="sharp"):
        """Compute distance matrix between queries."""
        if self.use_triton:
            return self._compute_triton_distance_matrix(queries, method)
        # elif self.use_cuda:
        #     return self._compute_cuda_distance_matrix(queries, method)
        else:
            return self._compute_cpu_distance_matrix(queries, method)
    
    def _compute_triton_distance_matrix(self, queries, method):
        """Compute distance matrix using Triton."""
        try:
            return self.triton_ops.compute_distance_matrix(queries, method)
        except Exception as e:
            print(f"Triton computation failed ({e}), falling back to CPU")
            return self._compute_cpu_distance_matrix(queries, method)
    
    # def _compute_cuda_distance_matrix(self, queries, method):
    #     """Compute distance matrix using CUDA."""
    #     return self.cuda_ops.compute_distance_matrix(queries, method)
    
    def _compute_cpu_distance_matrix(self, queries, method):
        """Compute distance matrix using CPU (fallback)."""
        n = len(queries)
        if n < 2:
            return np.array([])
        
        # Convert to sets for efficient intersection
        prompt_sets = [set(q) for q in queries]
        distance_matrix = np.zeros((n, n), dtype=np.float32)
        
        for i in range(n):
            for j in range(i, n):
                similarity = self._compute_similarity(prompt_sets[i], prompt_sets[j], method)
                distance = 1.0 - similarity
                distance_matrix[i, j] = distance
                distance_matrix[j, i] = distance
        
        return distance_matrix
    
    def _compute_similarity(self, set1, set2, method):
        """Compute similarity between two sets."""
        if len(set1) == 0 or len(set2) == 0:
            return 0.0
        
        intersection_size = len(set1.intersection(set2))
        max_size = max(len(set1), len(set2))
        base_similarity = intersection_size / max_size
        
        if method == "direct":
            return base_similarity
        elif method == "square":
            return base_similarity ** 2
        elif method == "sharp":
            a = 2.0
            exp_a = np.exp(a)
            similarity = (np.exp(a * base_similarity * 2) - 1) / (exp_a - 1)
            return np.clip(similarity, 0.0, 1.0)
        else:
            return base_similarity
    
    def get_condensed_matrix(self, distance_matrix):
        """Convert full distance matrix to condensed format for scipy."""
        n = distance_matrix.shape[0]
        return distance_matrix[np.triu_indices(n, k=1)]