# coding: utf-8
# Copyright (c) Pymatgen Development Team.
# Distributed under the terms of the MIT License.


import os
import unittest

from pymatgen.io.qchem.sets import *
from pymatgen.util.testing import PymatgenTest

__author__ = "Samuel Blau, Brandon Wood, Shyam Dwaraknath"
__copyright__ = "Copyright 2018, The Materials Project"
__version__ = "0.1"


test_dir = os.path.join(PymatgenTest.TEST_FILES_DIR, "molecules")


class QChemDictSetTest(PymatgenTest):
    def test_init(self):
        test_molecule = QCInput.from_file(os.path.join(test_dir, "new_qchem_files/pcm.qin")).molecule
        test_DictSet = QChemDictSet(
            molecule=test_molecule,
            job_type="opt",
            basis_set="6-31G*",
            scf_algorithm="diis",
        )
        self.assertEqual(
            test_DictSet.rem,
            {
                "job_type": "opt",
                "gen_scfman": "true",
                "basis": "6-31g*",
                "max_scf_cycles": 200,
                "method": "wb97xv",
                "scf_algorithm": "diis",
                "xc_grid": "3",
                "geom_opt_max_cycles": 200,
                "symmetry": "false",
                "sym_ignore": "true",
                "resp_charges": "true",
            },
        )
        self.assertEqual(test_DictSet.pcm, {})
        self.assertEqual(test_DictSet.solvent, {})
        self.assertEqual(test_DictSet.smx, {})
        self.assertEqual(test_DictSet.molecule, test_molecule)

    def test_full_init(self):
        test_molecule = QCInput.from_file(os.path.join(test_dir, "new_qchem_files/pcm.qin")).molecule

        test_DictSet = QChemDictSet(
            molecule=test_molecule,
            job_type="opt",
            basis_set="6-31g*",
            scf_algorithm="diis",
            dft_rung=1,
            pcm_dielectric=10.0,
            max_scf_cycles=35,
        )
        self.assertEqual(
            test_DictSet.rem,
            {
                "job_type": "opt",
                "gen_scfman": "true",
                "basis": "6-31g*",
                "max_scf_cycles": 35,
                "method": "b3lyp",
                "geom_opt_max_cycles": 200,
                "scf_algorithm": "diis",
                "xc_grid": "3",
                "solvent_method": "pcm",
                "symmetry": "false",
                "sym_ignore": "true",
                "resp_charges": "true",
            },
        )
        self.assertEqual(
            test_DictSet.pcm,
            {
                "heavypoints": "194",
                "hpoints": "194",
                "radii": "uff",
                "theory": "cpcm",
                "vdwscale": "1.1",
            },
        )
        self.assertEqual(test_DictSet.solvent, {"dielectric": 10.0})
        self.assertEqual(test_DictSet.molecule, test_molecule)

        test_DictSet = QChemDictSet(
            molecule=test_molecule,
            job_type="opt",
            basis_set="6-31g*",
            scf_algorithm="diis",
            dft_rung=1,
            smd_solvent="water",
            max_scf_cycles=35,
        )
        self.assertEqual(
            test_DictSet.rem,
            {
                "job_type": "opt",
                "gen_scfman": "true",
                "basis": "6-31g*",
                "max_scf_cycles": 35,
                "method": "b3lyp",
                "geom_opt_max_cycles": 200,
                "scf_algorithm": "diis",
                "xc_grid": "3",
                "solvent_method": "smd",
                "ideriv": "1",
                "symmetry": "false",
                "sym_ignore": "true",
                "resp_charges": "true",
            },
        )
        self.assertEqual(test_DictSet.smx, {"solvent": "water"})

    def test_overwrite_input(self):
        test_molecule = QCInput.from_file(os.path.join(test_dir, "new_qchem_files/pcm.qin")).molecule
        overwrite_inputs = {
            "rem": {
                "method": "b3lyp",
                "basis": "6-31g*",
                "thresh": 10,
                "xc_grid": "000150000302",
            }
        }
        test_OptSet = OptSet(molecule=test_molecule, overwrite_inputs=overwrite_inputs)
        act_rem = {
            "job_type": "opt",
            "gen_scfman": "true",
            "basis": "6-31g*",
            "max_scf_cycles": 200,
            "method": "b3lyp",
            "scf_algorithm": "diis",
            "xc_grid": "000150000302",
            "geom_opt_max_cycles": 200,
            "thresh": 10,
            "symmetry": "false",
            "sym_ignore": "true",
            "resp_charges": "true",
        }
        self.assertDictEqual(act_rem, test_OptSet.rem)

    def test_double_solvation(self):
        test_molecule = QCInput.from_file(os.path.join(test_dir, "new_qchem_files/pcm.qin")).molecule
        raised_error = False
        dict_set = None
        try:
            dict_set = QChemDictSet(
                molecule=test_molecule,
                job_type="opt",
                basis_set="6-31g*",
                scf_algorithm="diis",
                dft_rung=1,
                pcm_dielectric=10.0,
                smd_solvent="water",
                max_scf_cycles=35,
            )
        except ValueError:
            raised_error = True

        self.assertTrue(raised_error)
        self.assertEqual(dict_set, None)

    def test_pcm_write(self):
        test_molecule = QCInput.from_file(os.path.join(test_dir, "new_qchem_files/pcm.qin")).molecule
        dict_set = QChemDictSet(
            molecule=test_molecule,
            job_type="opt",
            basis_set="6-31g*",
            scf_algorithm="diis",
            dft_rung=5,
            pcm_dielectric=10.0,
            max_scf_cycles=35,
        )
        dict_set.write("mol.qin")
        test_dict = QCInput.from_file("mol.qin").as_dict()
        rem = {
            "job_type": "opt",
            "basis": "6-31G*",
            "max_scf_cycles": "35",
            "method": "wb97mv",
            "geom_opt_max_cycles": "200",
            "gen_scfman": "true",
            "scf_algorithm": "diis",
            "xc_grid": "3",
            "solvent_method": "pcm",
            "symmetry": "false",