import os
import unittest

from common import (
    j_loader,
)

from deepmd.utils.compat import (
    convert_input_v0_v1,
    convert_input_v1_v2,
)


class TestConvertInput(unittest.TestCase):
    def test_convert_smth(self):
        jdata0 = j_loader(os.path.join("compat_inputs", "water_se_a_v0.json"))
        jdata1 = j_loader(os.path.join("compat_inputs", "water_se_a_v1.json"))
        jdata = convert_input_v0_v1(jdata0, warning=False, dump=None)
        self.assertEqual(jdata, jdata1)

    def test_convert_nonsmth(self):
        jdata0 = j_loader(os.path.join("compat_inputs", "water_v0.json"))
        jdata1 = j_loader(os.path.join("compat_inputs", "water_v1.json"))
        jdata = convert_input_v0_v1(jdata0, warning=False, dump=None)
        self.assertEqual(jdata, jdata1)

    def test_convert_v1_v2(self):
        jdata0 = j_loader(os.path.join("compat_inputs", "water_v1.json"))
        jdata1 = j_loader(os.path.join("compat_inputs", "water_v2.json"))
        jdata = convert_input_v1_v2(jdata0, warning=False, dump=None)
        self.assertDictAlmostEqual(jdata, jdata1)

    def assertDictAlmostEqual(self, d1, d2, msg=None, places=7):
        self.assertEqual(d1.keys(), d2.keys())
        for kk, vv in d1.items():
            if isinstance(vv, dict):
                self.assertDictAlmostEqual(d1[kk], d2[kk], msg=msg)
            else:
                self.assertAlmostEqual(d1[kk], d2[kk], places=places, msg=msg)

    def test_json_yaml_equal(self):
        inputs = ("water_v1", "water_se_a_v1")

        for i in inputs:
            jdata = j_loader(os.path.join("yaml_inputs", f"{i}.json"))
            ydata = j_loader(os.path.join("yaml_inputs", f"{i}.yaml"))
            self.assertEqual(jdata, ydata)

        with self.assertRaises(TypeError):
            j_loader("path_with_wrong.extension")
