#!/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
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):
        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:
                    return {"errors": self.errors, "actions": None}
                t.theory_directives["cgmin"] = ""
                if t.operation == "optimize":
                    for t in nwi.tasks:
                        if t.operation.startswith("freq"):
                            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"])


@@ -312,16 +312,58 @@
 20/13%22%0A%0A
+import glob%0Aimport logging%0Aimport tarfile%0A
 %0Afrom cu
@@ -1626,32 +1626,49 @@
  correct(self):%0A
+        backup()%0A
         actions 
@@ -2191,56 +2191,65 @@
     
-return %7B%22errors%22: self.errors, %22actions%22: None%7D%0A
+nwi.tasks.pop(self.ntasks - 1)%0A                else:%0A    
     
@@ -2302,32 +2302,36 @@
                 
+    
 if t.operation =
@@ -2360,24 +2360,28 @@
             
+ 
+   
 for t in nwi
@@ -2388,16 +2388,20 @@
 .tasks:%0A
+    
         
@@ -2447,24 +2447,102 @@
 th(%22freq%22):%0A
+                                #You cannot calculate hessian with cgmin.%0A    
             
@@ -3500,8 +3500,360 @@
 ame%22%5D)%0A%0A
+%0Adef backup():%0A    error_num = max(%5B0%5D + %5Bint(f.split(%22.%22)%5B1%5D)%0A                           for f in glob.glob(%22error.*.tar.gz%22)%5D)%0A    filename = %22error.%7B%7D.tar.gz%22.format(error_num + 1)%0A    logging.info(%22Backing up run to %7B%7D.%22.format(filename))%0A    tar = tarfile.open(filename, %22w:gz%22)%0A    for f in glob.glob(%22*.nw*%22):%0A        tar.add(f)%0A    tar.close()%0A

Add backup for nwchem handler. Improve handling of convergence errors. Now the run will skip past tasks which cannot be converged with cgmin.