#! /usr/bin/env python

import unittest

from tbplas.builder.orbital_map import GeneralOrbitalMap, RegularOrbitalMap
from tbplas.base.utils import TestHelper


class MyTest(unittest.TestCase):

    def test_general(self):
        th = TestHelper(self)
        orb_map = GeneralOrbitalMap()

        # Initialization
        self.assertEqual(orb_map.counter, 0)

        # Add orbital
        def _test():
            orb_map.add((1, 1, 0))
        th.test_raise(_test, AssertionError, "Length of")

        def _test():
            orb_map.add((1, 1, 0, -1))
        th.test_raise(_test, AssertionError, "out of range")

        orb_map.add((1, 1, 0, 0))
        orb_map.add((0, 1, 1, 0))
        self.assertEqual(orb_map.counter, 2)

        # pc2sc
        def _test():
            orb_map.pc2sc((1, 1, 0))
        th.test_raise(_test, AssertionError, "Length of")

        def _test():
            orb_map.pc2sc((1, 1, 0, -1))
        th.test_raise(_test, AssertionError, "out of range")

        def _test():
            orb_map.pc2sc((1, 1, 0, 1))
        th.test_raise(_test, KeyError, r"^\(\d+,")

        self.assertEqual(orb_map.pc2sc((1, 1, 0, 0)), 0)
        self.assertEqual(orb_map.pc2sc((0, 1, 1, 0)), 1)

        # sc2pc
        def _test():
            orb_map.sc2pc(-1)
        th.test_raise(_test, AssertionError, r"should >= 0$")

        def _test():
            orb_map.sc2pc(3)
        th.test_raise(_test, IndexError, "list index out of range")

        self.assertTupleEqual(orb_map.sc2pc(0), (1, 1, 0, 0))
        self.assertTupleEqual(orb_map.sc2pc(1), (0, 1, 1, 0))

    def test_regular(self):
        th = TestHelper(self)

        # Initialization
        def _test():
            RegularOrbitalMap((-1, 1), 0)
        th.test_raise(_test, AssertionError, "Length of")

        def _test():
            RegularOrbitalMap((-1, 1, 1), 0)
        th.test_raise(_test, AssertionError, r"^Dim\[0\]")

        def _test():
            RegularOrbitalMap((1, -1, 1), 0)
        th.test_raise(_test, AssertionError, r"^Dim\[1\]")

        def _test():
            RegularOrbitalMap((1, 1, -1), 0)
        th.test_raise(_test, AssertionError, r"^Dim\[2\]")

        def _test():
            RegularOrbitalMap((1, 1, 1), 0)
        th.test_raise(_test, AssertionError, "num_orb_pc")

        # pc2sc
        orb_map = RegularOrbitalMap((3, 3, 1), 2)

        def _test():
            orb_map.pc2sc(-1, 1, 1, 0)
        th.test_raise(_test, AssertionError, r"out of range$")

        def _test():
            orb_map.pc2sc(1, -1, 1, 0)
        th.test_raise(_test, AssertionError, r"out of range$")

        def _test():
            orb_map.pc2sc(1, 1, -1, 0)
        th.test_raise(_test, AssertionError, r"out of range$")

        def _test():
            orb_map.pc2sc(1, 1, 1, -1)
        th.test_raise(_test, AssertionError, r"out of range$")

        self.assertEqual(orb_map.pc2sc(1, 1, 0, 0), 8)
        self.assertEqual(orb_map.pc2sc(1, 1, 0, 1), 9)

        # sc2pc
        def _test():
            orb_map.sc2pc(-1)
        th.test_raise(_test, AssertionError, r"should >= 0$")

        self.assertTupleEqual(orb_map.sc2pc(8), (1, 1, 0, 0))
        self.assertTupleEqual(orb_map.sc2pc(9), (1, 1, 0, 1))

    def test_vac_dummy(self):
        """Unittest of RegularOrbitalMapVac will be done with SuperCell."""
        pass


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