import numpy as np
from scipy.cluster.hierarchy import linkage
from .distance_matrix import DistanceCalculator
from .tree_nodes import NodeManager
from .docs_reordering import DocReorderer

class HierarchicalClusterer:
    """Main class for hierarchical clustering of prompts."""
    
    def __init__(self, similarity_method="sharp", linkage_method="average", use_triton=True):
        self.similarity_method = similarity_method
        self.linkage_method = linkage_method
        self.distance_calculator = DistanceCalculator(use_triton=use_triton)
        self.node_manager = NodeManager()
        self.prompt_reorderer = DocReorderer()
    
    def fit_transform(self, prompts):
        """Perform clustering and return results."""
        n = len(prompts)
        
        if n < 2:
            print("Need at least two prompts for clustering.")
            return self._handle_single_prompt(prompts)
        
        # Compute distance matrix
        print("Computing distance matrix...")
        distance_matrix = self.distance_calculator.compute_distance_matrix(
            prompts, self.similarity_method
        )
        
        # Perform hierarchical clustering
        print("Performing hierarchical clustering...")
        condensed_distances = self.distance_calculator.get_condensed_matrix(distance_matrix)
        linkage_matrix = linkage(condensed_distances, method=self.linkage_method)
        
        # Build tree structure
        print("Building tree structure...")
        self._build_tree(prompts, linkage_matrix)
        
        # Clean up tree
        self.node_manager.cleanup_empty_nodes()
        
        # Generate reordered prompts
        print("Generating reordered prompts...")
        reordered_prompts = self.prompt_reorderer.reorder_prompts(
            prompts, self.node_manager.unique_nodes
        )
        
        # Print statistics
        stats = self.node_manager.get_node_stats()
        print(f"Clustering completed: {stats['total_nodes']} unique nodes, "
              f"{stats['leaf_nodes']} leaf nodes")
        
        return ClusteringResult(
            linkage_matrix=linkage_matrix,
            cluster_nodes=self.node_manager.cluster_nodes,
            unique_nodes=self.node_manager.unique_nodes,
            reordered_prompts=reordered_prompts,
            original_prompts=prompts,
            stats=stats
        )
    
    def _handle_single_prompt(self, prompts):
        """Handle case with less than 2 prompts."""
        for i, prompt in enumerate(prompts):
            self.node_manager.create_leaf_node(i, prompt)
        
        return ClusteringResult(
            linkage_matrix=np.empty((0, 4)),
            cluster_nodes=self.node_manager.cluster_nodes,
            unique_nodes=self.node_manager.unique_nodes,
            reordered_prompts=prompts.copy() if hasattr(prompts, 'copy') else list(prompts),
            original_prompts=prompts,
            stats=self.node_manager.get_node_stats()
        )
    
    def _build_tree(self, prompts, linkage_matrix):
        """Build the clustering tree from linkage matrix."""
        n = len(prompts)
        
        # Create leaf nodes
        for i, prompt in enumerate(prompts):
            self.node_manager.create_leaf_node(i, prompt)
        
        # Process internal nodes
        for i, (idx1, idx2, distance, _) in enumerate(linkage_matrix):
            new_node_id = n + i
            self.node_manager.create_internal_node(
                new_node_id, int(idx1), int(idx2), distance
            )

class ClusteringResult:
    """Container for clustering results."""
    
    def __init__(self, linkage_matrix, cluster_nodes, unique_nodes, 
                 reordered_prompts, original_prompts, stats):
        self.linkage_matrix = linkage_matrix
        self.cluster_nodes = cluster_nodes
        self.unique_nodes = unique_nodes
        self.reordered_prompts = reordered_prompts
        self.original_prompts = original_prompts
        self.stats = stats
    
    def print_tree(self):
        """Print the unique cluster tree."""
        print("\n--- Unique Cluster Tree Nodes ---")
        for node_id in sorted(self.unique_nodes.keys()):
            node = self.unique_nodes[node_id]
            print(node)
            print(f"  Content: {node.doc_ids}")
            print(f"  Original indices: {sorted(node.original_indices)}")
            if not node.is_leaf:
                print(f"  Children: {node.children}")
                print(f"  Merge distance: {node.merge_distance:.4f}")
            print("-" * 40)