import os.path as op

from ase.db.row import AtomsRow
from ase.io.jsonio import read_json


class Collection:
    """Collection of atomic configurations and associated data.

    Example of use:

    >>> from ase.collections import s22
    >>> len(s22)
    22
    >>> s22.names[:3]
    ['Ammonia_dimer', 'Water_dimer', 'Formic_acid_dimer']
    >>> dimer = s22['Water_dimer']
    >>> dimer.get_chemical_symbols()
    ['O', 'H', 'H', 'O', 'H', 'H']
    >>> s22.data['Ammonia_dimer']
    {'cc_energy': -0.1375}
    >>> sum(len(atoms) for atoms in s22)
    414
    """

    def __init__(self, name):
        """Create a collection lazily.

        Will read data from json file when needed.

        A collection can be iterated over to get the Atoms objects and indexed
        with names to get individual members.

        Attributes:

        name: str
            Name of collection.
        data: dict
            Data dictionary.
        filename: str
            Location of json file.
        names: list
            Names of configurations in the collection.
        """

        self.name = name
        self._names = []
        self._systems = {}
        self._data = {}
        self.filename = op.join(op.dirname(__file__), f'{name}.json')

    def __getitem__(self, name):
        self._read()
        return self._systems[name].copy()

    def has(self, name):
        # Not __contains__() because __iter__ yields the systems.
        self._read()
        return name in self._systems

    def __iter__(self):
        for name in self.names:
            yield self[name]

    def __len__(self):
        return len(self.names)

    def __str__(self):
        return '<{}-collection, {} systems: {}, {}, ...>'.format(
            self.name, len(self), *self.names[:2])

    def __repr__(self):
        return f'Collection({self.name!r})'

    @property
    def names(self):
        self._read()
        return list(self._names)

    @property
    def data(self):
        self._read()
        return self._data

    def _read(self):
        if self._names:
            return
        bigdct = read_json(self.filename)
        for id in bigdct['ids']:
            dct = bigdct[id]
            kvp = dct['key_value_pairs']
            name = str(kvp['name'])
            self._names.append(name)
            self._systems[name] = AtomsRow(dct).toatoms()
            del kvp['name']
            self._data[name] = {str(k): v for k, v in kvp.items()}