from .prefix_computers import CPUPrefixComputer
class PromptOptimizer:
    """Optimizes query order within groups to maximize prefix sharing."""
    
    
    @staticmethod
    def optimize_group_order(prompts, group_indices):
        """Optimize the order within a group to maximize prefix sharing."""
        if len(group_indices) <= 1:
            return group_indices
        
        group_prompts = [prompts[i] for i in group_indices]
        n = len(group_prompts)
        
        # Find the query with longest common prefix with others as starting point
        best_start = 0
        best_total_prefix = 0
        
        for i in range(n):
            total_prefix = sum(CPUPrefixComputer._longest_common_prefix_length(group_prompts[i], group_prompts[j])
                              for j in range(n) if i != j)
            if total_prefix > best_total_prefix:
                best_total_prefix = total_prefix
                best_start = i
        
        # Start with best query and greedily add others
        ordered_indices = [best_start]
        remaining = set(range(n)) - {best_start}
        
        while remaining:
            last_added = ordered_indices[-1]
            best_next = max(remaining,
                           key=lambda x: CPUPrefixComputer._longest_common_prefix_length(
                               group_prompts[last_added], group_prompts[x]))
            ordered_indices.append(best_next)
            remaining.remove(best_next)
        
        return [group_indices[i] for i in ordered_indices]