#! /usr/bin/env python
import unittest

import numpy as np

from tbplas.cython import super as core
from tbplas import (gen_lattice_vectors, PrimitiveCell, SuperCell,
                    SCInterHopping, TestHelper)


def make_cell():
    """Make 2-band model of graphene."""
    vectors = gen_lattice_vectors(a=2.46, b=2.46, gamma=60)
    cell = PrimitiveCell(vectors)
    cell.add_orbital((0.0, 0.0), 0.0)
    cell.add_orbital((1. / 3, 1. / 3), 0.0)
    cell.add_hopping((0, 0), 0, 1, -2.7)
    cell.add_hopping((1, 0), 1, 0, -2.7)
    cell.add_hopping((0, 1), 1, 0, -2.7)
    return cell


def make_cell_orb():
    """Make 2-band model of graphene without hopping terms."""
    vectors = gen_lattice_vectors(a=2.46, b=2.46, gamma=60)
    cell = PrimitiveCell(vectors)
    cell.add_orbital((0.0, 0.0), 0.0, label="C_pz")
    cell.add_orbital((1. / 3, 1. / 3), 0.0, label="C_pz")
    return cell


def make_cell_empty():
    """Make an empty cell."""
    vectors = gen_lattice_vectors(a=2.46, b=2.46, gamma=60)
    cell = PrimitiveCell(vectors)
    return cell


class MyTest(unittest.TestCase):

    def test_init(self):
        """Test SuperCell.__init__."""
        th = TestHelper(self)

        def _test():
            _cell = make_cell_empty()
            SuperCell(_cell, dim=(3, 3, 1))
        th.test_raise(_test, AssertionError, "Num_orb_pc should > 0")

        def _test():
            _cell = make_cell_orb()
            SuperCell(_cell, dim=(3, 3, 1))
        th.test_raise(_test, AssertionError, "Num_hop_pc should > 0")

        for dim in ((-1, 1, 1), (2, -1, 1), (2, 1, -1)):
            def _test():
                SuperCell(make_cell(), dim=dim)
            th.test_raise(_test, AssertionError, "dim")

        sc = SuperCell(make_cell(), dim=(3, 3, 1))
        self.assertTupleEqual(tuple(sc._dim), (3, 3, 1))
        self.assertTupleEqual(tuple(sc._pbc), (False, False, False))
        self.assertEqual(sc._vacancy_set, set())
        self.assertEqual(sc._cached_data.vac_id_sc.size, 0)
        self.assertEqual(sc._cached_data.orb_id_pc.shape, (18, 4))

    def test_vacancies_illegal(self):
        """
        Test if error handling of vacancy operations works as expected.

        :return: None.
        """
        th = TestHelper(self)

        # Feed a list of vacancies with an illegal one.
        # Errors should be raised
        sc = SuperCell(make_cell(), dim=(3, 3, 1))
        sc.lock()

        def _test():
            sc.add_vacancies([0, 1, 2])
        th.test_raise(_test, AssertionError, "Modifying a locked object")

        def _test():
            sc.set_vacancies([0, 1, 2])
        th.test_raise(_test, AssertionError, "Modifying a locked object")

        def _test():
            sc.unlock()
            sc.add_vacancies(np.array([[1, 2]]))
        th.test_raise(_test, AssertionError, "vacancies should be a vector")

        def _test():
            sc.unlock()
            sc.set_vacancies(np.array([[1, 2]]))
        th.test_raise(_test, AssertionError, "vacancies should be a vector")

        def _test():
            sc.unlock()
            sc.add_vacancies([120])
        th.test_raise(_test, AssertionError, r"^vac_id \d+ out of range")

        def _test():
            sc.unlock()
            sc.set_vacancies([120])
        th.test_raise(_test, AssertionError, r"^vac_id \d+ out of range")

    def test_vacancies(self):
        """
        Test if set_vacancies works as expected for legal input.

        :return: None.
        """
        th = TestHelper(self)
        vacancies = [(0, 0, 0, 1), (0, 1, 0, 0), (1, 1, 0, 1)]
        vac_id_sc = np.array([1, 2, 9])
        orb_id_pc = [(i_a, i_b, 0, i_o)
                     for i_a in range(3)
                     for i_b in range(3)
                     for i_o in range(2)]
        for vac in vacancies:
            orb_id_pc.remove(vac)
        orb_id_pc = np.array(orb_id_pc)

        # First we feed some effective input to set_vacancies.
        sc = SuperCell(make_cell(), dim=(3, 3, 1))
        sc.set_vacancies(vac_id_sc)
        sc._sync_array()
        self.assertSetEqual(sc._vacancy_set, set(vac_id_sc))
        th.test_equal_array(sc._cached_data.vac_id_sc, vac_id_sc)
        th.test_equal_array(sc._cached_data.orb_id_pc, orb_id_pc)

        # Then we feed a blank to set_vacancies.
        sc = SuperCell(make_cell(), dim=(3, 3, 1))
        sc.set_vacancies([])
        sc._sync_array()
        self.assertSetEqual(sc._vacancy_set, set())
        self.assertEqual(sc._cached_data.vac_id_sc.size, 0)
        self.assertEqual(sc._cached_data.orb_id_pc.shape, (18, 4))

        # Finally, we feed a list of vacancies with redundant items.
        # These redundant items should be detected and removed.
        sc = SuperCell(make_cell(), dim=(3, 3, 1))
        vacancies = [1, 2, 1]
        sc.set_vacancies(vacancies)
        sc._sync_array()
        reduced_vacancies = {1, 2}
        self.assertSetEqual(sc._vacancy_set, reduced_vacancies)
        self.assertEqual(sc._cached_data.vac_id_sc.size, 2)
        self.assertEqual(sc._cached_data.orb_id_pc.shape, (16, 4))

    def test_add_hopping(self):
        """Test add_hopping method of SuperCell."""
        sc = SuperCell(make_cell(), dim=(3, 3, 1))
        th = TestHelper(self)

        def _test():
            sc.lock()
            sc.add_hopping((0, 0, 0), 1, 0, 1.0)
        th.test_raise(_test, AssertionError, "Modifying a locked object")

        def _test():
            sc.unlock()
            sc.add_hopping((0, 0, 0), 20, 0, 1.0)
        th.test_raise(_test,
                      AssertionError,  "Orb_i 20 out of range")

        def _test():
            sc.unlock()
            sc.add_hopping((0, 0, 0), 0, 20, 1.0)
        th.test_raise(_test, AssertionError, "Orb_j 20 out of range")

        def _test():
            sc.unlock()
            sc.add_hopping((0, 0, 0), 1, 0, 1.0)
        th.test_no_raise(_test, AssertionError)

    def test_set_modifier(self):
        """Test set_orb_pos_modifier"""
        sc = SuperCell(make_cell(), dim=(3, 3, 1))
        th = TestHelper(self)

        def _test():
            sc.lock()
            sc.set_orb_pos_modifier(make_cell)
        th.test_raise(_test, AssertionError, "Modifying a locked object")

        def _test():
            sc.unlock()
            sc.set_orb_pos_modifier(make_cell)
        th.test_no_raise(_test, AssertionError)

    def test_trim(self):
        """Test trim method. Also tested in tests for sample."""
        sc = SuperCell(make_cell(), dim=(3, 3, 1))
        th = TestHelper(self)

        def _test():
            sc.lock()
            sc.trim()
        th.test_raise(_test, AssertionError, "Modifying a locked object")

        def _test():
            sc.unlock()
            sc.add_vacancies([8])
            sc.trim()
        th.test_no_raise(_test, AssertionError)

    def test_sc_properties(self):
        """
        Test methods of SuperCell decorated with '@property'.

        :return: None
        """
        sc = SuperCell(make_cell(), dim=(3, 3, 1))
        th = TestHelper(self)
        self.assertEqual(sc.num_orb_pc, sc.prim_cell.num_orb)
        self.assertEqual(sc.num_orb_sc, sc._cached_data.orb_id_pc.shape[0])
        self.assertEqual(sc.num_vac, 0)
        sc_lat_vec_ref = sc.prim_cell.lat_vec.copy()
        sc_lat_vec_ref[0] *= 3
        sc_lat_vec_ref[1] *= 3
        th.test_equal_array(sc.sc_lat_vec, sc_lat_vec_ref)

    def test_get_orb_eng(self):
        """
        Test if SuperCell.get_orb_eng works as expected.

        :return: None
        """
        sc = SuperCell(make_cell(), dim=(3, 3, 1))
        orb_eng = sc.get_orb_eng()
        self.assertTupleEqual(orb_eng.shape, (18,))

    def test_get_orb_pos(self):
        """
        Test if SuperCell.get_orb_pos works as expected.

        :return: None
        """
        sc = SuperCell(make_cell(), dim=(3, 3, 1))
        orb_pos = sc.get_orb_pos()
        self.assertTupleEqual(orb_pos.shape, (18, 3))

    def test_get_hop_dr(self):
        """
        Test if SuperCell.get_hop and SuperCell.get_dr work as expected.

        :return: None
        """
        # Test hop_modifier
        sc = SuperCell(make_cell(), dim=(3, 3, 1))
        orb_map = sc.get_orb_map()

        id_pc_bra = [
            (0, 0, 0, 0),  # same item
            (1, 1, 0, 1),  # conjugate item
            (0, 0, 0, 1),  # new term
            (0, 2, 0, 0),  # new term
        ]
        id_pc_ket = [
            (0, 0, 0, 1),  # same item
            (1, 1, 0, 0),  # conjugate item
            (1, 2, 0, 0),  # new term
            (2, 0, 0, 1),  # new term
        ]
        id_sc_bra = orb_map.pc2sc_array(np.array(id_pc_bra))
        id_sc_ket = orb_map.pc2sc_array(np.array(id_pc_ket))
        for i in range(id_sc_bra.shape[0]):
            sc.add_hopping((0, 0, 0), id_sc_bra.item(i), id_sc_ket.item(i), energy=0.0)

        # Inspect the arrays
        hop_ind, hop_eng, dr = sc.get_hop()
        self.assertEqual(hop_ind.shape[0], dr.shape[0])

    def test_fast_algo(self):
        """
        Test the fast algorithm to build hopping terms.

        :return: None
        """
        # Test the core function to split hopping terms
        pc_hop_ind = np.array([
            (0, 0, 0, 1, 0),
            (0, 0, 0, 1, 1),
            (1, 0, 0, 1, 0),
            (0, -1, 0, 1, 1),
            (0, 0, 2, 1, 1)], dtype=np.int32
        )
        pc_hop_eng = np.array([1.0, -1.0, 2.5, 1.3, 0.6], dtype=np.complex128)
        pbc = np.array([0, 1, 0], dtype=np.int32)

        ind_pbc, eng_pbc, ind_free, eng_free = \
            core.split_pc_hop(pc_hop_ind, pc_hop_eng, pbc)
        ind_pbc_ref = pc_hop_ind[[0, 1, 3]]
        eng_pbc_ref = pc_hop_eng[[0, 1, 3]]
        ind_free_ref = pc_hop_ind[[2, 4]]
        eng_free_ref = pc_hop_eng[[2, 4]]

        th = TestHelper(self)
        th.test_equal_array(ind_pbc, ind_pbc_ref)
        th.test_equal_array(eng_pbc, eng_pbc_ref)
        th.test_equal_array(ind_free, ind_free_ref)
        th.test_equal_array(eng_free, eng_free_ref)

    def test_void(self):
        """
        Test if SuperCell behaves well under void primitive cells.

        :return: None
        """
        th = TestHelper(self)

        # Create supercell with empty primitive cell
        def _test():
            SuperCell(make_cell_orb(), dim=(3, 3, 1))
        th.test_raise(_test, AssertionError, r"Num_hop_pc should > 0")

        def _test():
            SuperCell(make_cell_empty(), dim=(3, 3, 1))
        th.test_raise(_test, AssertionError, r"Num_orb_pc should > 0")

        # Data updating
        prim_cell = make_cell()
        super_cell = SuperCell(prim_cell, dim=(3, 3, 1))
        prim_cell.unlock()
        prim_cell.remove_orbitals([0, 1])

        def _test():
            super_cell._sync_array()
        th.test_raise(_test, AssertionError, r"Num_orb_pc should > 0")

    def test_sc_inter_get_hop(self):
        """
        Test 'get_hop' method of 'SCInterHopping' class.

        :return: None
        """
        pc1 = make_cell()
        pc2 = make_cell()
        pc2.set_orbital(orb_i=0, position=(1./3, 1./3, 0.0))
        pc2.set_orbital(orb_i=1, position=(2./3, 2./3, 0.0))
        sc1 = SuperCell(pc1, dim=(3, 3, 1), pbc=(True, True, False))
        sc2 = SuperCell(pc2, dim=(3, 3, 1), pbc=(True, True, False))
        orb_map1 = sc1.get_orb_map()
        orb_map2 = sc2.get_orb_map()
        th = TestHelper(self)

        # Exception handling
        def _test():
            inter_hop2 = SCInterHopping(sc_bra=sc1, sc_ket=sc2)
            inter_hop2.get_hop()
        th.test_raise(_test, AssertionError, f"Num_hop should > 0")

        def _test():
            inter_hop2 = SCInterHopping(sc_bra=sc1, sc_ket=sc2)
            inter_hop2.add_hopping(rn=(0, 0, 0), orb_i=0, orb_j=1, energy=-1.2)
            inter_hop2.add_hopping(rn=(1, 0, 0), orb_i=0, orb_j=1, energy=-1.2)
            inter_hop2.get_hop(check_dup=True)
        th.test_raise(_test, AssertionError, r"Duplicate terms detected \(0, 1\)")

        # Normal case
        inter_hop = SCInterHopping(sc_bra=sc1, sc_ket=sc2)
        inter_hop.add_hopping(rn=(0, 0, 0), energy=-1.2,
                              orb_i=orb_map1.pc2sc((0, 0, 0, 0)),
                              orb_j=orb_map2.pc2sc((1, 2, 0, 1)))
        inter_hop.add_hopping(rn=(0, 0, 0), energy=1.5,
                              orb_i=orb_map1.pc2sc((2, 1, 0, 1)),
                              orb_j=orb_map2.pc2sc((1, 2, 0, 0)))
        inter_hop.add_hopping(rn=(0, 0, 0), energy=-0.7,
                              orb_i=orb_map1.pc2sc((2, 0, 0, 0)),
                              orb_j=orb_map2.pc2sc((0, 2, 0, 1)))

        hop_ind_ref = np.array([(0, 11),  (15, 10), (12, 5)])
        hop_eng_ref = np.array([-1.2, 1.5, -0.7])
        dr_ref = []
        orb_pos1 = sc1.get_orb_pos()
        orb_pos2 = sc2.get_orb_pos()
        for i in range(hop_ind_ref.shape[0]):
            dr_ref.append(orb_pos2[hop_ind_ref.item(i, 1)] - orb_pos1[hop_ind_ref.item(i, 0)])

        hop_ind, hop_eng, dr = inter_hop.get_hop()
        th.test_equal_array(hop_ind, hop_ind_ref)
        th.test_equal_array(hop_eng, hop_eng_ref)
        th.test_equal_array(dr, np.array(dr_ref))

    def test_orb_map_check_id_pc(self):
        th = TestHelper(self)
        sc = SuperCell(make_cell(), dim=(3, 3, 1))
        orb_map = sc.get_orb_map()

        # Tuple
        def _test():
            orb_map._check_id_pc((3, 0, 0, 0))
        th.test_raise(_test, AssertionError, r"^.+\[0\] out of range$")

        def _test():
            orb_map._check_id_pc((0, -1, 0, 0))
        th.test_raise(_test, AssertionError, r"^.+\[1\] out of range$")

        def _test():
            orb_map._check_id_pc((0, 0, 5, 0))
        th.test_raise(_test, AssertionError, r"^.+\[2\] out of range$")

        def _test():
            orb_map._check_id_pc((0, 0, 0, 5))
        th.test_raise(_test, AssertionError, r"^.+\[3\] out of range$")

        # Numpy array
        def _test():
            orb_map._check_id_pc(np.array((3, 0, 0, 0)))
        th.test_raise(_test, AssertionError, r"^.+\[0\] out of range$")

        def _test():
            orb_map._check_id_pc(np.array((0, -1, 0, 0)))
        th.test_raise(_test, AssertionError, r"^.+\[1\] out of range$")

        def _test():
            orb_map._check_id_pc(np.array((0, 0, 5, 0)))
        th.test_raise(_test, AssertionError, r"^.+\[2\] out of range$")

        def _test():
            orb_map._check_id_pc(np.array((0, 0, 0, 5)))
        th.test_raise(_test, AssertionError, r"^.+\[3\] out of range$")

        def _test():
            orb_map._check_id_pc((1, 0, 0, 1))
        th.test_no_raise(_test, AssertionError)

        def _test():
            orb_map._check_id_pc(np.array((1, 0, 0, 1)))
        th.test_no_raise(_test, AssertionError)

    def test_orb_map_sc2pc(self):
        vacancies = [1, 2, 9]
        sc = SuperCell(make_cell(), dim=(3, 3, 1))
        sc.set_vacancies(vacancies)
        orb_map = sc.get_orb_map()
        th = TestHelper(self)

        # Feed a wrong input and check if errors are raised.
        def _test():
            orb_map.sc2pc(21)
        th.test_raise(_test, AssertionError, r"id_sc .+ out of range")

        # Then test the normal case
        for i in range(sc.num_orb_sc):
            th.test_equal_array(np.array(orb_map.sc2pc(i)),
                                sc._cached_data.orb_id_pc[i])

    def test_orb_map_pc2sc(self):
        vacancies = [1, 2, 9]
        sc = SuperCell(make_cell(), dim=(3, 3, 1))
        sc.set_vacancies(vacancies)
        orb_map = sc.get_orb_map()
        th = TestHelper(self)

        # Test if the function can detect errors
        def _test():
            orb_map.pc2sc((1, 1, 0, 1))
        th.test_raise(_test, AssertionError, r"id_pc .+ is a vacancy")

        # Then test the normal case
        for i in range(sc.num_orb_sc):
            self.assertEqual(orb_map.pc2sc(sc._cached_data.orb_id_pc[i]), i)

    def test_orb_map_sc2pc_array(self):
        vacancies = [1, 2, 9]
        sc = SuperCell(make_cell(), dim=(3, 3, 1))
        sc.set_vacancies(vacancies)
        orb_map = sc.get_orb_map()
        th = TestHelper(self)

        # Assertion
        def _test():
            orb_map.sc2pc_array(np.eye(3))
        th.test_raise(_test, AssertionError, "id_sc_array should be a vector")

        def _test():
            orb_map.sc2pc_array(np.array([-1, 18], dtype=np.int64))
        th.test_raise(_test, AssertionError, r"id_sc .+ out of range")

        # Normal case
        id_sc = np.linspace(0, sc.num_orb_sc-1, sc.num_orb_sc,
                            dtype=np.int64)
        id_pc = orb_map.sc2pc_array(id_sc)
        th.test_equal_array(id_pc, sc._cached_data.orb_id_pc)

    def test_orb_map_pc2sc_array(self):
        vacancies = [1, 2, 9]
        sc = SuperCell(make_cell(), dim=(3, 3, 1))
        sc.set_vacancies(vacancies)
        orb_map = sc.get_orb_map()
        th = TestHelper(self)

        # Error handling
        def _test():
            orb_map.pc2sc_array(np.array([-1, 0, 0, 1], dtype=np.int32))
        th.test_raise(_test, AssertionError, r"Illegal shape of id_pc_array")

        def _test():
            orb_map.pc2sc_array(np.array([[-1, 0, 0, 1, 5]], dtype=np.int32))
        th.test_raise(_test, AssertionError, r"Illegal shape of id_pc_array")

        def _test():
            orb_map.pc2sc_array(np.array([[-1, 0, 0, 1]], dtype=np.int32))
        th.test_raise(_test, AssertionError, r"out of range")

        def _test():
            orb_map.pc2sc_array(np.array([[0, 0, 0, 1]], dtype=np.int32))
        th.test_raise(_test, AssertionError, r"is a vacancy")

        # Normal case
        id_sc_test = orb_map.pc2sc_array(sc._cached_data.orb_id_pc)
        id_sc_ref = np.linspace(0, sc.num_orb_sc-1, sc.num_orb_sc,
                                dtype=np.int64)
        th.test_equal_array(id_sc_test, id_sc_ref)


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