#! /usr/bin/env python

import unittest

import numpy as np
import matplotlib.pyplot as plt

from tbplas import (gen_lattice_vectors, gen_kpath, gen_kmesh,
                    PrimitiveCell, ANG, NM, frac2cart, TestHelper,
                    extend_prim_cell, HopDict, Visualizer, reshape_prim_cell,
                    PCHopNotFoundError)
import tbplas as tb


def make_cell():
    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")
    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_only():
    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():
    vectors = gen_lattice_vectors(a=2.46, b=2.46, gamma=60)
    cell = PrimitiveCell(vectors)
    return cell


def make_cell_rect():
    cell = PrimitiveCell(lat_vec=np.eye(3, 3), unit=1.0)
    cell.add_orbital((0, 0), energy=-1.0)
    cell.add_orbital((0.5, 0.5), energy=1.0)
    cell.add_hopping((0, 0), 0, 1, -1.0)
    return cell


class MyTest(unittest.TestCase):

    def test_init(self):
        """
        Test initialization of 'PrimitiveCell' class.

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

        # Check if feeding an illegal lattice vector and origin raises the right
        # exception

        def _test():
            PrimitiveCell(lat_vec=np.eye(2, 2))
        th.test_raise(_test, AssertionError, r"Lat_vec")

        def _test():
            PrimitiveCell(lat_vec=np.zeros((3, 3)))
        th.test_raise(_test, AssertionError, r".*")

        def _test():
            PrimitiveCell(origin=np.zeros(5))
        th.test_raise(_test, AssertionError, r"Length of origin")

        def _test():
            PrimitiveCell(unit=0.0)
        th.test_raise(_test, AssertionError, r"Unit")

        # Check unit conversion
        vectors = gen_lattice_vectors(a=1.5, b=2.0, c=2.5)
        cell = PrimitiveCell(vectors, unit=ANG)
        self.assertAlmostEqual(cell.lat_vec.item(0, 0), 0.15)
        self.assertAlmostEqual(cell.lat_vec.item(1, 1), 0.20)
        self.assertAlmostEqual(cell.lat_vec.item(2, 2), 0.25)
        cell = PrimitiveCell(vectors, unit=NM)
        self.assertAlmostEqual(cell.lat_vec.item(0, 0), 1.5)
        self.assertAlmostEqual(cell.lat_vec.item(1, 1), 2.0)
        self.assertAlmostEqual(cell.lat_vec.item(2, 2), 2.5)
        origin = np.array([1.0, 2.0, 3.0])
        cell = PrimitiveCell(origin=origin, unit=ANG)
        th.test_equal_array(cell.origin, origin*0.1)
        cell = PrimitiveCell(origin=origin, unit=NM)
        th.test_equal_array(cell.origin, origin)

        # Check other attributes
        th.test_equal_array(cell.lat_vec, np.eye(3))
        th.test_equal_array(cell.origin, origin)
        self.assertListEqual(cell.orbitals, [])
        self.assertEqual(cell.orb_pos.size, 0)
        self.assertEqual(cell.orb_eng.size, 0)
        self.assertEqual(cell.hop_ind.size, 0)
        self.assertEqual(cell.hop_eng.size, 0)
        self.assertEqual(cell.num_orb, 0)
        self.assertEqual(cell.num_hop, 0)
        self.assertEqual(cell.orb_pos_nm.size, 0)
        self.assertEqual(cell.orb_pos_ang.size, 0)
        self.assertEqual(cell.dr.size, 0)
        self.assertEqual(cell.dr_nm.size, 0)
        self.assertEqual(cell.dr_ang.size, 0)

    def test_lock(self):
        """
        Test if calling the functions that may change the structure of
        model raises errors when the model is locked.

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

        # add_orbital
        def _test():
            cell = make_cell()
            cell.lock()
            cell.add_orbital((1.2, 0.5), 0.1)
        th.test_raise(_test, AssertionError, r"Modifying a locked object")

        def _test():
            cell = make_cell()
            cell.lock()
            cell.add_orbital_cart((1.2, 0.5), 0.1)
        th.test_raise(_test, AssertionError, r"Modifying a locked object")

        # set_orbital
        def _test():
            cell = make_cell()
            cell.lock()
            cell.set_orbital(1, energy=0.25)
        th.test_raise(_test, AssertionError, r"Modifying a locked object")

        def _test():
            cell = make_cell()
            cell.lock()
            cell.set_orbital_cart(1, energy=0.25)
        th.test_raise(_test, AssertionError, r"Modifying a locked object")

        # remove_orbitals
        def _test():
            cell = make_cell()
            cell.lock()
            cell.remove_orbitals([0])
        th.test_raise(_test, AssertionError, r"Modifying a locked object")

        # add_hopping
        def _test():
            cell = make_cell()
            cell.lock()
            cell.add_hopping((-1, 0), 0, 0, 2.5)
        th.test_raise(_test, AssertionError, r"Modifying a locked object")

        # add_hopping_dict
        def _test():
            cell = make_cell()
            cell.lock()
            hop_dict = HopDict(2)
            hop_dict[(-1, 0, 0)] = np.random.rand(2, 2)
            cell.add_hopping((-1, 0), 0, 0, 2.5)
        th.test_raise(_test, AssertionError, r"Modifying a locked object")

        # remove_hopping
        def _test():
            cell = make_cell()
            cell.lock()
            cell.remove_hopping((0, 0), 0, 1)
        th.test_raise(_test, AssertionError, r"Modifying a locked object")

        # trim
        def _test():
            cell = make_cell()
            cell = extend_prim_cell(cell, (3, 3, 1))
            cell.remove_orbitals([8])
            cell.lock()
            cell.trim()
        th.test_raise(_test, AssertionError, r"Modifying a locked object")

        # apply_pbc
        def _test():
            cell = make_cell()
            cell = extend_prim_cell(cell, (3, 3, 1))
            cell.lock()
            cell.apply_pbc((False, True, False))
        th.test_raise(_test, AssertionError, r"Modifying a locked object")

        # reset_lattice
        def _test():
            cell = make_cell()
            cell.lock()
            cell.reset_lattice(lat_vec=np.eye(3, 3))
        th.test_raise(_test, AssertionError, r"Modifying a locked object")

    def test_unlock(self):
        """
        Test if calling the functions that may change the structure of
        model raises errors when the model is unlocked.

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

        # add_orbital
        def _test():
            cell = make_cell()
            cell.lock()
            cell.unlock()
            cell.add_orbital((1.2, 0.5), 0.1)
        th.test_no_raise(_test, AssertionError)

        def _test():
            cell = make_cell()
            cell.lock()
            cell.unlock()
            cell.add_orbital_cart((1.2, 0.5), 0.1)
        th.test_no_raise(_test, AssertionError)

        # set_orbital
        def _test():
            cell = make_cell()
            cell.lock()
            cell.unlock()
            cell.set_orbital(1, energy=0.25)
        th.test_no_raise(_test, AssertionError)

        def _test():
            cell = make_cell()
            cell.lock()
            cell.unlock()
            cell.set_orbital_cart(1, energy=0.25)
        th.test_no_raise(_test, AssertionError)

        # remove_orbitals
        def _test():
            cell = make_cell()
            cell.lock()
            cell.unlock()
            cell.remove_orbitals([0])
        th.test_no_raise(_test, AssertionError)

        # add_hopping
        def _test():
            cell = make_cell()
            cell.lock()
            cell.unlock()
            cell.add_hopping((-1, 0), 0, 0, 2.5)
        th.test_no_raise(_test, AssertionError)

        # add_hopping_dict
        def _test():
            cell = make_cell()
            cell.lock()
            cell.unlock()
            hop_dict = HopDict(2)
            hop_dict[(-1, 0, 0)] = np.random.rand(2, 2)
            cell.add_hopping((-1, 0), 0, 0, 2.5)
        th.test_no_raise(_test, AssertionError)

        # remove_hopping
        def _test():
            cell = make_cell()
            cell.lock()
            cell.unlock()
            cell.remove_hopping((0, 0), 0, 1)
        th.test_no_raise(_test, AssertionError)

        # trim
        def _test():
            cell = make_cell()
            cell = extend_prim_cell(cell, (3, 3, 1))
            cell.remove_orbitals([8])
            cell.lock()
            cell.unlock()
            cell.trim()
        th.test_no_raise(_test, AssertionError)

        # apply_pbc
        def _test():
            cell = make_cell()
            cell = extend_prim_cell(cell, (3, 3, 1))
            cell.lock()
            cell.unlock()
            cell.apply_pbc((False, True, False))
        th.test_no_raise(_test, AssertionError)

        # reset_lattice
        def _test():
            cell = make_cell()
            cell.lock()
            cell.unlock()
            cell.reset_lattice(lat_vec=np.eye(3, 3))
        th.test_no_raise(_test, AssertionError)

    def test_add_orbital(self):
        """
        Test if add_orbital works as expected.

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

        # Test if feeding an illegal position raises the right exception
        cell = make_cell()

        def _test():
            cell.add_orbital((1.0,), energy=1.0)
        th.test_raise(_test, AssertionError, r"Length of pos")

        # Test the normal case
        cell = make_cell()
        orbitals = cell.orbitals
        self.assertEqual(cell.num_orb, 2)
        self.assertAlmostEqual(np.linalg.norm(orbitals[0].position), 0.0)
        self.assertAlmostEqual(np.linalg.norm(orbitals[1].position),
                               4.714045207910317e-1)
        self.assertAlmostEqual(orbitals[0].energy, 0.0)
        self.assertAlmostEqual(orbitals[1].energy, 0.0)
        self.assertEqual(orbitals[0].label, "C_pz")
        self.assertEqual(orbitals[1].label, "C_pz")

        # Test adding orbital with Cartesian coordinates
        vectors = gen_lattice_vectors(a=2.46, b=2.46, gamma=60)
        cell_ref = PrimitiveCell(vectors)
        cell_ref.add_orbital((1. / 3, 1. / 3), 0.0, label="C_pz")
        cell_ref.add_orbital((2. / 3, 2. / 3), 0.0, label="C_pz")
        orb_pos_ref = frac2cart(cell_ref.lat_vec, cell_ref.orb_pos)

        cell_test = PrimitiveCell(vectors)
        cell_test.add_orbital_cart(orb_pos_ref[0] * 10, unit=ANG, label="C_pz")
        cell_test.add_orbital_cart(orb_pos_ref[1], unit=NM, label="C_pz")

        th.test_equal_array(cell_ref.orb_pos, cell_test.orb_pos, almost=True)
        orb_pos_test = frac2cart(cell_test.lat_vec, cell_test.orb_pos)
        th.test_equal_array(orb_pos_ref, orb_pos_test, almost=True)

    def test_set_orbital(self):
        """
        Test if set_orbital and set_orbital_cart works as expected.

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

        # Test error handling for illegal orbital index
        cell = make_cell()

        def _test():
            cell.set_orbital(3, energy=4.2)
        th.test_raise(_test, AssertionError,
                      r"Orbital index .+ out of range")

        def _test():
            cell.set_orbital_cart(3, energy=4.2)
        th.test_raise(_test, AssertionError,
                      r"Orbital index .+ out of range")

        # Test error handling for illegal position
        cell = make_cell()

        def _test():
            cell.set_orbital(0, position=(1.5,))
        th.test_raise(_test, AssertionError, r"Length of pos")

        def _test():
            cell.set_orbital_cart(0, position=(1.5,))
        th.test_raise(_test, AssertionError, r"Length of pos")

        # Test setting both position and energy
        cell = make_cell()
        cell.set_orbital(1, position=(1.5, 0.0), energy=4.2)
        self.assertEqual(cell.num_orb, 2)
        self.assertAlmostEqual(cell.orbitals[1].energy, 4.2)
        self.assertAlmostEqual(np.linalg.norm(cell.orbitals[-1].position),
                               1.5)

        # Test setting position
        cell = make_cell()
        cell.set_orbital(1, position=(1.5, 0.0))
        self.assertEqual(cell.num_orb, 2)
        self.assertAlmostEqual(cell.orbitals[1].energy, 0.0)
        self.assertAlmostEqual(np.linalg.norm(cell.orbitals[-1].position),
                               1.5)

        # Test setting energy
        cell = make_cell()
        cell.set_orbital(1, energy=4.2)
        self.assertEqual(cell.num_orb, 2)
        self.assertAlmostEqual(cell.orbitals[1].energy, 4.2)
        self.assertAlmostEqual(np.linalg.norm(cell.orbitals[-1].position),
                               4.714045207910317e-1)

        # Test setting label
        cell = make_cell()
        cell.set_orbital(1, label="C_pz_2")
        self.assertEqual(cell.orbitals[1].label, "C_pz_2")

        # Test setting orbital with Cartesian coordinates
        cell_ref = make_cell()
        cell_ref.set_orbital(0, position=(1. / 3, 1. / 3))
        cell_ref.set_orbital(1, position=(2. / 3, 2. / 3))
        orb_pos_ref = frac2cart(cell_ref.lat_vec, cell_ref.orb_pos)

        cell_test = make_cell()
        cell_test.set_orbital_cart(0, position=orb_pos_ref[0]*10, unit=ANG)
        cell_test.set_orbital_cart(1, position=orb_pos_ref[1], unit=NM)

        th.test_equal_array(cell_ref.orb_pos, cell_test.orb_pos, almost=True)
        orb_pos_test = frac2cart(cell_test.lat_vec, cell_test.orb_pos)
        th.test_equal_array(orb_pos_ref, orb_pos_test, almost=True)

        # Void case
        cell = make_cell_empty()

    def test_get_orbital(self):
        """
        Test get_orbital.

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

        # Error handling
        cell = make_cell()

        def _test():
            cell.get_orbital(3)
        th.test_raise(_test, AssertionError,
                      r"Orbital index .+ out of range")

        # Normal case
        cell = make_cell()
        orbital = cell.get_orbital(1)
        self.assertAlmostEqual(orbital.energy, 0.0)
        self.assertAlmostEqual(np.linalg.norm(cell.orbitals[1].position),
                               4.714045207910317e-1)

        # Void case
        cell = make_cell_empty()

        def _test():
            cell.get_orbital(3)
        th.test_raise(_test, AssertionError,
                      r"Orbital index .+ out of range")

    def test_remove_orbitals(self):
        """
        Test if remove_orbitals works as expected.

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

        # Test if feeding an illegal orbital index raises the right error
        cell = make_cell()

        def _test():
            cell.remove_orbitals([3])
        th.test_raise(_test, AssertionError,
                      r"Orbital index .+ out of range")

        # removing orbital #0
        cell = make_cell()
        cell.remove_orbitals([0])
        self.assertEqual(cell.num_orb, 1)
        self.assertEqual(cell.num_hop, 0)

        # removing orbital #1
        cell = make_cell()
        cell.remove_orbitals([1])
        self.assertEqual(cell.num_orb, 1)
        self.assertEqual(cell.num_hop, 0)

        # adding orbital #2
        cell = make_cell()
        cell.add_orbital((0.5, 0.5), energy=-0.1)
        cell.add_hopping((0, 0), 0, 2, -1.5)
        cell.add_hopping((0, 0), 1, 2, -1.9)

        # removing orbital #0
        cell.remove_orbitals([0])
        self.assertEqual(cell.num_orb, 2)
        self.assertEqual(cell.num_hop, 1)
        self.assertAlmostEqual(cell.get_hopping((0, 0, 0), 0, 1), -1.9)

        # removing orbital #1
        cell = make_cell()
        cell.add_orbital((0.5, 0.5), energy=-0.1)
        cell.add_hopping((0, 0), 0, 2, -1.5)
        cell.add_hopping((0, 0), 1, 2, -1.9)
        cell.remove_orbitals([1])
        self.assertEqual(cell.num_orb, 2)
        self.assertEqual(cell.num_hop, 1)
        self.assertAlmostEqual(cell.get_hopping((0, 0, 0), 0, 1), -1.5)

        # removing orbital #2
        cell = make_cell()
        cell.add_orbital((0.5, 0.5), energy=-0.1)
        cell.add_hopping((0, 0), 0, 2, -1.5)
        cell.add_hopping((0, 0), 1, 2, -1.9)
        cell.remove_orbitals([2])
        self.assertEqual(cell.num_orb, 2)
        self.assertEqual(cell.num_hop, 3)
        self.assertAlmostEqual(cell.get_hopping((0, 0, 0), 0, 1), -2.7)

    def test_add_hopping(self):
        """
        Test if add_hopping works as expected.

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

        # Error handling
        cell = make_cell()

        def _test():
            cell.add_hopping(rn=(2,), orb_i=2, orb_j=0, energy=-2.8)
        th.test_raise(_test, AssertionError, r"Length of rn")

        def _test():
            cell.add_hopping(rn=(2, 0), orb_i=2, orb_j=0, energy=-2.8)
        th.test_raise(_test, AssertionError, r"Orb_i .+ out of range")

        def _test():
            cell.add_hopping(rn=(2, 0), orb_i=0, orb_j=2, energy=-2.8)
        th.test_raise(_test, AssertionError, r"Orb_j .+ out of range")

        def _test():
            cell.add_hopping(rn=(0, 0), orb_i=0, orb_j=0, energy=-2.8)
        th.test_raise(_test, AssertionError, r".+ is a diagonal term")

        # The normal case
        cell = make_cell()
        hop_ind, hop_eng = cell.hop_ind, cell.hop_eng
        self.assertEqual(cell.num_hop, 3)
        th.test_equal_array(hop_ind[0], np.array([0, 0, 0, 0, 1]))
        th.test_equal_array(hop_ind[1], np.array([1, 0, 0, 1, 0]))
        th.test_equal_array(hop_ind[2], np.array([0, 1, 0, 1, 0]))
        self.assertAlmostEqual(hop_eng.item(0), -2.7)
        self.assertAlmostEqual(hop_eng.item(1), -2.7)
        self.assertAlmostEqual(hop_eng.item(2), -2.7)

        # Updating an existing hopping term
        cell = make_cell()
        cell.add_hopping(rn=(0, 0), orb_i=0, orb_j=1, energy=-2.8)
        self.assertEqual(cell.num_hop, 3)
        self.assertAlmostEqual(cell.get_hopping((0, 0, 0), 0, 1), -2.8)

        # Updating an existing conjugate hopping term
        cell = make_cell()
        cell.add_hopping(rn=(-1, 0), orb_i=0, orb_j=1, energy=-2.7)
        cell.add_hopping(rn=(0, 0), orb_i=1, orb_j=0, energy=-2.8)
        self.assertEqual(cell.num_hop, 3)
        self.assertAlmostEqual(cell.get_hopping((-1, 0, 0), 0, 1), -2.7)
        self.assertAlmostEqual(cell.get_hopping((0, 0, 0), 0, 1), -2.8)

    def test_get_hopping(self):
        """
        Test get_hopping.

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

        # Error handling
        cell = make_cell()

        def _test():
            cell.add_hopping(rn=(2,), orb_i=2, orb_j=0, energy=-2.8)
        th.test_raise(_test, AssertionError, r"Length of rn")

        def _test():
            cell.add_hopping(rn=(2, 0), orb_i=2, orb_j=0, energy=-2.8)
        th.test_raise(_test, AssertionError, r"Orb_i .+ out of range")

        def _test():
            cell.add_hopping(rn=(2, 0), orb_i=0, orb_j=2, energy=-2.8)
        th.test_raise(_test, AssertionError, r"Orb_j .+ out of range")

        def _test():
            cell.add_hopping(rn=(0, 0), orb_i=0, orb_j=0, energy=-2.8)
        th.test_raise(_test, AssertionError, r".+ is a diagonal term")

        def _test():
            cell.get_hopping((-2, 0), 0, 1)
        th.test_raise(_test, PCHopNotFoundError,
                      r"hopping term .+ not found")

        # The normal case
        energy = cell.get_hopping((0, 0), 0, 1)
        self.assertAlmostEqual(energy, -2.7)

        # Conjugate terms
        self.assertAlmostEqual(cell.get_hopping((0, 0, 0), 1, 0), -2.7)
        self.assertAlmostEqual(cell.get_hopping((-1, 0, 0), 0, 1), -2.7)
        self.assertAlmostEqual(cell.get_hopping((0, -1, 0), 0, 1), -2.7)

        # Void case
        cell = make_cell_orb_only()
        th.test_raise(_test, PCHopNotFoundError,
                      r"hopping term .+ not found")
        cell = make_cell_empty()
        th.test_raise(_test, AssertionError,
                      r"Orb_i .+ out of range")

    def test_remove_hopping(self):
        """
        Test if remove_hopping works as expected.

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

        # Error handling
        cell = make_cell()

        def _test():
            cell.add_hopping(rn=(2,), orb_i=2, orb_j=0, energy=-2.8)
        th.test_raise(_test, AssertionError, r"Length of rn")

        def _test():
            cell.add_hopping(rn=(2, 0), orb_i=2, orb_j=0, energy=-2.8)
        th.test_raise(_test, AssertionError, r"Orb_i .+ out of range")

        def _test():
            cell.add_hopping(rn=(2, 0), orb_i=0, orb_j=2, energy=-2.8)
        th.test_raise(_test, AssertionError, r"Orb_j .+ out of range")

        def _test():
            cell.add_hopping(rn=(0, 0), orb_i=0, orb_j=0, energy=-2.8)
        th.test_raise(_test, AssertionError, r".+ is a diagonal term")

        def _test():
            cell.remove_hopping((-2, 0), 0, 1)
        th.test_raise(_test, PCHopNotFoundError,
                      r"hopping term .+ not found")
        self.assertEqual(cell.num_orb, 2)
        self.assertEqual(cell.num_hop, 3)

        # The normal case
        cell = make_cell()
        cell.remove_hopping((0, 0), 0, 1)
        self.assertEqual(cell.num_orb, 2)
        self.assertEqual(cell.num_hop, 2)

        # Remove conjugate part
        cell = make_cell()
        cell.remove_hopping((-1, 0), 0, 1)
        cell.remove_hopping((0, -1), 0, 1)
        self.assertEqual(cell.num_orb, 2)
        self.assertEqual(cell.num_hop, 1)

        # Void case
        cell = make_cell_orb_only()
        th.test_raise(_test, PCHopNotFoundError,
                      r"hopping term .+ not found")
        cell = make_cell_empty()
        th.test_raise(_test, AssertionError,
                      r"Orb_i .+ out of range")

    def test_check_dim(self):
        """
        Test function 'check_dim'.

        :return: None
        """
        cell = make_cell()
        cell2 = extend_prim_cell(cell, dim=(1, 1, 1))
        th = TestHelper(self)

        def _test():
            cell2.check_dim()
        th.test_raise(_test, AssertionError,
                      r"Duplicate terms detected \(0, 1\)")
        cell2 = extend_prim_cell(cell, dim=(2, 1, 1))
        th.test_raise(_test, AssertionError,
                      r"Duplicate terms detected \(0, 1\)")
        cell2 = extend_prim_cell(cell, dim=(1, 2, 1))
        th.test_raise(_test, AssertionError,
                      r"Duplicate terms detected \(0, 1\)")
        cell2 = extend_prim_cell(cell, dim=(2, 2, 1))
        th.test_no_raise(_test, AssertionError)

    def test_sync_array(self):
        """
        Test if sync_array works as expected.

        :return: None
        """
        # DO NOT add or remove anything
        cell = make_cell()
        self.assertEqual(cell.orb_pos.shape, (2, 3))
        self.assertEqual(cell.orb_eng.shape, (2,))
        self.assertEqual(cell.hop_ind.shape, (3, 5))
        self.assertEqual(cell.hop_eng.shape, (3,))

        # Add orbital
        cell = make_cell()
        cell.add_orbital((0.0, 0.5), 0.15)
        self.assertEqual(cell.orb_pos.shape, (3, 3))
        self.assertEqual(cell.orb_eng.shape, (3,))
        self.assertEqual(cell.hop_ind.shape, (3, 5))
        self.assertEqual(cell.hop_eng.shape, (3,))

        # Add hopping term
        cell = make_cell()
        cell.add_hopping((0, 2), 0, 1, -1.5)
        self.assertEqual(cell.orb_pos.shape, (2, 3))
        self.assertEqual(cell.orb_eng.shape, (2,))
        self.assertEqual(cell.hop_ind.shape, (4, 5))
        self.assertEqual(cell.hop_eng.shape, (4,))

        # Removing orbitals
        cell = make_cell()
        cell.remove_orbitals([0])
        self.assertEqual(cell.orb_pos.shape, (1, 3))
        self.assertEqual(cell.orb_eng.shape, (1,))
        self.assertEqual(cell.hop_ind.shape, (0,))
        self.assertEqual(cell.hop_eng.shape, (0,))
        cell.remove_orbitals([0])
        self.assertEqual(cell.orb_pos.size, 0)
        self.assertEqual(cell.orb_eng.size, 0)
        self.assertEqual(cell.hop_ind.size, 0)
        self.assertEqual(cell.hop_eng.size, 0)

        # Removing hopping terms
        cell = make_cell()
        cell.remove_hopping((0, 0), 0, 1)
        self.assertEqual(cell.orb_pos.shape, (2, 3))
        self.assertEqual(cell.orb_eng.shape, (2,))
        self.assertEqual(cell.hop_ind.shape, (2, 5))
        self.assertEqual(cell.hop_eng.shape, (2,))
        cell.remove_hopping((-1, 0), 0, 1)
        self.assertEqual(cell.orb_pos.shape, (2, 3))
        self.assertEqual(cell.orb_eng.shape, (2,))
        self.assertEqual(cell.hop_ind.shape, (1, 5))
        self.assertEqual(cell.hop_eng.shape, (1,))
        cell.remove_hopping((0, 1), 1, 0)
        self.assertEqual(cell.orb_pos.shape, (2, 3))
        self.assertEqual(cell.orb_eng.shape, (2,))
        self.assertEqual(cell.hop_ind.shape, (0,))
        self.assertEqual(cell.hop_eng.shape, (0,))

        # Void cases
        cell = make_cell_orb_only()
        cell._sync_array()
        self.assertEqual(cell.orb_pos.shape, (2, 3))
        self.assertEqual(cell.orb_eng.shape, (2,))
        self.assertEqual(cell.hop_ind.size, 0)
        self.assertEqual(cell.hop_eng.size, 0)
        cell = make_cell_empty()
        cell._sync_array()
        self.assertEqual(cell.orb_pos.size, 0)
        self.assertEqual(cell.orb_eng.size, 0)
        self.assertEqual(cell.hop_ind.size, 0)
        self.assertEqual(cell.hop_eng.size, 0)

    def test_plot(self):
        """
        Test plotting of orbitals and hopping terms.

        :return: None
        """
        cell = make_cell()
        cell.plot()
        cell = make_cell_orb_only()
        cell.plot()
        cell = make_cell_empty()
        cell.plot()

        # User-defined orbital colors
        def orb_color(orbital_list):
            scatter_color = ["r" if i % 2 == 0 else "b"
                             for i in range(len(orbital_list))]
            return scatter_color

        cell = make_cell()
        cell.plot(orb_color=orb_color, hop_color="gray")

    def test_print(self):
        """
        Test printing attributes.

        :return: None
        """
        cell = make_cell()
        cell.print()
        cell = make_cell_orb_only()
        cell.print()
        cell = make_cell_empty()
        cell.print()

    def test_print_hk(self):
        """
        Test the print_hk method.

        :return: None
        """
        th = TestHelper(self)
        cell = make_cell()
        cell.print_hk(convention=1)
        cell.print_hk(convention=2)

        # Void cases
        def _test():
            cell.print_hk()
        cell = make_cell_orb_only()
        th.test_raise(_test, AssertionError, r"Num_hop should > 0")
        cell = make_cell_empty()
        th.test_raise(_test, AssertionError,
                      r"Num_orb should > 0")

        def _test():
            cell.print_hk(3)
        cell = make_cell()
        th.test_raise(_test, AssertionError, r"Illegal convention 3")

    def test_calc_bands_dos(self):
        """
        Test band structure and dos calculation.

        :return: None
        """
        # Calculate band structure.
        cell = make_cell()
        k_points = np.array([
            [0.0, 0.0, 0.0],
            [2./3, 1./3, 0.0],
            [1./2, 0.0, 0.0],
            [0.0, 0.0, 0.0],
        ])
        k_path, k_idx = gen_kpath(k_points, [40, 40, 40])
        k_len, bands = tb.calc_bands(cell, k_path)
        num_bands = bands.shape[1]
        for i in range(num_bands):
            plt.plot(k_len, bands[:, i], color="red", linewidth=1.0)
        plt.show()

        # Calculate DOS.
        k_points = gen_kmesh((120, 120, 1))
        energies, dos = tb.calc_dos(cell, k_points)
        plt.plot(energies, dos)
        plt.show()

    def test_lattice(self):
        """
        Test lattice operation functions.

        :return: None
        """
        lat_vec = np.zeros((3, 3))
        lat_vec[0, 0] = 1.0
        lat_vec[1, 1] = 2.0
        lat_vec[2, 2] = 3.0
        cell = PrimitiveCell(lat_vec, unit=1.0)
        self.assertAlmostEqual(cell.get_lattice_volume(), 6.0)
        self.assertAlmostEqual(cell.get_lattice_area("a"), 6.0)
        self.assertAlmostEqual(cell.get_lattice_area("b"), 3.0)
        self.assertAlmostEqual(cell.get_lattice_area("c"), 2.0)
        th = TestHelper(self)
        recip_test = cell.get_reciprocal_vectors()
        recip_ref = np.zeros((3, 3))
        recip_ref[0, 0] = 2 * np.pi
        recip_ref[1, 1] = 2 * np.pi / 2.0
        recip_ref[2, 2] = 2 * np.pi / 3.0
        th.test_equal_array(recip_test, recip_ref)

        def _test():
            cell.get_lattice_area("d")
        th.test_raise(_test, AssertionError, f"Wrong direction d")

    def test_orb_pos_cart(self):
        """
        Test orb_pos_nm and orb_pos_ang.

        :return: None
        """
        th = TestHelper(self)
        cell = make_cell()
        orb_pos_ref = frac2cart(cell.lat_vec, cell.orb_pos)
        orb_pos_test = cell.orb_pos_nm
        th.test_equal_array(orb_pos_ref, orb_pos_test, almost=True)
        orb_pos_ref *= 10.0
        orb_pos_test = cell.orb_pos_ang
        th.test_equal_array(orb_pos_ref, orb_pos_test, almost=True)

        # Void cases
        cell = make_cell_orb_only()
        orb_pos_ref = frac2cart(cell.lat_vec, cell.orb_pos)
        orb_pos_test = cell.orb_pos_nm
        th.test_equal_array(orb_pos_ref, orb_pos_test, almost=True)
        orb_pos_ref *= 10.0
        orb_pos_test = cell.orb_pos_ang
        th.test_equal_array(orb_pos_ref, orb_pos_test, almost=True)

        cell = make_cell_empty()
        orb_pos_ref = frac2cart(cell.lat_vec, cell.orb_pos)
        orb_pos_test = cell.orb_pos_nm
        th.test_equal_array(orb_pos_ref, orb_pos_test, almost=True)
        orb_pos_ref *= 10.0
        orb_pos_test = cell.orb_pos_ang
        th.test_equal_array(orb_pos_ref, orb_pos_test, almost=True)

    def test_dr_cart(self):
        """
        Test dr, dr_nm and dr_ang.

        :return: None.
        """
        th = TestHelper(self)
        cell = make_cell()
        dr_ref = frac2cart(cell.lat_vec, cell.dr)
        dr_test = cell.dr_nm
        th.test_equal_array(dr_ref, dr_test, almost=True)
        dr_ref *= 10.0
        dr_test = cell.dr_ang
        th.test_equal_array(dr_ref, dr_test, almost=True)

        # Void cases
        cell = make_cell_orb_only()
        dr_ref = frac2cart(cell.lat_vec, cell.dr)
        dr_test = cell.dr_nm
        th.test_equal_array(dr_ref, dr_test, almost=True)
        dr_ref *= 10.0
        dr_test = cell.dr_ang
        th.test_equal_array(dr_ref, dr_test, almost=True)
        cell = make_cell_empty()
        dr_ref = frac2cart(cell.lat_vec, cell.dr)
        dr_test = cell.dr_nm
        th.test_equal_array(dr_ref, dr_test, almost=True)
        dr_ref *= 10.0
        dr_test = cell.dr_ang
        th.test_equal_array(dr_ref, dr_test, almost=True)

    def test_trim(self):
        """Test trim method"""
        cell = extend_prim_cell(make_cell(), (3, 3, 1))
        cell.remove_orbitals([8, 14])
        cell.plot()
        cell.trim()
        cell.plot()

    def test_apply_pbc(self):
        """
        Test function 'apply_pbc'.

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

        # Assertion
        def _test():
            make_cell().apply_pbc((True, True))
        th.test_raise(_test, AssertionError, "Length of")

        # Normal case
        lat_frac = np.array([[1, 0, 0], [-1, 2, 0], [0, 0, 1]])
        rect_cell = reshape_prim_cell(make_cell(), lat_frac)

        # GNR along AM direction
        gnr = extend_prim_cell(rect_cell, dim=(3, 3, 1))
        gnr.plot(with_conj=False)
        gnr.apply_pbc(pbc=(False, True, False))
        gnr.plot(with_conj=False)
        gnr.trim()
        gnr.plot(with_conj=False)

        k_points = np.array([
            [0.0, -0.5, 0.0],
            [0.0, 0.0, 0.0],
            [0.0, 0.5, 0.0],
        ])
        k_label = ["X", "G", "X"]
        k_path, k_idx = gen_kpath(k_points, [40, 40])
        k_len, bands = tb.calc_bands(gnr, k_path)
        Visualizer().plot_bands(k_len, bands, k_idx, k_label)

        # GNR along ZZ direction
        gnr = extend_prim_cell(rect_cell, dim=(3, 3, 1))
        gnr.plot(with_conj=False)
        gnr.apply_pbc(pbc=(True, False, False))
        gnr.plot(with_conj=False)
        gnr.trim()
        gnr.plot(with_conj=False)

        k_points = np.array([
            [-0.5, 0.0, 0.0],
            [0.0, 0.0, 0.0],
            [0.5, 0.0, 0.0],
        ])
        k_label = ["X", "G", "X"]
        k_path, k_idx = gen_kpath(k_points, [40, 40])
        k_len, bands = tb.calc_bands(gnr, k_path)
        Visualizer().plot_bands(k_len, bands, k_idx, k_label)

        # Void cases
        for i1 in (True, False):
            for i2 in (True, False):
                for i3 in (True, False):
                    cell = make_cell_orb_only()
                    cell.apply_pbc(pbc=(i1, i2, i3))
                    cell.plot()
        for i1 in (True, False):
            for i2 in (True, False):
                for i3 in (True, False):
                    cell = make_cell_empty()
                    cell.apply_pbc(pbc=(i1, i2, i3))
                    cell.plot()

    def test_reset_lattice(self):
        """Test reset_lattice method."""
        c = np.sqrt(2.0) * 0.5
        new_lat_vec = np.array([
            [c, c, 0.0],
            [-c, c, 0.0],
            [0, 0, 1]])
        new_origin = np.array([0.5, 0.5, 0.0])

        # Fix orbital
        cell = make_cell_rect()
        cell.plot()
        cell.reset_lattice(lat_vec=new_lat_vec, origin=new_origin, unit=1.0)
        cell.plot()

        # Free orbital
        cell = make_cell_rect()
        cell.reset_lattice(lat_vec=new_lat_vec, origin=new_origin, unit=1.0,
                           fix_orb=False)
        cell.plot()


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