from strucscan.engine.generalengine import GeneralEngine
from strucscan.core import datatree
from strucscan.utils import PROJECT_PATH, STRUCT_FILE_FORMAT
from strucscan.resources.properties import *

from ase import io

import subprocess
import os


class DummyEngine(GeneralEngine):
    def __init__(self, input_dict):
        """
        dummy engine object
        :param input_dict: (dict) input dictionary. Please follow the examples in strucscan.resources.inputyaml
        """
        GeneralEngine.__init__(self, input_dict)
        self.resultfilename = "final"
        self.final_struct_fname = "final.cfg"
        self.struct_file_format = "cfg"

    def get_name(self):
        """
        :return: (str) engine name that is given by user in input_dict
        """
        return self.name

    def get_dummy_infile(self):
        """
        :return: (str list)
        """
        infile = ["This is a dummy infile.\n"]
        return infile

    def get_dummy_potential_file(self):
        """
        :return: (str list)
        """
        potfile = ["This is a dummy potential file.\n"]
        return potfile

    def make_inputfiles(self, machine_info, jobobject):
        """
        - dummy method that creates / writes engine specific input files
        - for VASP, e.g. this method needs to write INCAR, KPOINTS, POSCAR, POTCAR
        - calls strucscan.engine.generalengine.GeneralEngine.get_absolute_jobpath
        - calls strucscan.engine.generalengine.GeneralEngine.make_machinefile
        - calls strucscan.engine.generalengine.GeneralEngine.write_structure

        :param machine_info: (dict) dictionary of form {"queuename": str, "ncores": int, "nnodes": int}
        :param jobobject: (strucscan.core.jobobject.JobObject object) object that contains information about job
        :return: (str) machine_script_fname
        """
        jobpath = jobobject.get_jobpath()
        atoms = jobobject.basis_ref_atoms
        property = jobobject.property

        if not os.path.exists(jobpath):
            os.makedirs(jobpath)

        # write machinefile
        ncores = int(machine_info["ncores"])
        nnodes = int(machine_info["nnodes"])
        ntotalcores = ncores * nnodes
        nsteps = len(atoms)
        jobname = self.subjobname(self.species, property)
        machinefile, machinefilename = self.make_machinefile(machine_info, jobname, ntotalcores, property, nsteps)

        # write structure
        if isinstance(atoms, list):
            # this is in the case of, e.g. murnaghan calculation, transformation path, ...
            for i, atom in enumerate(atoms):
                self.write_structure(atom, jobpath, structfilename="structure-{:d}.{}".format(i, self.struct_file_format))
            atoms = atoms[0]
        else:
            self.write_structure(atoms, jobpath)

        # write infile
        infile = self.get_dummy_infile()
        with open(jobpath + "/infile.in", "w") as f:
            for line in infile:
                f.write(line)

        # write potential file
        potfile = self.get_dummy_potential_file()
        with open(jobpath + "/potfile.pot", "w") as f:
            for line in potfile:
                f.write(line)

        with open(jobpath + "/" + machinefilename, "w") as f:
            for line in machinefile:
                f.write(line)

        return machinefilename

    def get_absolute_jobpath(self, property, jobobject, structpath=None):
        """
        - dummy method to return absolute path to job directory
        - this is engine specific: 'DATA_TREE_PATH/ENGINE_SIGNATURE/RELATIVE_JOBPATH'
        - while 'DATA_TREE_PATH' is configured in .strucscan and 'RELATIVE_JOBPATH' is generated by
        strucscan.datatree.get_relative_jobpath, 'ENGINE_SIGNATURE' needs to be generated by this method
        - example: 'DATA_TREE_PATH/VASP_5_4_PBE/AlNi/static__L1_2_NiAl'

        :param property: (str) name of property
        :param structpath: (str) absolute path to structure file
        :return: (str) absolute path to job directory
        """
        engine_signature = "DUMMY"
        jobpath = datatree.get_relative_jobpath(self.input_dict["species"], property, jobobject,
                                                structpath=structpath)
        absolute_jobpath = "{}/{}/{}". \
            format(PROJECT_PATH(),
                   engine_signature,
                   jobpath
                   )
        return absolute_jobpath

    def make_machinefile(self, machine_info, jobname, ntotalcores, property, nsteps):
        """
        - dummy method to create / write machine script

        :param machine_info: (dict) machine configuration script taken from 'input_dict'
        :param jobname: (str) name of job, e.g. 'static__fcc__Al'
        :param ntotalcores: (int) number of total cores, that is #cores per node * #of nodes
        :param property: (str) name of property
        :param nsteps: (int) number of single calculation of, e.g. E-V curves, transformation path, ...
        :return: (str list, str) tuple of (list of lines in machine script, machine script file name)
        """
        machine_script, machine_script_fname = self.scheduler.configure_machine_script(machine_info, jobname=jobname)

        modules = []
        if ntotalcores > 1:  # parallel
            # Strucscan loads any prerequesites like module to be load
            # and the final binary call from the machine specific config.yaml.
            config = self.machine_configuration_dict["DUMMY"]["parallel"]
        else:  # serial
            config = self.machine_configuration_dict["DUMMY"]["serial"]
        call = config.split("\n")[-2]

        machine_script.append("echo \"property: %s\" >> start.dat\n" % property)
        machine_script.append("\n")

        if modules != []:
            for module in modules:
                machine_script.append("module load %s\n" % module)
        machine_script.append("\n")

        if nsteps > 1:
            for i in range(nsteps):
                machine_script.append("cp structure-{:d}.cfg structure.cfg\n".format(i))
                machine_script.append(call + "\n")
                machine_script.append("mv final.cfg final-{:d}.cfg\n".format(i))
                machine_script.append("mv log.out log-{:d}.out\n".format(i))
                machine_script.append("rm structure-{:d}.cfg\n".format(i))
                # you can add further lines to machinefile to, e.g. move or gzip files
                machine_script.append("\n")
            machine_script.append("rm structure.cfg\n")
        else:
            machine_script.append(call)
            # you can add further lines to machinefile to, e.g. move or gzip files

        machine_script.append("\n")
        machine_script.append("STOP=`date`\n")
        machine_script.append("echo \"stop: $STOP  $HOSTNAME\" >> end.dat ")
        return machine_script, machine_script_fname

    @staticmethod
    def write_structure(atoms, jobpath, structfilename="structure.cfg"):
        """
        - dummy method to write engine specific structure file

        :param atoms: (ASE atoms object) atoms object with assigned chemical symbols, magnetic moments and scaled positions
        :param jobpath: (str) absolute path to job directory
        :param structfilename: (str) name of structure file that is written to disk. for VASP, structfilename is 'POSCAR'
        :return: 0
        """
        io.write(jobpath + "/" + structfilename, atoms, format="cfg")
        return

    @staticmethod
    def read_final_structure(path, resultfilename="final.cfg"):
        """
        - dummy method to read final structure file

        :param path: (str) absolute path to result directory
        :param resultfilename: (str) name of final structure file. for VASP, structfilename is 'OUTCAR.gz'
        :return: (ASE atoms object) final atoms object
        """
        fname = path
        if resultfilename not in fname:
            fname += "/" + resultfilename
        try:
            final_struct = io.read(fname, format="cfg")
            return final_struct
        except Exception:
            raise FileNotFoundError("{} not found.".format(fname))

    def has_resultfile(self, files):
        try:
            for file in files:
                if ("log" in file) and (".out" in file):
                    return True
            else:
                return False
        except NotADirectoryError:
            raise

    def check_if_finished(self, files):
        """
        - dummy method to check if the calculation in job directory with files with files is finished

        :param files: (str list) str list of all files in directory
        :return: (bool)
        """
        result_filename = self.resultfilename
        for file in files:
            if ("log" in file) and (".out" in file):
                result_filename = file
        cmd = subprocess.Popen("zgrep \"This is a dummy log file.\" %s" % result_filename, shell=True,
                               stderr=subprocess.PIPE, stdout=subprocess.PIPE)
        output, err = cmd.communicate()
        if str(err) != "":
            if str(output)[2:-1] != "":
                return True
        return

