"""
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
import numpy as np
import faiss
from tqdm import tqdm

def get_duplicate_indices_hash(
    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 returning hashable objects.

    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):
        # Generate the combined hash by applying all given hash functions
        combined_hash = "_".join(str(f(item)) for f in hash_funcs)
        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)

def get_duplicate_indices_faiss(
    partition: Iterable[Any], 
    global_descriptor: 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 returning vectors (e.g. numeric arrays).
    
    Uses faiss for fast L2 distance similarity matching.

    Args:
        partition: An iterable of items (e.g., list of Atoms objects or dictionaries).
        global_descriptor: A single hash callable, or an iterable of hash callables. 
                    Each callable should take an item from the partition and return a numeric array or vector.
        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 group identifier string to a list of original indices 
        from the partition representing the same structure.
    """
    if callable(global_descriptor):
        global_descriptor = [global_descriptor]
        
    partition_list = list(partition)
    if not partition_list:
        return {}
        
    vectors = []
    
    for idx, item in enumerate(partition_list):
        # Generate the combined vector by applying all given hash functions
        vec_parts = [np.atleast_1d(f(item)).flatten() for f in global_descriptor]
        vectors.append(np.concatenate(vec_parts))
        
    # Convert to float32 contiguous array as required by faiss
    vectors_np = np.array(vectors, dtype=np.float32)
    
    if len(vectors_np) == 0 or vectors_np.size == 0:
        return {}

    d = vectors_np.shape[1]
    index = faiss.IndexFlatL2(d)
    index.add(vectors_np)
    
    # Range search for duplicates with a very small squared L2 distance threshold
    radius = 1e-5
    lims, D, I = index.range_search(vectors_np, radius)
    
    hash_groups = {}
    visited = set()
    
    for i in range(len(partition_list)):
        if i in visited:
            continue
            
        start = lims[i]
        end = lims[i+1]
        neighbors = I[start:end]
        
        group = []
        for n in neighbors:
            if n not in visited:
                group.append(int(n))
                visited.add(int(n))
                
        if len(group) > 0:
            group.sort()
            # Use the first element's index to create a group key string
            key = f"group_{group[0]}"
            hash_groups[key] = group
            
    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)

def estimate_faiss_threshold(
    partition: Iterable[Any], 
    global_descriptor: Union[Callable[[Any], Any], Iterable[Callable[[Any], Any]]],
    max_displacement: float = 0.1,
    random_seed: int = 42,
    percentile: float = 95.0,
    use_mass_factor: bool = False,
    n_repetitions: int = 5
) -> float:
    """
    Estimates a safe FAISS L2 squared distance threshold for duplicate detection 
    by applying uniform noise to the atomic positions of a batch of structures.
    The noise amplitude is scaled inverse-proportionally to the square root of 
    the atomic masses to simulate realistic vibrational displacements.

    Args:
        partition: An iterable of items (e.g., list of Atoms objects or dictionaries).
        global_descriptor: A single hash callable, or an iterable of hash callables returning vectors.
        max_displacement: Maximum displacement in Angstroms for the lightest atom.
        random_seed: Random seed for reproducible noise generation.
        percentile: Percentile of the observed squared distances to use (default 95.0) 
                    to avoid outliers causing huge threshold variations.
        use_mass_factor: If True, scales the noise amplitude inverse-proportionally 
                         to the square root of the atomic masses. Default is False.
        n_repetitions: Number of noise samples to apply to each structure to build a 
                       more robust statistical distribution of displacements (default 5).

    Returns:
        float: The recommended squared L2 distance threshold for faiss.RangeSearch, 
               with a 10% safety margin added to the determined percentile distance.
    """
    if callable(global_descriptor):
        global_descriptor = [global_descriptor]
        
    partition_list = list(partition)
    if not partition_list:
        return 1e-5
        
    distances_sq = []
    
    # Initialize a local random state for reproducible noise
    rng = np.random.RandomState(random_seed)
    
    # Try to import periodic table for masses, fallback to Z*2 approximation
    try:
        from ase.data import atomic_masses
        has_masses = True
    except ImportError:
        has_masses = False
        
    for item in tqdm(partition_list, desc="Estimating FAISS threshold"):
        apm = getattr(item, 'AtomPositionManager', None)
        if apm is None:
            continue
            
        # Get clean vectors
        vec_parts_clean = [np.atleast_1d(f(item)).flatten() for f in global_descriptor]
        vec_clean = np.concatenate(vec_parts_clean).astype(np.float32)
        
        # Save original positions
        orig_positions = np.copy(apm.atomPositions)
        try:
            Z = np.asarray(apm.get_atomic_numbers(), dtype=np.int32)
            if has_masses:
                masses = np.array([atomic_masses[z] if z < len(atomic_masses) else z*2.0 for z in Z])
            else:
                masses = np.where(Z == 1, 1.0, Z * 2.0)
                
            for _ in range(n_repetitions):
                # Generate uniform noise in [-1, 1]
                noise = rng.uniform(-1.0, 1.0, size=orig_positions.shape)
                    
                if use_mass_factor:
                    # Scale noise inversely proportional to sqrt(mass).
                    # Lighter atoms vibrate more: Amplitude ~ 1/sqrt(M)
                    mass_factor = 1.0 / np.sqrt(np.maximum(masses, 1.0))
                    # Normalize mass factor so that max displacement is applied to the lightest atom
                    mass_factor = mass_factor / np.max(mass_factor)
                    displacement = noise * max_displacement * mass_factor[:, np.newaxis]
                else:
                    displacement = noise * max_displacement
                    
                # Apply displacement
                apm.atomPositions = orig_positions + displacement
                
                # Recompute descriptor
                vec_parts_noisy = [np.atleast_1d(f(item)).flatten() for f in global_descriptor]
                vec_noisy = np.concatenate(vec_parts_noisy).astype(np.float32)
                
                # Squared L2 distance
                dist_sq = np.sum((vec_clean - vec_noisy) ** 2)
                distances_sq.append(dist_sq)
            
        finally:
            # Restore positions
            apm.atomPositions = orig_positions
            
    if not distances_sq:
        return 1e-5
        
    # Use a high percentile instead of the absolute max to avoid extreme outliers
    base_threshold = np.percentile(distances_sq, percentile)
    
    # Add a safety margin (e.g., 10%) to the robust distance 
    recommended_threshold = float(base_threshold * 1.1)
    
    # If the structures didn't change (e.g., descriptors are invariant or displacement was 0)
    if recommended_threshold < 1e-6:
        recommended_threshold = 1e-5
        
    return recommended_threshold
