# coding: utf-8

from __future__ import unicode_literals, division, print_function
import unittest
import random
from custodian.custodian import Job, ErrorHandler, Custodian, Validator
from custodian.custodian import ValidationError, NonRecoverableError, ReturnCodeError
from custodian.custodian import MaxCorrectionsError, MaxCorrectionsPerJobError, MaxCorrectionsPerHandlerError
import os
import glob
import shutil
import subprocess
import ruamel.yaml as yaml

"""
Created on Jun 1, 2012
"""

__author__ = "Shyue Ping Ong"
__copyright__ = "Copyright 2012, The Materials Project"
__version__ = "0.1"
__maintainer__ = "Shyue Ping Ong"
__email__ = "shyue@mit.edu"
__date__ = "Jun 1, 2012"


class ExitCodeJob(Job):

    def __init__(self, exitcode=0):
        self.exitcode = exitcode

    def setup(self):
        pass

    def run(self):
        return subprocess.Popen('exit {}'.format(self.exitcode), shell=True)

    def postprocess(self):
        pass


class ExampleJob(Job):

    def __init__(self, jobid, params=None):
        if params is None:
            params = {"initial": 0, "total": 0}
        self.jobid = jobid
        self.params = params

    def setup(self):
        self.params["initial"] = 0
        self.params["total"] = 0

    def run(self):
        sequence = [random.uniform(0, 1) for i in range(100)]
        self.params["total"] = self.params["initial"] + sum(sequence)

    def postprocess(self):
        pass

    @property
    def name(self):
        return "ExampleJob{}".format(self.jobid)


class ExampleHandler(ErrorHandler):

    def __init__(self, params):
        self.params = params

    def check(self):
        return self.params["total"] < 50

    def correct(self):
        self.params["initial"] += 1
        return {"errors": "total < 50", "actions": "increment by 1"}


class ExampleHandler1b(ExampleHandler):
    """
    This handler always can apply a correction, but will only apply it twice before raising.
    """
    max_num_corrections = 2
    raise_on_max = True


class ExampleHandler1c(ExampleHandler):
    """
    This handler always can apply a correction, but will only apply it twice and then not anymore.
    """
    max_num_corrections = 2
    raise_on_max = False


class ExampleHandler2(ErrorHandler):
    """
    This handler always result in an error.
    """

    def __init__(self, params):
        self.params = params
        self.has_error = False

    def check(self):
        return True

    def correct(self):
        self.has_error = True
        return {"errors": "Unrecoverable error", "actions": None}


class ExampleHandler2b(ExampleHandler2):
    """
    This handler always result in an error. No runtime error though
    """
    raises_runtime_error = False

    def correct(self):
        self.has_error = True
        return {"errors": "Unrecoverable error", "actions": []}


class ExampleValidator1(Validator):

    def __init__(self):
        pass

    def check(self):
        return False


class ExampleValidator2(Validator):

    def __init__(self):
        pass

    def check(self):
        return True


class CustodianTest(unittest.TestCase):

    def setUp(self):
        self.cwd = os.getcwd()
        os.chdir(os.path.abspath(os.path.dirname(__file__)))

    def test_exitcode_error(self):
        c = Custodian([], [ExitCodeJob(0)])
        c.run()
        c = Custodian([], [ExitCodeJob(1)])
        self.assertRaises(ReturnCodeError, c.run)
        self.assertTrue(c.run_log[-1]["nonzero_return_code"])
        c = Custodian([], [ExitCodeJob(1)],
                      terminate_on_nonzero_returncode=False)
        c.run()

    def test_run(self):
        njobs = 100
        params = {"initial": 0, "total": 0}
        c = Custodian([ExampleHandler(params)],
                      [ExampleJob(i, params) for i in range(njobs)],
                      max_errors=njobs)
        output = c.run()
        self.assertEqual(len(output), njobs)
        d = ExampleHandler(params).as_dict()

    def test_run_interrupted(self):
        njobs = 100
        params = {'initial': 0, 'total': 0}
        c = Custodian([ExampleHandler(params)],
                      [ExampleJob(i, params) for i in range(njobs)],
                      max_errors=njobs)

        self.assertEqual(c.run_interrupted(), njobs)
        self.assertEqual(c.run_interrupted(), njobs)

        total_done = 1
        while total_done < njobs:
            c.jobs[njobs - 1].run()
            if params['total'] > 50:
                self.assertEqual(c.run_interrupted(), njobs - total_done)
                total_done += 1

    def test_unrecoverable(self):
        njobs = 100
        params = {"initial": 0, "total": 0}
        h = ExampleHandler2(params)
        c = Custodian([h],
                      [ExampleJob(i, params) for i in range(njobs)],
                      max_errors=njobs)
        self.assertRaises(NonRecoverableError, c.run)
        self.assertTrue(h.has_error)
        h = ExampleHandler2b(params)
        c = Custodian([h],
                      [ExampleJob(i, params) for i in range(njobs)],
                      max_errors=njobs)
        c.run()
        self.assertTrue(h.has_error)
        self.assertEqual(c.run_log[-1]["handler"], h)

    def test_max_errors(self):
        njobs = 100
        params = {"initial": 0, "total": 0}
        h = ExampleHandler(params)
        c = Custodian([h],
                      [ExampleJob(i, params) for i in range(njobs)],
                      max_errors=1, max_errors_per_job=10)
        self.assertRaises(MaxCorrectionsError, c.run)
        self.assertTrue(c.run_log[-1]["max_errors"])

    def test_max_errors_per_job(self):
        njobs = 100
        params = {"initial": 0, "total": 0}
        h = ExampleHandler(params)
        c = Custodian([h],
                      [ExampleJob(i, params) for i in range(njobs)],
                      max_errors=njobs, max_errors_per_job=1)
        self.assertRaises(MaxCorrectionsPerJobError, c.run)
        self.assertTrue(c.run_log[-1]["max_errors_per_job"])

    def test_max_errors_per_handler_raise(self):
        njobs = 100
        params = {"initial": 0, "total": 0}
        h = ExampleHandler1b(params)
        c = Custodian([h],
                      [ExampleJob(i, params) for i in range(njobs)],
                      max_errors=njobs*10, max_errors_per_job=1000)
        self.assertRaises(MaxCorrectionsPerHandlerError, c.run)
        self.assertEqual(h.n_applied_corrections, 2)
        self.assertEqual(len(c.run_log[-1]["corrections"]), 2)
        self.assertTrue(c.run_log[-1]["max_errors_per_handler"])
        self.assertEqual(c.run_log[-1]["handler"], h)

    def test_max_errors_per_handler_warning(self):
        njobs = 100
        params = {"initial": 0, "total": 0}
        c = Custodian([ExampleHandler1c(params)],
                      [ExampleJob(i, params) for i in range(njobs)],
                      max_errors=njobs*10, max_errors_per_job=1000)
        c.run()
        self.assertTrue(all(len(r["corrections"]) <= 2 for r in c.run_log))

    def test_validators(self):
        njobs = 100
        params = {"initial": 0, "total": 0}
        c = Custodian([ExampleHandler(params)],
                      [ExampleJob(i, params) for i in range(njobs)],
                      [ExampleValidator1()],
                      max_errors=njobs)
        output = c.run()
        self.assertEqual(len(output), njobs)

        njobs = 100
        params = {"initial": 0, "total": 0}
        v = ExampleValidator2()
        c = Custodian([ExampleHandler(params)],
                      [ExampleJob(i, params) for i in range(njobs)],
                      [v],
                      max_errors=njobs)
        self.assertRaises(ValidationError, c.run)
        self.assertEqual(c.run_log[-1]["validator"], v)

    def test_from_spec(self):
        spec = """jobs:
- jb: custodian.vasp.jobs.VaspJob
  params:
    final: False
    suffix: .relax1
- jb: custodian.vasp.jobs.VaspJob
  params:
    final: True
    suffix: .relax2
    settings_override: {"file": "CONTCAR", "action": {"_file_copy": {"dest": "POSCAR"}}}
jobs_common_params:
  $vasp_cmd: ["mpirun", "-machinefile", "$PBS_NODEFILE", "-np", "24", "/opt/vasp/5.4.1/bin/vasp"]
handlers:
- hdlr: custodian.vasp.handlers.VaspErrorHandler
- hdlr: custodian.vasp.handlers.AliasingErrorHandler
- hdlr: custodian.vasp.handlers.MeshSymmetryErrorHandler
validators:
- vldr: custodian.vasp.validators.VasprunXMLValidator
custodian_params:
  $scratch_dir: $TMPDIR"""

        os.environ["TMPDIR"] = "/tmp/random"
        os.environ["PBS_NODEFILE"] = "whatever"
        d = yaml.safe_load(spec)
        c = Custodian.from_spec(d)
        self.assertEqual(c.jobs[0].vasp_cmd[2], "whatever")
        self.assertEqual(c.scratch_dir, "/tmp/random")
        self.assertEqual(len(c.jobs), 2)
        self.assertEqual(len(c.handlers), 3)
        self.assertEqual(len(c.validators), 1)

    def tearDown(self):
        for f in glob.glob("custodian.*.tar.gz"):
            os.remove(f)
        try:
            os.remove("custodian.json")
        except OSError:
            pass  # Ignore if file cannot be found.
        os.chdir(self.cwd)


class CustodianCheckpointTest(unittest.TestCase):

    def setUp(self):
        self.cwd = os.getcwd()
        os.chdir(os.path.join(os.path.dirname(__file__), "..", "..",
                              "test_files", "checkpointing"))
        shutil.copy(os.path.join('backup.tar.gz'),
                    'custodian.chk.3.tar.gz')

    def test_checkpoint_loading(self):
        njobs = 5
        params = {"initial": 0, "total": 0}
        c = Custodian([ExampleHandler(params)],
                      [ExampleJob(i, params) for i in range(njobs)],
                      [ExampleValidator1()],
                      max_errors=100, checkpoint=True)
        self.assertEqual(len(c.run_log), 3)
        self.assertEqual(len(c.run()), 5)

    def tearDown(self):
        os.remove("custodian.json")
        os.chdir(self.cwd)


if __name__ == "__main__":
    unittest.main()
