import os
from collections import OrderedDict

from pymatgen.core import Specie, Element, Structure
from pymatgen.core.periodic_table import _pt_data as periodic_table
import numpy as np

from .core import LammpsBox, LammpsPotentials

class LammpsInput:
    def __init__(self, lammps_script, lammps_data, additional_files=None):
        self.lammps_data = lammps_data
        self.lammps_script = lammps_script
        self.additional_files = additional_files or []

    def write_input(self, output_dir, input_filename="lammps.in", make_dir=True):
        if make_dir and not os.path.exists(output_dir):
            os.makedirs(output_dir)

        self.lammps_data.write_file(os.path.join(output_dir, self.lammps_script['read_data']))
        self.lammps_script.write_file(os.path.join(output_dir, input_filename))

        for file_buffer, filename in self.additional_files:
            with open(os.path.join(output_dir, filename)) as f:
                f.write(file_buffer)


class LammpsScript(OrderedDict):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

    def __str__(self):
        lines = ""
        for k1, v1 in self.items():
            if isinstance(v1, dict):
                v1 = v1.values()
            if isinstance(v1, list):
                for x in v1:
                    lines = "".join([lines, "{} ".format(k1)])
                    lines = "".join([lines, str(x), os.linesep])
            else:
                lines = "".join([lines, "{} ".format(k1)])
                lines = "".join([lines, " {}{}".format(str(v1), os.linesep)])
        return lines

    @property
    def log_filename(self):
        log = self.get('log', 'log.lammps')
        if log is None:
            log = "log.lammps"
        elif isinstance(log, (list, tuple)):
            log = log[-1]
        return str(log).split(' ')[0]

    @property
    def data_filenames(self):
        data = self.get('read_data', None)
        if data is None:
            return []
        elif isinstance(data, (list, tuple)):
            return [str(d).split()[0] for d in data]
        else:
            return [str(data).split()[0]]

    @property
    def dump_filename(self):
        dump = self.get('dump', None)
        if dump is None:
            return None
        elif isinstance(dump, (list, tuple)):
            dump = log[-1]

        dump = str(dump).split()
        if len(dump) < 4:
            raise ValueError('Invalid dump command: %s' % dump)
        return dump[4]

    def write_file(self, filename):
        with open(filename, 'w') as f:
            f.write(str(self))


class LammpsData:
    HEADER_TAGS = {
        'atoms', 'bonds', 'angles', 'dihedrals', 'impropers',
        'atom types', 'bond types', 'angle types', 'dihedral types', 'improper types',
        'extra bond per atom', 'extra angle per atom', 'extra dihedral per atom', 'extra improper per atom', 'extra special per atom',
        'ellipsoids', 'lines', 'triangles', 'bodies', 'xlo xhi', 'ylo yhi', 'zlo zhi', 'xy xz yz'
    }
    # Otherwise implied integer
    FLOAT_HEADER_TAGS = {'xlo xhi', 'ylo yhi', 'zlo zhi', 'xy xz yz'}
    REQUIRED_HEADER_TAGS = {'xlo xhi', 'ylo yhi', 'zlo zhi', 'atoms', 'atom types'}

    REQUIRED_SECTION_TAGS = {'Atoms', 'Masses'}
    CHECKED_SECTION_TAGS = {'Atoms', 'Pair Coeffs', 'PairIJ Coeffs', 'Bond Coeffs', 'Angle Coeffs', 'Dihedral Coeffs', 'Improper Coeffs'}
    SECTIONS_TAGS = {
        'Atoms', 'Velocities', 'Masses', 'Ellipsoids', 'Lines', 'Triangles', 'Bodies',
        'Bonds', 'Angles', 'Dihedrals', 'Impropers',
        'Pair Coeffs', 'PairIJ Coeffs', 'Bond Coeffs', 'Angle Coeffs', 'Dihedral Coeffs', 'Improper Coeffs',
        'BondBond Coeffs', 'BondAngle Coeffs', 'MiddleBondTorsion Coeffs', 'EndBondTorsion Coeffs',
        'AngleTorsion Coeffs', 'AngleAngleTorsion Coeffs', 'BondBond13 Coeffs', 'AngleAngle Coeffs'
    }


    def __init__(self, name, symbol_indicies, masses, atoms, lammps_box, potentials=None, velocities=None):
        self.name = name
        self.symbol_indicies = symbol_indicies
        self.masses = masses
        self.atoms = atoms
        self.lammps_box = lammps_box
        self.potentials = potentials
        self.velocities = velocities
        if self.potentials:
            self.potentials.symbol_indicies = self.symbol_indicies

    @classmethod
    def from_structure(cls, structure, potentials=None, include_charge=False, include_velocities=False):
        lammps_box = LammpsBox.from_lattice(structure.lattice)
        name = 'pymatgen autogenerated data file'
        symbol_indicies = {}
        masses = {}
        atoms = []
        velocities = None
        if include_velocities:
            velocities = []
        counter = 1
        for atom in structure:
            if atom.specie not in symbol_indicies:
                symbol_indicies[atom.specie] = counter
                counter += 1

            if isinstance(atom.specie, Specie):
                element = atom.specie.element
                charge = atom.specie.oxi_state
            elif isinstance(atom.specie, Element):
                element = atom.specie
                charge = 0.0

            if atom.specie not in masses:
                masses[atom.specie] = element.atomic_mass

            if include_velocities:
                velocities.append(atom.properties.get('velocity', [0, 0, 0]))

            atoms.append([atom.specie, charge, atom.coords])
        return cls(name, symbol_indicies, masses, atoms, lammps_box,
                   potentials=potentials, velocities=velocities)

    @classmethod
    def _parse_data_file(cls, filename):
        def is_header(line):
            for header in cls.HEADER_TAGS:
                if header in line:
                    data = line.replace(header, '').strip().split()
                    if header in cls.FLOAT_HEADER_TAGS:
                        value = list(map(float, data))
                    else:
                        value = int(data[0])
                    return header, value
            return None

        def is_section(line):
            for section in cls.SECTIONS_TAGS:
                if section in line:
                    return section
            return None

        def parseline(line):
            line = line.strip()
            comment = ''
            if '#' in line:
                comment_index = line.find('#')
                comment = line[comment_index+1:].strip()
                line = line[:comment_index]
            return line, comment

        def data_to_array(data):
            if len(data) == 0:
                return np.array([])

            # either int or floats
            is_integer = np.vectorize(lambda f:f.is_integer())
            row_types = is_integer(data[0])
            for row in data:
                row_types  = row_types * is_integer(row)
            names = ['f%d' % i for i in range(len(row_types))]
            formats = [np.int64 if _ else np.float64 for _ in row_types]
            return np.rec.array(data, dtype={'names': names, 'formats': formats})

        headers = {}
        sections = {}

        with open(filename) as f:
            description = f.readline()
            f.readline() # second line blank

            in_header = True
            current_section = [None, None, []]

            for line in f:
                line, comment = parseline(line)
                if not line:
                    continue # Blank line or comment

                if in_header:
                    header_found = is_header(line)
                    if header_found:
                        header, value = header_found
                        headers[header] = value
                    elif is_section(line):
                        in_header = False
                        current_section[0] = is_section(line)
                        current_section[1] = comment or None
                        current_section[2] = []
                    else:
                        raise ValueError('line: %s not recognized' % line)
                elif is_section(line):
                    if current_section[0]:
                        section, check, data = current_section
                        sections[section] = {'data': data_to_array(data), 'check': check}
                    current_section[0] = is_section(line)
                    current_section[1] = comment or None
                    current_section[2] = []
                else: # Reading data from section
                    current_section[2].append(tuple(map(float, line.split())))
            if current_section[0]:
                section, check, data = current_section
                sections[section] = {'data': data_to_array(data), 'check': check}
        return description, headers, sections

    @classmethod
    def _validate_data_file(cls, headers, sections):
        for required_header_tag in cls.REQUIRED_HEADER_TAGS:
            if required_header_tag not in headers:
                return 'Required header tag %s not included' % required_header_tag

        for required_section_tag in cls.REQUIRED_SECTION_TAGS:
            if required_section_tag not in sections:
                return 'Required section tag %s not included' % required_section_tag
        # Obviously can do more checks

    @classmethod
    def from_file(cls, filename, atom_style=None):
        description, headers, sections = cls._parse_data_file(filename)
        errors = cls._validate_data_file(headers, sections)
        if errors:
            raise ValueError('data file is invalid: {}'.format(errors))

        xlo, xhi = headers['xlo xhi']
        ylo, yhi = headers['ylo yhi']
        zlo, zhi = headers['zlo zhi']
        xy, xz, yz = headers.get('xy xz yz', [0, 0, 0])
        lammps_box = LammpsBox(xhi, yhi, zhi, xlo, ylo, zlo, xy, xz, yz)

        # Guess symbol from closest atomic mass (yes not great for isotopes)
        elements = np.array(
            [tuple([Element(element).atomic_mass, Element(element)]) for element in periodic_table],
            dtype={'names': ['mass', 'symbol'], 'formats': [np.float64, np.chararray]}
        )
        symbol_indicies = {}
        index_symbols = {} # Makes element lookup quicker
        masses = {}
        for index, atomic_mass, in sections['Masses']['data']:
            symbol = elements['symbol'][np.abs(elements['mass'] - atomic_mass).argmin()]
            symbol_indicies[Element(symbol)] = index
            index_symbols[index] = Element(symbol)
            masses[Element(symbol)] = atomic_mass

        # Default full format or use check format
        atoms = []
        sections['Atoms']['data'].sort(order='f0') # f0 is default name of first field
        for atom in sections['Atoms']['data']:
            if sections['Atoms'].get('check') == 'full':
                atom_type, charge, *position = atom['f2'], atom['f3'], atom['f4'], atom['f5'], atom['f6']
            else:
                index, mol, atom_type, charge, *position = atom

            element = index_symbols[atom_type]
            atoms.append([element, charge, position])

        velocities = None
        if 'Velocities' in sections:
            sections['Velocities']['data'].sort(order='f0') # f0 is default name of first field
            velocities = [[vx, vy, vz] for vx, vy, vz in sections['Velocities']['data'][['f1', 'f2', 'f3']]]

        # Get Potentials
        # TODO only gets pair potentials for now and no reason to keep str
        pair_potentials = {}
        if 'PairIJ Coeffs' in sections:
            for s1, s2, *parameters in sections['PairIJ Coeffs']['data']:
                pair_potentials[(index_symbols[s1], index_symbols[s2])] = ' '.join(list(map(str, parameters)))

        potentials = LammpsPotentials(pair_potentials, symbol_indicies)

        return cls(description, symbol_indicies, masses, atoms, lammps_box,
                   potentials=potentials, velocities=velocities)

    @property
    def structure(self):
        lattice = self.lammps_box.lattice
        species = []
        positions = []
        for specie, charge, coords in self.atoms:
            positions.append(coords)
            if isinstance(specie, Specie):
                symbol = specie.element.symbol
            else:
                symbol = specie.symbol

            if charge:
                species.append(Specie(symbol, charge))
            else:
                species.append(Element(symbol))

        site_properties = {}
        if self.velocities:
            site_properties['velocities'] = self.velocities

        return Structure(lattice, species, positions, coords_are_cartesian=True, site_properties=site_properties)


    def __str__(self):
        lammps_data_str = [
            '{}\n'.format(self.name),
            '{} atoms\n'.format(len(self.atoms)),
            '{} atom types\n'.format(len(self.symbol_indicies)),
            '{}\n'.format(str(self.lammps_box)),
            'Masses\n',
            '\n'.join(['{} {}'.format(self.symbol_indicies[specie], float(mass)) for specie, mass in self.masses.items()]) + '\n',
            'Atoms\n',
            '\n'.join(['{} {} {} {} {} {} {}'.format(i+1, 1, self.symbol_indicies[specie], charge, *coords) for i, (specie, charge, coords) in enumerate(self.atoms)]) + '\n',
        ]

        if self.velocities:
            lammps_data_str.extend([
                'Velocities\n',
                '\n'.join(['{} {} {} {}'.format(i+1, *velocity) for i, velocity in enumerate(self.velocities)]) + '\n'
            ])

        if self.potentials:
            lammps_data_str.append('{}\n'.format(str(self.potentials)))

        return '\n'.join(lammps_data_str)

    def write_file(self, filename):
        with open(filename, 'w') as f:
            f.write(str(self))
