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

# ============================================================
# Internal utilities
# ============================================================

def _load_npz(filename):
    """Load an .npz file from the aegon/data/ package directory.

in:
    filename (str): relative filename within aegon/data/ (e.g. 'libdata_lj.npz').
out:
    numpy.lib.npyio.NpzFile: loaded NPZ archive object.
Raises:
    ValueError if the file is not found in the package data directory.
"""
    try:
        with resources.as_file(resources.files("aegon").joinpath(f"data/{filename}")) as path:
            return np.load(path, allow_pickle=True)
    except FileNotFoundError:
        raise ValueError(f"Data file not found: data/{filename}")


def _atoms_from_entry(data, N, symbol):
    """Build an ASE Atoms object from a cluster database entry dict.

in:
    data (dict): database entry with keys 'positions', 'energy', and 'label'.
    N (int): number of atoms in the cluster.
    symbol (str): atomic symbol to assign to all atoms.
out:
    ase.Atoms: cluster with info keys 'energy' and 'label' set from the database entry.
"""
    atoms = Atoms(symbols=[symbol] * N, positions=data["positions"])
    atoms.info["energy"] = data["energy"]
    atoms.info["label"] = data["label"]
    return atoms

# ============================================================
# ClusterFactory Core
# ============================================================

class ClusterFactory:
    """Factory class for retrieving preoptimized reference clusters from bundled databases.

    Supports Lennard-Jones (LJ) and Sutton-Chen (SC) model potentials.
    Databases are loaded lazily and cached after first access.
    """

    # Cached databases to avoid reloading from disk
    _LJ_DATA = None
    _SC_DATA = {}   # key: symbol → db

    # Supported models (LJ and SC for now)
    _SUPPORTED_MODELS = {"LJ", "SC"}

    # ---------------------- Public API ----------------------

    @classmethod
    def get(cls, N, model, element=None):
        """Retrieve a preoptimized cluster from the database.

in:
    N (int): number of atoms in the desired cluster.
    model (str): potential model, either 'LJ' or 'SC'.
    element (str, optional): atomic symbol required for SC clusters; for LJ
        clusters defaults to 'Mo' if not provided.
out:
    ase.Atoms: preoptimized cluster with info keys 'energy' and 'label'.
Raises:
    ValueError if the model is unsupported, element is missing for SC,
    or the requested cluster size is not in the database.
"""
        model = model.upper()

        if model not in cls._SUPPORTED_MODELS:
            raise ValueError(f"Unsupported model '{model}'. Supported models: {cls._SUPPORTED_MODELS}")

        if model == "LJ":
            return cls._get_lj(N, element or "Mo")  # default LJ element is Mo

        elif model == "SC":
            if element is None:
                raise ValueError("Element must be specified for SC clusters.")
            return cls._get_sc(N, element)

    @classmethod
    def list_available(cls, model, element=None):
        """List available cluster sizes for a given potential model.

in:
    model (str): potential model, either 'LJ' or 'SC'.
    element (str, optional): atomic symbol required for SC clusters.
out:
    list of int: sorted list of available cluster sizes (number of atoms).
Raises:
    ValueError if the model is unsupported or element is missing for SC.
"""
        model = model.upper()

        if model not in cls._SUPPORTED_MODELS:
            raise ValueError(f"Unsupported model '{model}'.")

        if model == "LJ":
            return cls._list_available_lj()

        elif model == "SC":
            if element is None:
                raise ValueError("Element must be specified for SC clusters.")
            return cls._list_available_sc(element)

    # ---------------------- LJ subsystem ----------------------

    @classmethod
    def _load_lj_database(cls):
        """Load and cache the LJ cluster database from the package data directory.

in:
    None.
out:
    numpy.lib.npyio.NpzFile: the loaded LJ NPZ database.
"""
        if cls._LJ_DATA is None:
            cls._LJ_DATA = _load_npz("libdata_lj.npz")
        return cls._LJ_DATA

    @classmethod
    def _get_lj(cls, N, symbol):
        """Retrieve an LJ cluster of size N with the given atomic symbol.

in:
    N (int): number of atoms.
    symbol (str): atomic symbol for all atoms.
out:
    ase.Atoms: preoptimized LJ cluster.
"""
        db = cls._load_lj_database()
        key = f"LJ{N:03d}"

        if key not in db:
            raise ValueError(f"No LJ cluster available for N = {N}")

        data = db[key].item()
        return _atoms_from_entry(data, N, symbol)

    @classmethod
    def _list_available_lj(cls):
        """Return sorted list of LJ cluster sizes in the database.

in:
    None.
out:
    list of int: sorted cluster sizes.
"""
        db = cls._load_lj_database()
        return sorted(int(k[2:]) for k in db.keys())

    # ---------------------- SC subsystem ----------------------

    @classmethod
    def _load_sc_database(cls, element):
        """Load and cache the Sutton-Chen cluster database for a given element.

in:
    element (str): atomic symbol (capitalized automatically).
out:
    numpy.lib.npyio.NpzFile: the loaded SC NPZ database for the element.
"""
        element = element.capitalize()

        if element not in cls._SC_DATA:
            filename = f"SC_{element}_clusters_data.npz"
            cls._SC_DATA[element] = _load_npz(filename)

        return cls._SC_DATA[element]

    @classmethod
    def _get_sc(cls, N, element):
        """Retrieve an SC cluster of size N for the given element.

in:
    N (int): number of atoms.
    element (str): atomic symbol for the metal element.
out:
    ase.Atoms: preoptimized SC cluster.
"""
        element = element.capitalize()
        db = cls._load_sc_database(element)

        key = f"SC{N:03d}"
        if key not in db:
            raise ValueError(f"No SC cluster available for element {element} with N = {N}")

        data = db[key].item()
        return _atoms_from_entry(data, N, element)

    @classmethod
    def _list_available_sc(cls, element):
        """Return sorted list of SC cluster sizes for a given element.

in:
    element (str): atomic symbol (capitalized automatically).
out:
    list of int: sorted cluster sizes.
"""
        element = element.capitalize()
        db = cls._load_sc_database(element)
        return sorted(int(k[2:]) for k in db.keys())
