#!/usr/bin/env python

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

from __future__ import division

__author__ = "Shyue Ping Ong"
__version__ = "0.1"
__maintainer__ = "Shyue Ping Ong"
__email__ = "shyuep@gmail.com"
__status__ = "Beta"
__date__ = "5/20/13"


from custodian.custodian import ErrorHandler, backup
from pymatgen.serializers.json_coders import MSONable
from pymatgen.io.nwchemio import NwOutput, NwInput
from custodian.ansible.intepreter import Modder


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

    def __init__(self, output_filename="mol.nwout"):
        self.output_filename = output_filename

    def check(self):
        # Checks output file for errors.
        out = NwOutput(self.output_filename)
        self.errors = []
        self.input_file = out.job_info['input']
        if out.data[-1]["has_error"]:
            self.errors.extend(out.data[-1]["errors"])
        self.errors = list(set(self.errors))
        self.ntasks = len(out.data)
        return len(self.errors) > 0

    def _mod_input(self, search_string_func, mod_string_func):
        with open(self.input_file) as f:
            lines = []
            for l in f:
                if search_string_func(l):
                    lines.append(mod_string_func(l))
                else:
                    lines.append(l)

        with open(self.input_file, "w") as fout:
            fout.write("".join(lines))

    def correct(self):
        backup("*.nw*")
        actions = []
        nwi = NwInput.from_file(self.input_file)
        for e in self.errors:
            if e == "autoz error":
                action = {"_set": {"geometry_options": ["units",
                                                        "angstroms",
                                                        "noautoz"]}}
                actions.append(action)
            elif e == "Bad convergence":
                t = nwi.tasks[self.ntasks - 1]
                if "cgmin" in t.theory_directives:
                    nwi.tasks.pop(self.ntasks - 1)
                else:
                    t.theory_directives["cgmin"] = ""
                for t in nwi.tasks:
                    if t.operation.startswith("freq"):
                        #You cannot calculate hessian with cgmin.
                        t.theory_directives["nocgmin"] = ""
                action = {"_set": {"tasks": [t.to_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()
        for action in actions:
            nwi = m.modify_object(action, nwi)
        nwi.write_file(self.input_file)
        return {"errors": self.errors, "actions": actions}

    @property
    def is_monitor(self):
        return False

    def __str__(self):
        return "NwchemErrorHandler"

    @property
    def to_dict(self):
        return {"@module": self.__class__.__module__,
                "@class": self.__class__.__name__,
                "output_filename": self.output_filename}

    @classmethod
    def from_dict(cls, d):
        return cls(d["output_filename"])

