import numpy as np
from importlib import resources
from ase import Atoms

# Load the compressed Lennard-Jones cluster dataset included in the package
with resources.as_file(resources.files("aegon").joinpath("data/libdata_lj.npz")) as path:
    _LJ_DATA = np.load(path, allow_pickle=True)

def get_lj_cluster(N, symbol="Mo"):
    """Return a preoptimized Lennard-Jones (LJ) cluster as an ASE Atoms object.

in:
    N (int): number of atoms in the cluster (3 <= N <= 150).
    symbol (str): atomic symbol for all atoms in the cluster (default 'Mo').
out:
    ase.Atoms: ASE Atoms object with info keys 'e' (LJ energy) and 'i' (cluster label).
Raises:
    ValueError if the requested cluster size is not available in the database.

Notes:
    The data were obtained from the Wales database of optimized Lennard-Jones
    clusters (sizes 3-150). Each cluster includes optimized Cartesian coordinates,
    total LJ potential energy, and a cluster label (e.g. 'LJ038').
"""
    key = f"LJ{N:03d}"
    if key not in _LJ_DATA:
        raise ValueError(f"No Lennard-Jones structure available for N = {N}")

    data = _LJ_DATA[key].item()
    atoms = Atoms(symbols=[symbol] * N, positions=data["positions"])
    atoms.info['e'] = data["energy"]
    atoms.info['i'] = data["label"]
    return atoms

def list_available():
    """Return a sorted list of available LJ cluster sizes in the database.

in:
    None.
out:
    list of int: cluster sizes available in the bundled LJ database.
"""
    return sorted(int(k[2:]) for k in _LJ_DATA.keys())
