# coding: utf-8
# Copyright (c) Tingzheng Hou.
# Distributed under the terms of the MIT License.

"""
This module implements a core class PackmolWrapper for packing molecules
into a single box.

You need the Packmol package to run the code, see
http://m3g.iqm.unicamp.br/packmol or
http://leandro.iqm.unicamp.br/m3g/packmol/home.shtml
for download and setup instructions. You may need to manually
set the folder of the packmol executable to the PATH environment variable.
"""

import subprocess
import os

__author__ = "Tingzheng Hou"
__version__ = "1.0"
__maintainer__ = "Tingzheng Hou"
__email__ = "tingzheng_hou@berkeley.edu"
__date__ = "Feb 9, 2021"


class PackmolWrapper:
    """
    Wrapper for the Packmol software that can be used to pack various types of
    molecules into a one single unit.

    Args:
        path (str): The path to the directory for file i/o.
        structures (list): A list containing dict information of molecules.
            Each molecule have two keys, "name", the structure name,
            and "file", the path to the structure file.
        numbers (dict): A dict of the numbers of each molecule
        box (list): A list of xlo, ylo, zlo, xhi, yhi, zhi, in Å.
        tolerance (float): Tolerance for packmol.
        seed (int): Random seed for packmol.
        inputfile (str): Path to the input file. Default to 'packmol.inp'.
        outputfile (str): Path to the output file. Default to 'output.xyz'.

    Examples:

        >>> structures = [{"name": "structure_name", "file": "/path/to/xyz"}]
        >>> pw = PackmolWrapper("/path/to/work/dir",
        ...                     structures,
        ...                     {"structure_name": '2'},
        ...                     [0., 0., 0., 10., 10., 10.]
        ... )
        >>> pw.make_packmol_input()
        >>> pw.run_packmol()
    """
    def __init__(self, path, structures, numbers, box, tolerance=None,
                 seed=None, inputfile='packmol.inp', outputfile='output.xyz'):
        self.path = path
        self.input = os.path.join(self.path, inputfile)
        self.output = os.path.join(self.path, outputfile)
        self.screen = os.path.join(self.path, 'packmol.stdout')
        self.structures = structures
        self.numbers = numbers
        self.box = box
        if tolerance is None:
            self.tolerance = 2.0
        else:
            self.tolerance = tolerance
        if seed is None:
            self.seed = 123
        else:
            self.seed = seed

    def run_packmol(self):
        """Run packmol and write out the packed structure."""
        try:
            p = subprocess.run('packmol < {}'.format(self.input),
                               check=True,
                               shell=True,
                               stdout=subprocess.PIPE,
                               stderr=subprocess.PIPE)
        except subprocess.CalledProcessError as e:
            raise ValueError("Packmol failed with errorcode {}"
                             " and stderr: {}".format(e.returncode, e.stderr))
        else:
            with open(self.screen, 'w') as out:
                out.write(p.stdout.decode())

    def make_packmol_input(self):
        """Make a Packmol usable input file."""

        with open(self.input, 'w') as out:
            out.write("# " + ' + '.join(self.numbers[structure["name"]] + " "
                                        + structure["name"] for structure
                                        in self.structures) + "\n")
            out.write("# Packmol input generated by mdgo.\n")
            out.write('seed {}\n'.format(self.seed))
            out.write('tolerance {}\n\n'.format(self.tolerance))
            out.write('filetype xyz\n\n')
            out.write('output {}\n\n'.format(self.output))

            for structure in self.structures:
                out.write("structure {}\n".format(structure["file"]))
                out.write("  number {}\n".format(
                    self.numbers[structure["name"]]))
                out.write("  inside box {}\n".format(
                    " ".join(str(i) for i in self.box)))
                out.write("end structure\n\n")


if __name__ == "__main__":
    structures = [{"name": "EMC",
                   "file": "/Users/th/Downloads/test_selenium/EMC.lmp.xyz"}]
    pw = PackmolWrapper("/Users/th/Downloads/test_selenium/", structures,
                        {"EMC": '2'}, [0., 0., 0., 10., 10., 10.])
    pw.make_packmol_input()
    pw.run_packmol()
