# SPDX-License-Identifier: MIT
"""
Generate a functional boilerplate EZGA Python script for quick starting.
"""
from pathlib import Path
import typer
from rich.console import Console
from rich.panel import Panel

console = Console()

BOILERPLATE_CODE = '''# my_experiment.py
"""
Auto-generated EZGA Example Script.
Goal: Run a Gran-Canonical Evolution for a Cluster System using a Fast Calculator.
"""
# 0. EZGA Imports
from ezga.utils.generate_initial_population import generate_initial_population
from ezga.DoE.constraint import ConstraintFactory
from ezga.variation.mutation import * 
from ezga.variation.crossover import *
from ezga.evaluator.features import feature_composition_vector
from ezga.evaluator.objective import objective_formation_energy
from ezga.simulator.mace_calculator import ase_calculator
from ezga.factory import load_config
from ezga.hise.manager import run_hise

# Calculator import
from ase.calculators.emt import EMT

# 1. Elements
composition = ["Ni", "O"] 

# 2. System (Cluster)
pbc = [False, False, False]
hise_supercells = [(1,1,1)]
use_strain = False

# 2b. Composition (Variable Stoichiometry / Grand Canonical)
ConstraintFactory.set_name_mapping({s: i for i, s in enumerate(composition)})
constraints = [
    ConstraintFactory.within_range("Ni", 1, 20),
    ConstraintFactory.within_range("O",  1, 20),
]
use_grand_canonical = True

# 3. Features & Objectives
features = [feature_composition_vector(composition)]
objectives = [objective_formation_energy()]

# 4. Operators
mutations = [mutation_rattle(std=0.1, species=[s], seed=42, constraints=[]) for s in composition]
if use_strain: mutations += [mutation_random_strain(max_strain=0.1)]
mutations += [mutation_swap(composition[0], composition[1])]
if use_grand_canonical:
    mutations += [mutation_add(s) for s in composition] + [mutation_remove(s) for s in composition]
    
crossover = [crossover_voronoi_bisector_surface()]

# 5. Calculator (Fast Scouting with EMT)
# Switch to mace_mp() for production!
calculator = ase_calculator(calculator=EMT(), steps_max=40, fmax=0.05)

# 6. Global Config
cfg_dict = {
    "output_path": "runs/cluster_search",
    "max_generations": 50,
    "population": {
        "filter_duplicates": True,
        "collision_factor": 0.60,
        "constraints": constraints,
    },
    "generative": {'size': 0}, # Set to > 0 to enable Bayesian Optimization
    "variation": {"initial_mutation_rate": 1.0},
    "mutation_funcs": mutations,
    "crossover_funcs": crossover,
    "simulator": {"mode": "sampling", "calculator": calculator},
    "evaluator": {"features_funcs": features, "objectives_funcs": objectives},
    "hise": {"supercells": hise_supercells},
}

# 7. Run
if __name__ == "__main__":
    print(f"Starting EZGA for system: {composition}")
    run_hise(load_config(cfg_dict))
'''

def init_cli(
    output: Path = typer.Argument(Path("my_experiment.py"), help="Name of the generated python script"),
    force: bool = typer.Option(False, "--force", "-f", help="Overwrite the file if it already exists")
) -> None:
    """
    Scaffold a ready-to-use EZGA experiment script.
    
    Creates a functional Python script containing all the fundamental building 
    blocks (Constraints, Mutators, Calculators, Config) required to run an 
    evolutionary search from scratch.
    """
    if output.exists() and not force:
        console.print(f"[bold red]Error:[/] File '{output}' already exists. Use --force to overwrite.")
        raise typer.Exit(1)
        
    output.write_text(BOILERPLATE_CODE)
    
    msg = f"""[bold green]Success![/] Scaffolded EZGA experiment script at [bold]{output}[/].
    
[bold cyan]Next Steps:[/bold cyan]
1. Open [bold]{output}[/] and change `composition = ["Ni", "O"]` to your target elements.
2. Run it with: [bold]python {output}[/]"""
    console.print(Panel(msg, border_style="green"))
