"""
Utility module to find and group duplicate structures based on their hashed representations.
Designed to be simple and independent of the main code.
"""
from typing import Callable, Iterable, Dict, List, Any, Union
from collections import defaultdict

def get_duplicate_indices(
    partition: Iterable[Any], 
    hash_funcs: Union[Callable[[Any], Any], Iterable[Callable[[Any], Any]]],
    only_duplicates: bool = True
) -> Dict[str, List[int]]:
    """
    Identifies and groups indices from a partition that represent the same structure,
    based on one or more hash functions.

    Args:
        partition: An iterable of items (e.g., list of Atoms objects or dictionaries).
        hash_funcs: A single hash callable, or an iterable of hash callables. 
                    Each callable should take an item from the partition and return a hashable object (e.g. string).
        only_duplicates: If True, the resulting dictionary will only contain groups 
                         that have more than 1 index (i.e., actual duplicates). If False,
                         it returns all groups including unique items.

    Returns:
        A dictionary mapping a combined hash string to a list of original indices 
        from the partition representing the same structure.
    """
    if callable(hash_funcs):
        hash_funcs = [hash_funcs]
        
    hash_groups = defaultdict(list)
    
    for idx, item in enumerate(partition):
        print(idx, item.atoms.x.shape, item.atoms.N, item.AtomPositionManager.kdtree.data.shape, item.AtomPositionManager.atomLabelsList.shape)
        # Generate the combined hash by applying all given hash functions
        combined_hash = "_".join(str(f(item)) for f in hash_funcs)
        if combined_hash in hash_groups:
            print("Duplicate found!")
            print(hash_groups[combined_hash])
            print(idx)
        hash_groups[combined_hash].append(idx)
        
    if only_duplicates:
        # Filter groups to only include those with more than one item
        return {h: indices for h, indices in hash_groups.items() if len(indices) > 1}
    else:
        # Return all index groups
        return dict(hash_groups)
