"""
This module implements error handlers for Nwchem runs. Currently tested only
for B3LYP DFT jobs.
"""

import os

from pymatgen.io.nwchem import NwInput, NwOutput

from custodian.ansible.interpreter import Modder
from custodian.custodian import ErrorHandler
from custodian.utils import backup


class NwchemErrorHandler(ErrorHandler):
    """
    Error handler for Nwchem Jobs. Currently tested only for B3LYP DFT jobs
    generated by pymatgen.
    """

    def __init__(self, output_filename="mol.nwout") -> None:
        """Initialize with an output file name.

        Args:
            output_filename (str): This is the file where the stdout for nwchem
                is being redirected. The error messages that are checked are
                present in the stdout. Defaults to "mol.nwout", which is the
                default redirect used by :class:`custodian.nwchem.jobs
                .NwchemJob`.
        """
        self.output_filename = output_filename

    def check(self, directory="./"):
        """Check for errors."""
        out = NwOutput(os.path.join(directory, self.output_filename))
        self.errors = []
        self.input_file = out.job_info["input"]
        if out.data[-1]["has_error"]:
            self.errors += out.data[-1]["errors"]
        self.errors = list(set(self.errors))
        self.ntasks = len(out.data)
        return len(self.errors) > 0

    def correct(self, directory="./"):
        """Correct errors."""
        backup("*.nw*", directory=directory)
        actions = []
        nwi = NwInput.from_file(os.path.join(directory, self.input_file))
        for err in self.errors:
            if err == "autoz error":
                action = {"_set": {"geometry_options": ["units", "angstroms", "noautoz"]}}
                actions.append(action)
            elif err == "Bad convergence":
                task = nwi.tasks[self.ntasks - 1]
                if "cgmin" in task.theory_directives:
                    nwi.tasks.pop(self.ntasks - 1)
                else:
                    task.theory_directives["cgmin"] = ""
                for task in nwi.tasks:
                    if task.operation.startswith("freq"):
                        # You cannot calculate hessian with cgmin.
                        task.theory_directives["nocgmin"] = ""
                action = {"_set": {"tasks": [t.as_dict() for t in nwi.tasks]}}
                actions.append(action)
            else:
                # For unimplemented errors, this should just cause the job to
                # die.
                return {"errors": self.errors, "actions": None}

        m = Modder(directory=directory)
        for action in actions:
            nwi = m.modify_object(action, nwi)
        nwi.write_file(self.input_file)
        return {"errors": self.errors, "actions": actions}

    def __str__(self) -> str:
        return "NwchemErrorHandler"
