import time

import numpy as np

from ase.atom import Atom
from ase.atoms import Atoms
from ase.calculators.lammpsrun import Prism
from ase.data import atomic_masses, chemical_symbols
from ase.io import read
from ase.neighborlist import NeighborList


def twochar(name):
    if len(name) > 1:
        return name[:2]
    else:
        return name + ' '


class BondData:
    def __init__(self, name_value_hash):
        self.nvh = name_value_hash

    def name_value(self, aname, bname):
        name1 = twochar(aname) + '-' + twochar(bname)
        name2 = twochar(bname) + '-' + twochar(aname)
        if name1 in self.nvh:
            return name1, self.nvh[name1]
        if name2 in self.nvh:
            return name2, self.nvh[name2]
        return None, None

    def value(self, aname, bname):
        return self.name_value(aname, bname)[1]


class CutoffList(BondData):
    def max(self):
        return max(self.nvh.values())


class AnglesData:
    def __init__(self, name_value_hash):
        self.nvh = name_value_hash

    def name_value(self, aname, bname, cname):
        for name in [
            (twochar(aname) + '-' + twochar(bname) + '-' + twochar(cname)),
                (twochar(cname) + '-' + twochar(bname) + '-' + twochar(aname))]:
            if name in self.nvh:
                return name, self.nvh[name]
        return None, None


class DihedralsData:
    def __init__(self, name_value_hash):
        self.nvh = name_value_hash

    def name_value(self, aname, bname, cname, dname):
        for name in [
            (twochar(aname) + '-' + twochar(bname) + '-' +
             twochar(cname) + '-' + twochar(dname)),
            (twochar(dname) + '-' + twochar(cname) + '-' +
             twochar(bname) + '-' + twochar(aname))]:
            if name in self.nvh:
                return name, self.nvh[name]
        return None, None


class OPLSff:
    def __init__(self, fileobj=None, warnings=0):
        self.warnings = warnings
        self.data = {}
        if fileobj is not None:
            self.read(fileobj)

    def read(self, fileobj, comments='#'):

        def read_block(name, symlen, nvalues):
            """Read a data block.

            name: name of the block to store in self.data
            symlen: length of the symbol
            nvalues: number of values expected
            """

            if name not in self.data:
                self.data[name] = {}
            data = self.data[name]

            def add_line():
                line = fileobj.readline().strip()
                if not len(line):  # end of the block
                    return False
                line = line.split('#')[0]  # get rid of comments
                if len(line) > symlen:
                    symbol = line[:symlen]
                    words = line[symlen:].split()
                    if len(words) >= nvalues:
                        if nvalues == 1:
                            data[symbol] = float(words[0])
                        else:
                            data[symbol] = [float(word)
                                            for word in words[:nvalues]]
                return True

            while add_line():
                pass

        read_block('one', 2, 3)
        read_block('bonds', 5, 2)
        read_block('angles', 8, 2)
        read_block('dihedrals', 11, 4)
        read_block('cutoffs', 5, 1)

        self.bonds = BondData(self.data['bonds'])
        self.angles = AnglesData(self.data['angles'])
        self.dihedrals = DihedralsData(self.data['dihedrals'])
        self.cutoffs = CutoffList(self.data['cutoffs'])

    def write_lammps(self, atoms, prefix='lammps'):
        """Write input for a LAMMPS calculation."""
        self.prefix = prefix

        if hasattr(atoms, 'connectivities'):
            connectivities = atoms.connectivities
        else:
            btypes, blist = self.get_bonds(atoms)
            atypes, alist = self.get_angles()
            dtypes, dlist = self.get_dihedrals(alist, atypes)
            connectivities = {
                'bonds': blist,
                'bond types': btypes,
                'angles': alist,
                'angle types': atypes,
                'dihedrals': dlist,
                'dihedral types': dtypes}

            self.write_lammps_definitions(atoms, btypes, atypes, dtypes)
            self.write_lammps_in()

        self.write_lammps_atoms(atoms, connectivities)

    def write_lammps_in(self):
        with open(self.prefix + '_in', 'w') as fileobj:
            self._write_lammps_in(fileobj)

    def _write_lammps_in(self, fileobj):
        fileobj.write("""# LAMMPS relaxation (written by ASE)

units           metal
atom_style      full
boundary        p p p
#boundary       p p f

""")
        fileobj.write('read_data ' + self.prefix + '_atoms\n')
        fileobj.write('include  ' + self.prefix + '_opls\n')
        fileobj.write("""
kspace_style    pppm 1e-5
#kspace_modify  slab 3.0

neighbor        1.0 bin
neigh_modify    delay 0 every 1 check yes

thermo          1000
thermo_style    custom step temp press cpu pxx pyy pzz pxy pxz pyz ke pe etotal vol lx ly lz atoms

dump            1 all xyz 1000 dump_relax.xyz
dump_modify     1 sort id

restart         100000 test_relax

min_style       fire
minimize        1.0e-14 1.0e-5 100000 100000
""")  # noqa: E501

    def write_lammps_atoms(self, atoms, connectivities):
        """Write atoms input for LAMMPS"""
        with open(self.prefix + '_atoms', 'w') as fileobj:
            self._write_lammps_atoms(fileobj, atoms, connectivities)

    def _write_lammps_atoms(self, fileobj, atoms, connectivities):
        # header
        fileobj.write(fileobj.name + ' (by ' + str(self.__class__) + ')\n\n')
        fileobj.write(str(len(atoms)) + ' atoms\n')
        fileobj.write(str(len(atoms.types)) + ' atom types\n')
        blist = connectivities['bonds']
        if len(blist):
            btypes = connectivities['bond types']
            fileobj.write(str(len(blist)) + ' bonds\n')
            fileobj.write(str(len(btypes)) + ' bond types\n')
        alist = connectivities['angles']
        if len(alist):
            atypes = connectivities['angle types']
            fileobj.write(str(len(alist)) + ' angles\n')
            fileobj.write(str(len(atypes)) + ' angle types\n')
        dlist = connectivities['dihedrals']
        if len(dlist):
            dtypes = connectivities['dihedral types']
            fileobj.write(str(len(dlist)) + ' dihedrals\n')
            fileobj.write(str(len(dtypes)) + ' dihedral types\n')

        # cell
        p = Prism(atoms.get_cell())
        xhi, yhi, zhi, xy, xz, yz = p.get_lammps_prism()
        fileobj.write(f'\n0.0 {xhi}  xlo xhi\n')
        fileobj.write(f'0.0 {yhi}  ylo yhi\n')
        fileobj.write(f'0.0 {zhi}  zlo zhi\n')

        if p.is_skewed():
            fileobj.write(f"{xy} {xz} {yz}  xy xz yz\n")

        # atoms
        fileobj.write('\nAtoms\n\n')
        tag = atoms.get_tags()
        if atoms.has('molid'):
            molid = atoms.get_array('molid')
        else:
            molid = [1] * len(atoms)
        for i, r in enumerate(
                p.vector_to_lammps(atoms.get_positions())):
            atype = atoms.types[tag[i]]
            if len(atype) < 2:
                atype = atype + ' '
            q = self.data['one'][atype][2]
            fileobj.write('%6d %3d %3d %s %s %s %s' % ((i + 1, molid[i],
                                                        tag[i] + 1,
                                                        q) + tuple(r)))
            fileobj.write(' # ' + atoms.types[tag[i]] + '\n')

        # velocities
        velocities = atoms.get_velocities()
        if velocities is not None:
            velocities = p.vector_to_lammps(atoms.get_velocities())
            fileobj.write('\nVelocities\n\n')
            for i, v in enumerate(velocities):
                fileobj.write('%6d %g %g %g\n' %
                              (i + 1, v[0], v[1], v[2]))

        # masses
        fileobj.write('\nMasses\n\n')
        for i, typ in enumerate(atoms.types):
            cs = atoms.split_symbol(typ)[0]
            fileobj.write('%6d %g # %s -> %s\n' %
                          (i + 1,
                           atomic_masses[chemical_symbols.index(cs)],
                           typ, cs))

        # bonds
        if blist:
            fileobj.write('\nBonds\n\n')
            for ib, bvals in enumerate(blist):
                fileobj.write('%8d %6d %6d %6d ' %
                              (ib + 1, bvals[0] + 1, bvals[1] + 1,
                               bvals[2] + 1))
                if bvals[0] in btypes:
                    fileobj.write('# ' + btypes[bvals[0]])
                fileobj.write('\n')

        # angles
        if alist:
            fileobj.write('\nAngles\n\n')
            for ia, avals in enumerate(alist):
                fileobj.write('%8d %6d %6d %6d %6d ' %
                              (ia + 1, avals[0] + 1,
                               avals[1] + 1, avals[2] + 1, avals[3] + 1))
                if avals[0] in atypes:
                    fileobj.write('# ' + atypes[avals[0]])
                fileobj.write('\n')

        # dihedrals
        if dlist:
            fileobj.write('\nDihedrals\n\n')
            for i, dvals in enumerate(dlist):
                fileobj.write('%8d %6d %6d %6d %6d %6d ' %
                              (i + 1, dvals[0] + 1,
                               dvals[1] + 1, dvals[2] + 1,
                               dvals[3] + 1, dvals[4] + 1))
                if dvals[0] in dtypes:
                    fileobj.write('# ' + dtypes[dvals[0]])
                fileobj.write('\n')

    def update_neighbor_list(self, atoms):
        cut = 0.5 * max(self.data['cutoffs'].values())
        self.nl = NeighborList([cut] * len(atoms), skin=0,
                               bothways=True, self_interaction=False)
        self.nl.update(atoms)
        self.atoms = atoms

    def get_bonds(self, atoms):
        """Find bonds and return them and their types"""
        cutoffs = CutoffList(self.data['cutoffs'])
        self.update_neighbor_list(atoms)

        types = atoms.get_types()
        tags = atoms.get_tags()
        cell = atoms.get_cell()
        bond_list = []
        bond_types = []
        for i, atom in enumerate(atoms):
            iname = types[tags[i]]
            indices, offsets = self.nl.get_neighbors(i)
            for j, offset in zip(indices, offsets):
                if j <= i:
                    continue  # do not double count
                jname = types[tags[j]]
                cut = cutoffs.value(iname, jname)
                if cut is None:
                    if self.warnings > 1:
                        print(f'Warning: cutoff {iname}-{jname} not found')
                    continue  # don't have it
                dist = np.linalg.norm(atom.position - atoms[j].position -
                                      np.dot(offset, cell))
                if dist > cut:
                    continue  # too far away
                name, _val = self.bonds.name_value(iname, jname)
                if name is None:
                    if self.warnings:
                        print(f'Warning: potential {iname}-{jname} not found')
                    continue  # don't have it
                if name not in bond_types:
                    bond_types.append(name)
                bond_list.append([bond_types.index(name), i, j])
        return bond_types, bond_list

    def get_angles(self, atoms=None):
        cutoffs = CutoffList(self.data['cutoffs'])
        if atoms is not None:
            self.update_neighbor_list(atoms)
        else:
            atoms = self.atoms

        types = atoms.get_types()
        tags = atoms.get_tags()
        cell = atoms.get_cell()
        ang_list = []
        ang_types = []

        # center atom *-i-*
        for i, atom in enumerate(atoms):
            iname = types[tags[i]]
            indicesi, offsetsi = self.nl.get_neighbors(i)

            # search for first neighbor j-i-*
            for j, offsetj in zip(indicesi, offsetsi):
                jname = types[tags[j]]
                cut = cutoffs.value(iname, jname)
                if cut is None:
                    continue  # don't have it
                dist = np.linalg.norm(atom.position - atoms[j].position -
                                      np.dot(offsetj, cell))
                if dist > cut:
                    continue  # too far away

                # search for second neighbor j-i-k
                for k, offsetk in zip(indicesi, offsetsi):
                    if k <= j:
                        continue  # avoid double count
                    kname = types[tags[k]]
                    cut = cutoffs.value(iname, kname)
                    if cut is None:
                        continue  # don't have it
                    dist = np.linalg.norm(atom.position -
                                          np.dot(offsetk, cell) -
                                          atoms[k].position)
                    if dist > cut:
                        continue  # too far away
                    name, _val = self.angles.name_value(jname, iname,
                                                       kname)
                    if name is None:
                        if self.warnings > 1:
                            print(
                                f'Warning: angles {jname}-{iname}-{kname} not '
                                'found'
                            )
                        continue  # don't have it
                    if name not in ang_types:
                        ang_types.append(name)
                    ang_list.append([ang_types.index(name), j, i, k])

        return ang_types, ang_list

    def get_dihedrals(self, ang_types, ang_list):
        'Dihedrals derived from angles.'

        cutoffs = CutoffList(self.data['cutoffs'])

        atoms = self.atoms
        types = atoms.get_types()
        tags = atoms.get_tags()
        cell = atoms.get_cell()

        dih_list = []
        dih_types = []

        def append(name, i, j, k, L):
            if name not in dih_types:
                dih_types.append(name)
            index = dih_types.index(name)
            if (([index, i, j, k, L] not in dih_list) and
                    ([index, L, k, j, i] not in dih_list)):
                dih_list.append([index, i, j, k, L])

        for angle in ang_types:
            L, i, j, k = angle
            iname = types[tags[i]]
            jname = types[tags[j]]
            kname = types[tags[k]]

            # search for l-i-j-k
            indicesi, offsetsi = self.nl.get_neighbors(i)
            for L, offsetl in zip(indicesi, offsetsi):
                if L == j:
                    continue  # avoid double count
                lname = types[tags[L]]
                cut = cutoffs.value(iname, lname)
                if cut is None:
                    continue  # don't have it
                dist = np.linalg.norm(atoms[i].position - atoms[L].position -
                                      np.dot(offsetl, cell))
                if dist > cut:
                    continue  # too far away
                name, _val = self.dihedrals.name_value(lname, iname,
                                                      jname, kname)
                if name is None:
                    continue  # don't have it
                append(name, L, i, j, k)

            # search for i-j-k-l
            indicesk, offsetsk = self.nl.get_neighbors(k)
            for L, offsetl in zip(indicesk, offsetsk):
                if L == j:
                    continue  # avoid double count
                lname = types[tags[L]]
                cut = cutoffs.value(kname, lname)
                if cut is None:
                    continue  # don't have it
                dist = np.linalg.norm(atoms[k].position - atoms[L].position -
                                      np.dot(offsetl, cell))
                if dist > cut:
                    continue  # too far away
                name, _val = self.dihedrals.name_value(iname, jname,
                                                      kname, lname)
                if name is None:
                    continue  # don't have it
                append(name, i, j, k, L)

        return dih_types, dih_list

    def write_lammps_definitions(self, atoms, btypes, atypes, dtypes):
        """Write force field definitions for LAMMPS."""
        with open(self.prefix + '_opls', 'w') as fd:
            self._write_lammps_definitions(fd, atoms, btypes, atypes, dtypes)

    def _write_lammps_definitions(self, fileobj, atoms, btypes, atypes,
                                  dtypes):
        fileobj.write('# OPLS potential\n')
        fileobj.write('# write_lammps' +
                      str(time.asctime(time.localtime(time.time()))))

        # bonds
        if len(btypes):
            fileobj.write('\n# bonds\n')
            fileobj.write('bond_style      harmonic\n')
            for ib, btype in enumerate(btypes):
                fileobj.write('bond_coeff %6d' % (ib + 1))
                for value in self.bonds.nvh[btype]:
                    fileobj.write(' ' + str(value))
                fileobj.write(' # ' + btype + '\n')

        # angles
        if len(atypes):
            fileobj.write('\n# angles\n')
            fileobj.write('angle_style      harmonic\n')
            for ia, atype in enumerate(atypes):
                fileobj.write('angle_coeff %6d' % (ia + 1))
                for value in self.angles.nvh[atype]:
                    fileobj.write(' ' + str(value))
                fileobj.write(' # ' + atype + '\n')

        # dihedrals
        if len(dtypes):
            fileobj.write('\n# dihedrals\n')
            fileobj.write('dihedral_style      opls\n')
            for i, dtype in enumerate(dtypes):
                fileobj.write('dihedral_coeff %6d' % (i + 1))
                for value in self.dihedrals.nvh[dtype]:
                    fileobj.write(' ' + str(value))
                fileobj.write(' # ' + dtype + '\n')

        # Lennard Jones settings
        fileobj.write('\n# L-J parameters\n')
        fileobj.write('pair_style lj/cut/coul/long 10.0 7.4' +
                      ' # consider changing these parameters\n')
        fileobj.write('special_bonds lj/coul 0.0 0.0 0.5\n')
        data = self.data['one']
        for ia, atype in enumerate(atoms.types):
            if len(atype) < 2:
                atype = atype + ' '
            fileobj.write('pair_coeff ' + str(ia + 1) + ' ' + str(ia + 1))
            for value in data[atype][:2]:
                fileobj.write(' ' + str(value))
            fileobj.write(' # ' + atype + '\n')
        fileobj.write('pair_modify shift yes mix geometric\n')

        # Charges
        fileobj.write('\n# charges\n')
        for ia, atype in enumerate(atoms.types):
            if len(atype) < 2:
                atype = atype + ' '
            fileobj.write('set type ' + str(ia + 1))
            fileobj.write(' charge ' + str(data[atype][2]))
            fileobj.write(' # ' + atype + '\n')


class OPLSStructure(Atoms):
    default_map = {
        'BR': 'Br',
        'Be': 'Be',
        'C0': 'Ca',
        'Li': 'Li',
        'Mg': 'Mg',
        'Al': 'Al',
        'Ar': 'Ar'}

    def __init__(self, filename=None, *args, **kwargs):
        Atoms.__init__(self, *args, **kwargs)
        if filename:
            self.read_extended_xyz(filename)
        else:
            self.types = []
            for atom in self:
                if atom.symbol not in self.types:
                    self.types.append(atom.symbol)
                atom.tag = self.types.index(atom.symbol)

    def append(self, atom):
        """Append atom to end."""
        self.extend(Atoms([atom]))

    def read_extended_xyz(self, fileobj, map={}):
        """Read extended xyz file with labeled atoms."""
        atoms = read(fileobj)
        self.set_cell(atoms.get_cell())
        self.set_pbc(atoms.get_pbc())

        types = []
        types_map = {}
        for atom, type in zip(atoms, atoms.get_array('type')):
            if type not in types:
                types_map[type] = len(types)
                types.append(type)
            atom.tag = types_map[type]
            self.append(atom)
        self.types = types

        # copy extra array info
        for name, array in atoms.arrays.items():
            if name not in self.arrays:
                self.new_array(name, array)

    def split_symbol(self, string, translate=default_map):

        if string in translate:
            return translate[string], string
        if len(string) < 2:
            return string, None
        return string[0], string[1]

    def get_types(self):
        return self.types

    def colored(self, elements):
        res = Atoms()
        res.set_cell(self.get_cell())
        for atom in self:
            elem = self.types[atom.tag]
            if elem in elements:
                elem = elements[elem]
            res.append(Atom(elem, atom.position))
        return res

    def update_from_lammps_dump(self, fileobj, check=True):
        atoms = read(fileobj, format='lammps-dump')

        if len(atoms) != len(self):
            raise RuntimeError('Structure in ' + str(fileobj) +
                               ' has wrong length: %d != %d' %
                               (len(atoms), len(self)))

        if check:
            for a, b in zip(self, atoms):
                # check that the atom types match
                if a.tag + 1 != b.number:
                    raise RuntimeError('Atoms index %d are of different '
                                       'type (%d != %d)'
                                       % (a.index, a.tag + 1, b.number))

        self.set_cell(atoms.get_cell())
        self.set_positions(atoms.get_positions())
        if atoms.get_velocities() is not None:
            self.set_velocities(atoms.get_velocities())
        # XXX what about energy and forces ???

    def read_connectivities(self, fileobj, update_types=False):
        """Read positions, connectivities, etc.

        update_types: update atom types from the masses
        """
        lines = fileobj.readlines()
        lines.pop(0)

        def next_entry():
            line = lines.pop(0).strip()
            if len(line) > 0:
                lines.insert(0, line)

        def next_key():
            while len(lines):
                line = lines.pop(0).strip()
                if len(line) > 0:
                    lines.pop(0)
                    return line
            return None

        next_entry()
        header = {}
        while True:
            line = lines.pop(0).strip()
            if len(line):
                w = line.split()
                if len(w) == 2:
                    header[w[1]] = int(w[0])
                else:
                    header[w[1] + ' ' + w[2]] = int(w[0])
            else:
                break

        while not lines.pop(0).startswith('Atoms'):
            pass
        lines.pop(0)

        natoms = len(self)
        positions = np.empty((natoms, 3))
        for i in range(natoms):
            w = lines.pop(0).split()
            assert int(w[0]) == (i + 1)
            positions[i] = np.array([float(w[4 + c]) for c in range(3)])
            # print(w, positions[i])

        key = next_key()

        velocities = None
        if key == 'Velocities':
            velocities = np.empty((natoms, 3))
            for i in range(natoms):
                w = lines.pop(0).split()
                assert int(w[0]) == (i + 1)
                velocities[i] = np.array([float(w[1 + c]) for c in range(3)])
            key = next_key()

        if key == 'Masses':
            ntypes = len(self.types)
            masses = np.empty(ntypes)
            for i in range(ntypes):
                w = lines.pop(0).split()
                assert int(w[0]) == (i + 1)
                masses[i] = float(w[1])

            if update_types:
                # get the elements from the masses
                # this ensures that we have the right elements
                # even when reading from a lammps dump file
                def newtype(element, types):
                    if len(element) > 1:
                        # can not extend, we are restricted to
                        # two characters
                        return element
                    count = 0
                    for type in types:
                        if type[0] == element:
                            count += 1
                    label = '0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZ'
                    return (element + label[count])

                symbolmap = {}
                typemap = {}
                types = []
                ams = atomic_masses[:]
                ams[np.isnan(ams)] = 0
                for i, mass in enumerate(masses):
                    m2 = (ams - mass)**2
                    symbolmap[self.types[i]] = chemical_symbols[m2.argmin()]
                    typemap[self.types[i]] = newtype(
                        chemical_symbols[m2.argmin()], types)
                    types.append(typemap[self.types[i]])
                for atom in self:
                    atom.symbol = symbolmap[atom.symbol]
                self.types = types

            key = next_key()

        def read_list(key_string, length, debug=False):
            if key != key_string:
                return [], key

            lst = []
            while len(lines):
                w = lines.pop(0).split()
                if len(w) > length:
                    lst.append([(int(w[1 + c]) - 1) for c in range(length)])
                else:
                    return lst, next_key()
            return lst, None

        bonds, key = read_list('Bonds', 3)
        angles, key = read_list('Angles', 4)
        dihedrals, key = read_list('Dihedrals', 5, True)

        self.connectivities = {
            'bonds': bonds,
            'angles': angles,
            'dihedrals': dihedrals
        }

        if 'bonds' in header:
            assert len(bonds) == header['bonds']
            self.connectivities['bond types'] = list(
                range(header['bond types']))
        if 'angles' in header:
            assert len(angles) == header['angles']
            self.connectivities['angle types'] = list(
                range(header['angle types']))
        if 'dihedrals' in header:
            assert len(dihedrals) == header['dihedrals']
            self.connectivities['dihedral types'] = list(range(
                header['dihedral types']))