from __future__ import annotations
from typing import Any, Dict, Optional, List, Union
from pathlib import Path
import numpy as np

from .factory import load_config, build_default_engine
from .core.engine import GeneticAlgorithm

class QuickGA:
    """
    A high-level facade for EZGA designed for simplicity and ease of use.
    
    QuickGA allows defining and running a Genetic Algorithm experiment with
    minimal boilerplate, using sensible defaults and shorthand string aliases
    for common components.
    """
    
    def __init__(
        self,
        composition: List[str],
        initial_structure: Optional[Union[str, Path]] = None,
        calculator: Union[str, Any] = "lj",
        mace_model_path: Optional[str] = None,
        max_generations: int = 50,
        output_path: str = "runs/quick_ga",
        seed: Optional[int] = None,
        **kwargs
    ):
        """
        Initialize a QuickGA experiment.
        
        Args:
            composition: List of atomic symbols (e.g., ["Ni", "O"]).
            initial_structure: Path to an initial XYZ or POSCAR file to seed the population.
            calculator: Shorthand string (e.g., "lj", "emt", "mace") or an ASE calculator instance.
            mace_model_path: Path to a MACE model file (required if calculator="mace").
            max_generations: Maximum number of generations to run.
            output_path: Directory to save results.
            seed: Random seed for reproducibility.
            **kwargs: Additional configuration parameters to override defaults.
        """
        self.composition = composition
        self.output_path = output_path
        
        # Determine unique species in composition for features and constraints
        unique_species = sorted(list(set(composition)))
        
        # Build the base config dictionary
        # Simple shorthand for MACE if a path is provided
        if calculator == "mace" and mace_model_path:
            calculator = f"mace:{mace_model_path}"

        self.cfg_dict = {
            "max_generations": max_generations,
            "output_path": output_path,
            "rng": seed,
            "foreigners": kwargs.pop("foreigners", 10), # Default to 10 random structures per gen if empty
            "simulator": {
                "calculator": calculator
            },
            "population": {
                "filter_duplicates": True,
                "dataset_path": initial_structure,
            },
            "generative": {
                "size": kwargs.pop("foreigners", 10), # Default to 10 random structures per gen if empty
                "start_gen": 0,
                "bo_kwargs": {
                    "bounds": np.array([[0.0, float(len(composition))] for _ in unique_species]),
                }
            },
            "mutation_funcs": kwargs.pop("mutation_funcs", ["rattle", "swap", "add", "remove"]),
            "evaluator": {
                "features_funcs": [{
                    "factory": "ezga.evaluator.features:feature_composition_vector",
                    "kwargs": {"IDs": unique_species}
                }]
            },
            **kwargs
        }
        
        # Ensure composition is mapped for constraints if using Grand Canonical
        # (This is a common pain point for users)
        from .DoE.constraint import ConstraintFactory
        ConstraintFactory.set_name_mapping({s: i for i, s in enumerate(unique_species)})
        
        # Internal state
        self._config = None
        self._engine = None

    @property
    def config(self):
        """Lazy-loaded GAConfig."""
        if self._config is None:
            self._config = load_config(self.cfg_dict)
        return self._config

    @property
    def engine(self) -> GeneticAlgorithm:
        """Lazy-loaded GeneticAlgorithm engine."""
        if self._engine is None:
            self._engine = build_default_engine(self.config)
        return self._engine

    def run(self) -> Any:
        """
        Run the evolution.
        
        Returns:
            The final population or result of the evolution.
        """
        return self.engine.run()

    def __repr__(self) -> str:
        return f"<QuickGA system={self.composition} out={self.output_path}>"
