from ase.calculators.calculator import (
    FileIOCalculator,
    CalculationFailed,
    equal,
    compare_atoms,
    PropertyNotPresent,
)
from xespresso.xio import read_espresso_asei
from xespresso.xespresso import Espresso
import os
import logging

logger = logging.getLogger(__name__)


class EspressoNscf(FileIOCalculator):

    package = "pw"

    def __init__(
        self,
        scf_directory,
        prefix,
        parallel="",
        queue=None,
        debug=False,
        kpts=(10, 10, 10),
        **kwargs
    ):
        print("{0:=^60}".format("nscf"))
        self.scf_directory = scf_directory
        self.prefix = prefix
        self.parallel = parallel
        self.queue = queue
        if debug:
            logger.setLevel(debug)
        self.load_scf()
        self.parameters["kpts"] = kpts
        self.parameters.update(kwargs)
        self.parameters = Espresso.check_input(self.parameters, self.prefix)
        #
        self.directory = os.path.join(self.scf_directory, "nscf")
        self.label = os.path.join(self.directory, self.prefix)
        self.asei = os.path.join(self.directory, "%s.nscf_asei" % self.prefix)
        # create working directory
        self.set_label(self.directory, self.prefix)

    def set_label(self, label, prefix):
        """
        set directory and prefix from label
        """
        self.directory = label
        if not prefix:
            self.prefix = os.path.split(label)[1]
        else:
            self.prefix = prefix
        if not os.path.exists(self.directory):
            os.makedirs(self.directory)
        self.label = os.path.join(self.directory, self.prefix)
        self.pwi = os.path.join(self.directory, "%s.pwi" % self.prefix)
        self.pwo = os.path.join(self.directory, "%s.pwo" % self.prefix)
        self.save_directory = os.path.join(self.directory, "%s.save" % self.prefix)
        logger.debug("Directory: %s" % (self.directory))
        logger.debug("Prefix: %s" % (self.prefix))

    def load_scf(self):
        """load parameters and results from privous scf calculation"""
        import copy

        asei = os.path.join(self.scf_directory, "%s.asei" % self.prefix)
        self.atoms, self.scf_parameters = read_espresso_asei(asei, "PW")
        self.parameters = copy.deepcopy(self.scf_parameters)
        self.parameters["input_data"]["CONTROL"]["calculation"] = "nscf"
        self.parameters["input_data"]["CONTROL"]["verbosity"] = "high"
        self.parameters["input_data"]["CONTROL"]["outdir"] = "../"
        self.parameters["input_data"]["SYSTEM"]["occupations"] = "tetrahedra"
        self.parameters["input_data"]["SYSTEM"].pop("degauss", None)

    def write_input(self, atoms, properties=None, system_changes=None):
        from xespresso.xio import write_espresso_asei, write_espresso_in
        from xespresso.scheduler import set_queue

        FileIOCalculator.write_input(self, atoms, properties, system_changes)
        set_queue(self, package=self.package, parallel=self.parallel, queue=self.queue)
        write_espresso_in(self.label + ".pwi", atoms, **self.parameters)
        logger.debug("write asei file: {}".format(self.asei))
        logger.debug("  with parameters: {}".format(self.parameters))
        write_espresso_asei(self.asei, self.state_info, self.parameters)
        logger.debug("Write input successfully")

    @property
    def state_info(self):
        from xespresso.utils import get_hash

        chargeFile = os.path.join(
            self.scf_directory, "%s.save/charge-density.dat" % self.prefix
        )
        state_info = get_hash(chargeFile)
        return state_info

    def check_state(
        self,
    ):
        # read information of the charge-density file
        self.state_parameters = self.parameters
        output, meg = self.read_convergence_post("pw")
        logger.debug("Check state: {}".format(meg))
        if output:
            if os.path.isfile(self.asei):
                system_changes = self.check_state_post(self.asei, package="PW")
                if not system_changes:
                    logger.debug("Use previous results!")
                    return False
            else:
                logger.debug("No asei output. Start a new calculation.")

        else:
            logger.debug("No pw output. Start a new calculation.")

        return True

    def check_state_post(self, asei, package):
        logger.debug("asei file: {}".format(asei))
        old_state_info, old_state_parameters = read_espresso_asei(asei, package)
        if not self.state_info == old_state_info:
            logger.debug("File in save changed")
            return True
        elif not self.state_parameters == old_state_parameters:
            logger.debug(self.state_parameters)
            logger.debug(old_state_parameters)
            logger.debug("Parameters changed")
            return True
        else:
            return False

    def read_convergence_post(self, package="pw"):
        """
        Read the status of the calculation.
        {
        '0': 'Done',
        }
        """
        output = self.label + ".%so" % package
        logger.debug("read output: {}".format(output))
        if not os.path.exists(output):
            # print('%s not exists'%output)
            return False, "No pwo output file"
        with open(output, "r") as f:
            lines = f.readlines()
            if len(lines) == 0:
                return False, "pwo file has nothing"
            nlines = len(lines)
            n = min([100, nlines])
            for line in lines[-n:-1]:
                if line.rfind("JOB DONE.") > -1:
                    logger.debug("JOB DONE.")
                    return True, line
        return False, line

    def run(self):
        if self.check_state():
            self.calculate()
