from .prefix_computers import PrefixComputer


class PromptGrouper:
    """Handles grouping of prompts based on prefix similarities."""
    
    def __init__(self, prefix_computer: PrefixComputer):
        self.prefix_computer = prefix_computer
    
    def group_by_prefix_similarity(self, prompts):
        """Group prompts using prefix similarity matrix."""
        n = len(prompts)
        if n <= 1:
            return [list(range(n))]
        
        prefix_matrix = self.prefix_computer.compute_prefix_matrix(prompts)
        return self._greedy_clustering(n, prefix_matrix)
    
    def _greedy_clustering(self, n, prefix_matrix):
        """Perform greedy clustering based on prefix matrix."""
        used = [False] * n
        groups = []
        
        for i in range(n):
            if used[i]:
                continue
            
            current_group = [i]
            used[i] = True
            
            while True:
                best_candidate = -1
                best_score = 0
                
                for j in range(n):
                    if used[j]:
                        continue
                    
                    min_prefix_with_group = min(prefix_matrix[j, member] for member in current_group)
                    if min_prefix_with_group > best_score:
                        best_score = min_prefix_with_group
                        best_candidate = j
                
                if best_candidate != -1 and best_score > 0:
                    current_group.append(best_candidate)
                    used[best_candidate] = True
                else:
                    break
            
            groups.append(current_group)
        
        return groups