# SPDX-License-Identifier: MIT
"""
Mutation Presets
================

This module provides high-level bundles of mutation operators to simplify
configuration. Instead of listing every operator individually, users can use
these presets to generate a standard suite of mutations for common scenarios.

Example:
    >>> from ezga.variation.presets import grand_canonical
    >>> mutations = grand_canonical(species=["Cu", "O"], constraints=...)

"""

from typing import List, Optional, Any, Dict, Union
import itertools

# Import primitive mutations
from ezga.variation.mutation import (
    mutation_rattle,
    mutation_swap,
    mutation_change_ID,
    mutation_add,
    mutation_remove
)

def grand_canonical(
    species: List[str],
    slab: bool = False,
    collision_tolerance: float = 2.0,
    rattle_std: float = 0.05,
    seed: int = 42,
    constraints: Optional[List[Dict[str, Any]]] = None
) -> List[Any]:
    """
    Returns a comprehensive list of operators for Grand Canonical exploration.
    
    Includes:
      - Rattle (conservative local optimization) for each species.
      - Swap (exchange identity) for all pairs.
      - Change ID (transmutation) for all pairs.
      - Add atom (Grand Canonical) for each species.
      - Remove atom (Grand Canonical) for each species.

    Args:
        species (List[str]): List of chemical symbols (e.g. ["Cu", "O"]).
        slab (bool): If True, 'mutation_add' respects slab geometry.
        collision_tolerance (float): Min distance for adding new atoms.
        rattle_std (float): Standard deviation for rattle operations.
        seed (int): RNG seed for rattle.
        constraints (Optional[List[Dict]]): Constraints to apply to ALL operators.

    Returns:
        List[Any]: A flattened list of configured mutation operators.
    """
    
    ops = []

    # 1. Rattle (Local Optimization)
    # Applied per species
    for s in species:
        ops.append(
            mutation_rattle(
                std=rattle_std,
                species=[s],
                seed=seed,
                constraints=constraints
            )
        )

    # 2. Pairs (Swap & Change ID)
    # Generate all ordered pairs (A, B) and (B, A)
    if len(species) > 1:
        pairs = list(itertools.permutations(species, 2))
        
        for (a, b) in pairs:
            # Swap
            ops.append(
                mutation_swap(
                    ID1=a,
                    ID2=b,
                    constraints=constraints
                )
            )
            # Transmute (Change ID)
            ops.append(
                mutation_change_ID(
                    ID_initial=a,
                    ID_final=b,
                    constraints=constraints
                )
            )

    # 3. Grand Canonical (Add & Remove)
    for s in species:
        # Add Atom
        # Bound is "all other species" to avoid adding into vacuum far from cluster
        # If single species, bound is itself
        others = [x for x in species]
        
        ops.append(
            mutation_add(
                species=[s],
                bound=others,
                collision_tolerance=collision_tolerance,
                slab=slab,
                constraints=constraints
            )
        )
        
        # Remove Atom
        ops.append(
            mutation_remove(
                species=s,
                constraints=constraints
            )
        )

    return ops
