"""
Macer: Machine-learning accelerated Atomic Computational Environment for automated Research workflows
Copyright (c) 2025 The Macer Package Authors
Author: Soungmin Bae <soungminbae@gmail.com>
License: MIT
"""

import importlib
from importlib.util import find_spec

# Mapping of force field names to their required python packages
# This allows checking availability without importing the package
FF_DEPENDENCIES = {
    "mace": "mace",
    "sevennet": "sevenn",
    "chgnet": "chgnet",
    "m3gnet": "matgl",
    "allegro": "nequip", # commonly nequip or allegro package
    "mattersim": "mattersim",
    "orb": "orb_models",
    "fairchem": "fairchem",
    "emt": "ase",
}

ALL_SUPPORTED_FFS = list(FF_DEPENDENCIES.keys())

def is_package_available(package_name):
    """Checks if a package is installed without importing it."""
    try:
        if package_name == "ase": return True # ase is required dependency
        return find_spec(package_name) is not None
    except ImportError:
        return False
    except Exception:
        # Handle cases where find_spec might crash on some environments
        return False

def get_available_ffs():
    """
    Returns a list of currently installed (available) force fields,
    with a preferred order for defaults. Uses lightweight checks.
    """
    available_ffs = []
    
    # Define the preferred order
    preferred_order = ["mattersim", "mace", "sevennet", "orb", "fairchem"]
    
    # Check preferred first
    for ff in preferred_order:
        if ff in FF_DEPENDENCIES and is_package_available(FF_DEPENDENCIES[ff]):
            available_ffs.append(ff)
            
    # Check others
    for ff in ALL_SUPPORTED_FFS:
        if ff not in available_ffs:
            if is_package_available(FF_DEPENDENCIES[ff]):
                available_ffs.append(ff)
                
    return available_ffs

def get_calculator(ff_name: str, **kwargs):
    """
    Creates an ASE Calculator instance for the given force field name.
    """
    # Import locally to avoid startup cost
    if ff_name == "mace":
        from .mace import get_mace_calculator as func
    elif ff_name == "sevennet":
        from .sevennet import get_sevennet_calculator as func
    elif ff_name == "chgnet":
        from .chgnet import get_chgnet_calculator as func
    elif ff_name == "m3gnet":
        from .m3gnet import get_m3gnet_calculator as func
    elif ff_name == "allegro":
        from .allegro import get_allegro_calculator as func
    elif ff_name == "mattersim":
        from .mattersim import get_mattersim_calculator as func
    elif ff_name == "orb":
        from .orb import get_orb_calculator as func
    elif ff_name == "fairchem":
        from .fairchem import get_fairchem_calculator as func
    elif ff_name == "emt":
        from .emt import get_emt_calculator as func
    else:
        raise ValueError(f"Unsupported force field: {ff_name}. Supported: {ALL_SUPPORTED_FFS}")
        
    return func(**kwargs)

