
_M3GNET_AVAILABLE = False
try:
    import matgl
    from matgl.ext.ase import PESCalculator
    _M3GNET_AVAILABLE = True
except (ImportError, Exception):
    try:
        # On some systems (e.g. macOS with Torch 2.4+), 
        # DGL / Graphbolt might fail to load some libraries.
        # We try to mock dgl.graphbolt if it's the culprit.
        import sys
        from unittest.mock import MagicMock
        sys.modules['dgl.graphbolt'] = MagicMock()
        import matgl
        from matgl.ext.ase import PESCalculator
        _M3GNET_AVAILABLE = True
    except Exception:
        pass

from macer.defaults import resolve_model_path

def get_m3gnet_calculator(model_path=None, device='cpu', **kwargs):
    if not _M3GNET_AVAILABLE:
        raise RuntimeError('M3GNet (matgl) is not installed. Please install with "pip install .[m3gnet]"')

    if model_path is None:
        model_path = 'M3GNet-MP-2021.2.8-PES'
    else:
        # Check if it's a known pretrained name or a local path
        resolved = resolve_model_path(model_path)
        if resolved and resolved != model_path:
             model_path = resolved
    
    potential = matgl.load_model(model_path)
    return PESCalculator(potential=potential)
