import numpy as np
from ase.data import atomic_numbers, chemical_symbols, atomic_masses

# XXX todo: atoms.masses[1] = 2.0 fails, but atoms.masses[1:2] = 2.0 succeeds
#           atoms.(anything)[3:6] = [three things] fails, but should succeed

# Presently everything is backed by dictionaries and lists
# We can see if dictionaries/lists will reduce e.g. GUI framerate.
# (Probably won't.)
# Nevertheless we can back things by numpy arrays instead if necessary.
#
# e.g. list comprehension of dictionary lookups:
#    [kinds[id].mass for id in element_ids]
# would become numpy array slice:
#    masses[element_ids]
# We just have to dynamically grow and shrink arrays with number of kinds.

class Kind:
    """Represents a single chemical species (label, Z, mass).

    A Kind also has an index that is used to internally
    enumerate the chemical species of an Atoms object."""
    def __init__(self, label, Z, mass, id):
        self.label = label
        self.Z = Z
        self.mass = mass
        self.id = id

    def __repr__(self):
        return '{}:Z={}:M={}:id={}'.format(self.label, self.Z,
                                           self.mass, self.id)


class AtomKinds:
    """A collection of distinct chemical species."""

    def __init__(self):
        self.kinds = {}  # label -> kind
        self.kindslist = []  # species number -> kind
        self.kinds_lookup_table = {}  # (Z, mass) -> kind

    def autogenerate_label(self, label):
        while label in self.kinds:
            label += '*'
        return label

    def assign_by_mass(self, species, mass):
        newspecies = np.array(species, copy=True)
        for i, id in enumerate(species):
            oldkind = self.kindslist[id]
            if oldkind.mass != mass:
                # Check whether this Z/mass combination already exists:
                newkind = self.kinds_lookup_table.get((oldkind.Z, mass))
                if newkind is None:
                    newlabel = self.autogenerate_label(oldkind.label)
                    newkind = self.new_kind(newlabel, oldkind.Z, mass)
                newspecies[i] = newkind.id
        return newspecies

    def assign_by_number(self, numbers):
        """Create new kinds as necessary for the given atomic numbers.

        Only valid for real elements or Z=0."""
        return self.assign([chemical_symbols[Z] for Z in numbers])

    def assign(self, labels):
        """Create new kinds as necessary from list of labels.

        Labels may be any string; duplicates will resolve to the same kind.

        Return the species IDs."""
        for label in np.unique(labels):
            if label not in self.kinds:
                # We could allow special syntax here:
                #   'O_semicore:Z=8:mass=16'
                Z = atomic_numbers.get(label, 0)
                self.new_kind(label, Z)
        return self.get_indices(labels)

    def has_kind(self, Z, mass=None):
        if mass is None:
            mass = atomic_masses[Z]
        return (Z, mass) in self.kinds_lookup_table

    def kind(self, Z, mass=None):
        if mass is None:
            mass = atomic_masses[Z]
        kind = self.kinds_lookup_table.get((Z, mass))
        if kind is None:
            sym = chemical_symbols[Z]
            kind = self.new_kind(Z, mass)
        return kind

    def new_kind(self, label, Z, mass=None):
        num = len(self.kinds)
        assert label not in self.kinds
        if mass is None:
            mass = atomic_masses[Z]
        kind = Kind(label, Z, mass, num)
        self.kinds[label] = kind
        self.kindslist.append(kind)
        self.kinds_lookup_table[(Z, mass)] = kind
        return kind

    def expand(self, species):
        return [self.kindslist[i] for i in species]

    def get_labels(self, species):
        return [self.kindslist[i].label for i in species]

    def get_masses(self, species):
        return np.array([k.mass for k in self.expand(species)])

    def get_numbers(self, species):
        return np.array([k.Z for k in self.expand(species)])

    def get_indices(self, labels):
        return np.array([self.kinds[label].id for label in labels])

    def __repr__(self):
        return 'AtomKinds({})'.format(self.kindslist)


class Magic(object):
    def __init__(self, atoms, get, set):
        self.atoms = atoms
        self.get = get
        self.set = set

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

    def __iter__(self):
        return iter(self.get(self.atoms.species_ids))

    def __setitem__(self, index, value):
        self.atoms.species_ids[index] = self.set([value])

    def __getitem__(self, index):
        # index may be integer or slice object
        species = self.atoms.species_ids[index]
        if hasattr(species, '__len__'):
            return self.get(species)
        else:
            return self.get([species])[0]

    def __repr__(self):
        return '{}({})'.format(self.__class__.__name__, list(self))


class Labels(Magic):
    def __init__(self, atoms):
        super(Labels, self).__init__(atoms,
                                     atoms.kinds.get_labels,
                                     atoms.kinds.assign)

class Numbers(Magic):
    def __init__(self, atoms):
        super(Numbers, self).__init__(atoms,
                                      atoms.kinds.get_numbers,
                                      atoms.kinds.assign_by_number)

class Masses(Magic):
    def __init__(self, atoms):
        super(Masses, self).__init__(atoms,
                                     atoms.kinds.get_masses,
                                     atoms.kinds.assign_by_mass)

    def __setitem__(self, index, value):
        species = self.atoms.species_ids[index]
        self.atoms.species_ids[index] = self.set(species, value)


class AtomsWithLabels:
    def __init__(self, labels):
        #natoms = len(labels)

        self.kinds = AtomKinds()
        self.species_ids = self.kinds.assign(labels)
        self.labels = Labels(self)
        self.numbers = Numbers(self)
        self.masses = Masses(self)

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

    def __repr__(self):
        return 'AtomsWithLabels({})'.format(self.get_labels())

    def get_labels(self):
        return self.kinds.get_labels(self.species_ids)

    def get_masses(self):
        return self.kinds.get_masses(self.species_ids)


def main():
    atoms = AtomsWithLabels(['H'] * 10)
    print(atoms)
    atoms.numbers[2:5] = 6  # Carbon
    print(atoms)
    atoms.masses[5:8] = 2.0  # Some elements now have mass 2; new label H*
    # (If we assign more new masses, labels become H**, H***, ...)
    print(atoms)
    atoms.labels[2:4] = 'O'  # Oxygen
    print(atoms)
    atoms.labels[8:] = 'Q'  # atom number 0
    print(atoms)
    atoms.numbers[8:] = 1  # discards 'Q', now just 'H'
    print(atoms)
    atoms.labels[9:] = 'H*'  # H* recognized! mass will be 2.0.
    print(atoms)
    print(atoms.get_masses())

    # Summary:
    # * Updates of the Atoms object will dynamically create new species
    # * Species are secretly enumerated, but this is supposed to be 'invisible'
    # * Overwriting Z discards label and mass (e.g. Z=79 ----> Au, always)
    # * Overwiting label discards Z and mass
    #   This is probably okay for real elements ('Au') and also
    #   'jellium' should default to Z=0, mass=1.0 as it does now.
    #   But we should make it possible to overwrite 'O' with 'O_special',
    #   which is still somehow keep Z=8.  How should we do this?
    # * Overwriting masses will append '*' to the label, except if the
    #   resulting species is equal to an existing one; in that case, it
    #   will retrieve the existing label of that species.
    # * Atoms with same label but distinct masses are Au, Au*, Au**, Au***, ...

    from ase.collections import g2

    symbols = []
    for atoms in g2:
        symbols.extend(atoms.get_chemical_symbols())
        symbols.append('jellium')
        if len(symbols) > 40:
            break


    a = AtomsWithLabels(symbols)
    #print(a)
    print(a.labels[3])
    print(a.labels[3:7])
    a.numbers[8:12] = 0

    a.labels[3] = 'Au'
    print(a.labels)
    hmm = a.labels[3]
    print(hmm)
    assert hmm == 'Au'
    a.labels[3:7] = 'Pu'
    a.labels[17] = 'jellium'
    #a.labels[15] = 'Au:ghost'
    assert all(label == 'Pu' for label in a.labels[3:7])

    a.masses[5:10] = 42.0
    print(a)

    a.masses[12:15] = 10.
    a.masses[:4] = 3.

    print(a)
    #print(a.get_masses())
    print(a.kinds)

    #print(a.labels[3:7])
    #a.labels[3:7] = 'Pu'
    #print(a.labels[:10])
    #a.labels[3] = 'Ag'
    #print(a)
    #a.labels[3:10] = 'Pu'
    #print(a)
    #a.labels[10:20:2] = 'Nd'
    #print(a)
    #print(a.labels[:5])
    #print(a.labels[5])

if __name__ == '__main__':
    main()