"""
This module contains some wrappers for using DFTB+ to perform calculations.

https://dftbplus.org/

Input files are defined using: https://github.com/dftbplus/hsd-python

Parameters can be found at: https://dftb.org/parameters
"""
from BigDFT.Calculators import Runner
from futile.Utils import write as safe_print


class DFTBLogfile(dict):
    """
    This class stores the values of the result.tag generated by DFTB+.

    Attributes:
        sys (BigDFT.Systems.System): a handle to the original system used
            in the calculation.
        energy (float): the energy of the system
    """
    def __init__(self, sys, logname):
        self.sys = sys
        self._process_values(logname)

    @property
    def energy(self):
        """
        The total energy of the system.
        """
        return self["total_energy"]

    def _process_values(self, logname):
        from operator import mul
        from functools import reduce
        from numpy import array

        with open(logname) as ifile:
            try:
                while(True):
                    line = next(ifile)
                    split = line.split(":")
                    name = split[0].strip()
                    dtype = split[1]
                    if dtype == "real":
                        trans = lambda x: float(x)
                    else:
                        trans = None
                    ndim = int(split[2])
                    if ndim == 0:
                        self[name] = trans(next(ifile))
                    else:
                        dims = [int(x) for x in split[3].split(",")]
                        total_data = reduce(mul, dims)
                        raw_data = []
                        while(True):
                            raw_data.extend([trans(x) for x in
                                             next(ifile).split()])
                            if len(raw_data) == total_data:
                                break
                        rdim = tuple(reversed(dims))
                        self[name] = array(raw_data).reshape(rdim)
            except StopIteration:
                pass


class DFTBCalculator(Runner):
    """
    A calculator that drives DFTB+ calculations through the command line.
    """
    import os

    def __init__(self, omp=os.environ.get('OMP_NUM_THREADS', '1'),
                 mpi_run=os.environ.get('BIGDFT_MPIRUN', ''),
                 dry_run=False, skip=False, verbose=True):
        from os.path import join

        # Use the initialization from the Runner class (so all options inside
        # __global_options)
        Runner.__init__(self, omp=str(omp), dry_run=dry_run, skip=skip,
                        mpi_run=mpi_run, verbose=verbose)

        self.command = self._global_options['mpi_run'] + " dftb+"

        if verbose:
            safe_print(
                'Initialize a Calculator with OMP_NUM_THREADS=%s '
                'and command %s' %
                (self._global_options['omp'], self.command))

    def pre_processing(self):
        """
        Process local run dictionary to create the input directory and identify
        the command to be passed

        Returns:
            :py:class:`dict`: dictionary containing the command to be passed to
            :meth:`process_run`
        """
        from hsd import dump
        from os.path import join

        # Update the run directory to include the name
        self._ensure_run_directory()
        run_dir = join(self.run_dir, self.run_options.get("name", "dftb"))
        self.run_options["run_dir"] = run_dir
        self._ensure_run_directory()

        # Setup the files
        sys = self.run_options["sys"]
        inp = self.run_options.get("input", {})

        # Write out the geometry file to our dictionary
        inp["Geometry"] = self._get_geometry(sys)

        # Ensure that the output options are correct
        if "Options" not in inp:
            inp["Options"] = {}
        inp["Options"]["WriteResultsTag"] = True

        # Write
        dump(inp, join(self.run_dir, "dftb_in.hsd"))

        return {'command': self._get_command()}

    def process_run(self, command):
        """
        Run the DFTB+ executable.
        """
        from os import environ, system
        from os.path import join

        # Set the number of omp threads only if the variable is not present
        # in the environment
        if 'OMP_NUM_THREADS' not in environ:
            environ['OMP_NUM_THREADS'] = self.run_options['omp']

        if self.run_options['verbose']:
            if self.run_dir != '.':
                safe_print('Run directory', self.run_dir)
            safe_print('Executing command: ', command)

        # Run the command
        system("cd " + self.run_dir + "; " + command + " > " + 
               self.run_options.get("name", "dftbrun") + ".log")

        return {'logname': join(self.run_dir, "results.tag")}

    def post_processing(self, logname, command):
        """
        Post processing the calculation.

        Returns:
            (BigDFT.Interop.DFTBInterop.DFTBLogfile): a representation of the
            detailed output.
        """
        try:
            return DFTBLogfile(sys=self.run_options["sys"], logname=logname)
        except IOError:
            raise ValueError("The logfile (", logname, ") does not exist.")

    def _get_geometry(self, sys):
        retval = {}

        # List of types
        type_names = []
        for frag in sys.values():
            for at in frag:
                type_names.append(at.sym)
        type_names = list(set(type_names))
        retval["TypeNames"] = type_names

        # A lookup table for types
        type_lookup = {}
        for i, sym in enumerate(type_names):
            type_lookup[sym] = i + 1

        # Periodic
        if sys.cell == None or sys.cell.get_boundary_condition() == "free":
            retval["Periodic"] = False
        else:
            retval["Periodic"] = True
            retval["LatticeVectors"] = sys.cell.cell
            retval["LatticeVectors.attrib"] = "Bohr"

        # Write the coordinates
        retval["TypesAndCoordinates.attrib"] = "Bohr"
        retval["TypesAndCoordinates"] = []
        for frag in sys.values():
            for at in frag:
                retval["TypesAndCoordinates"] += [[type_lookup[at.sym]] + 
                                                   at.get_position()]
  
        return retval

    def _get_command(self):
        if self._check_skip():
            return '''echo "skip"'''
        else:
            return self.command

    def _check_skip(self):
        if not self.run_options["skip"]:
            return False
        try:
            DFTBLogfile(sys=self.run_options["sys"], 
                        logname=join(self.run_dir, "results.tag"))
            return True
        except Exception as e:
            print(e)
            return False

    def _ensure_run_directory(self):
        from futile.Utils import ensure_dir
        run_dir = self.run_options.get('run_dir', '.')
        # Create the run_dir if not exist
        if ensure_dir(run_dir) and self.run_options['verbose']:
            safe_print("Create the sub-directory '%s'" % run_dir)

        self.run_dir = run_dir


def _example():
    """Example of using DFTB+ interoperability"""
    from BigDFT.IO import XYZReader
    from BigDFT.Systems import System
    from BigDFT.Fragments import Fragment
    from os.path import join
    from os import getcwd
    from copy import deepcopy

    # Create a system.
    reader = XYZReader("H")
    fsys = System()
    fsys["FRA:1"] = Fragment(xyzfile=reader)
    fsys["FRA:2"] = deepcopy(fsys["FRA:1"])
    fsys["FRA:2"].translate([0.7414, 0, 0])

    # Create an input file
    # https://github.com/dftbplus/hsd-python
    inp = {}
    inp["Driver"] = {}
    inp["Hamiltonian"] = {"DFTB": {"MaxAngularMomentum": {}, 
                                   "SlaterKosterFiles": {"Type2FileNames": {}}}}

    inp["Hamiltonian"]["DFTB"]["MaxAngularMomentum"] = {"H": "s"}
    skf = {"Suffix": ".skf", "Separator": "-", 
           "Prefix": join(getcwd(), "input", "3ob-3-1/")}
    inp["Hamiltonian"]["DFTB"]["SlaterKosterFiles"]["Type2FileNames"] = skf

    # Create a Calculator and Run
    calc = DFTBCalculator(mpi_run="mpirun -np 1")
    log = calc.run(sys=fsys, input=inp, name="H2", run_dir="scratch")

    # The full set of data from the `results.tag` file are available.
    log["mermin_energy"]


if __name__ == "__main__":
    _example()
