#! /usr/bin/env python

import unittest
import os

import numpy as np

import tbplas as tb
from tbplas import (gen_lattice_vectors, PrimitiveCell, SuperCell,
                    SCInterHopping, Sample, CachedSample, Timer, TestHelper,
                    RegularOrbitalMap)


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)
    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_test_set():
    # Make 3 supercells
    sc1 = SuperCell(make_cell(), dim=(3, 3, 1), pbc=(True, True, False))
    sc2 = SuperCell(make_cell(), dim=(3, 3, 1), pbc=(True, True, False))
    sc3 = SuperCell(make_cell(), dim=(3, 3, 1), pbc=(True, True, False))
    orb_map1 = sc1.get_orb_map()
    orb_map2 = sc2.get_orb_map()

    # Make 2 inter-hopping instances
    inter_hop1 = SCInterHopping(sc_bra=sc1, sc_ket=sc2)
    inter_hop1.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_hop1.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_hop1.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)))

    inter_hop2 = SCInterHopping(sc_bra=sc2, sc_ket=sc3)
    inter_hop2.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_hop2.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_hop2.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)))
    return sc1, sc2, sc3, inter_hop1, inter_hop2


class MyTest(unittest.TestCase):

    def test_init(self):
        """
        Test initialization of 'Sample' and 'CachedSample' class.

        :return: None
        """
        sc1 = SuperCell(make_cell(), dim=(3, 3, 1), pbc=(True, True, False))
        sc2 = SuperCell(make_cell(), dim=(3, 3, 1), pbc=(True, True, False))
        om1 = sc1.get_orb_map()
        om2 = sc2.get_orb_map()
        th = TestHelper(self)

        # Exception handling
        def _test():
            Sample()
        th.test_raise(_test, AssertionError,
                      r"No components assigned to sample")

        def _test():
            Sample(sc1, sc2, make_cell())
        th.test_raise(_test, AssertionError,
                      r"component .+ should be instance of SuperCell or"
                      r" SCInterHopping")

        def _test():
            sc3 = SuperCell(make_cell(), dim=(3, 3, 1), pbc=(True, True, False))
            inter_hop2 = SCInterHopping(sc1, sc3)
            Sample(sc1, sc2, inter_hop2)
        th.test_raise(_test, AssertionError,
                      r".+ of inter_hop .+ not included in args")

        def _test():
            sc3 = SuperCell(make_cell(), dim=(3, 3, 1), pbc=(True, True, False))
            inter_hop2 = SCInterHopping(sc3, sc1)
            Sample(sc1, sc2, inter_hop2)
        th.test_raise(_test, AssertionError,
                      r".+ of inter_hop .+ not included in args")

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

        def _test():
            Sample(sc1)
            Sample(sc1, sc2)
            Sample(sc1, sc2, inter_hop)
        th.test_no_raise(_test, AssertionError)

        # All arrays should be empty
        sample = Sample(sc1, sc2, inter_hop)
        self.assertEqual(sample.orb_eng.size, 0)
        self.assertEqual(sample.orb_pos.size, 0)
        self.assertEqual(sample.hop_ind.size, 0)
        self.assertEqual(sample.hop_eng.size, 0)
        self.assertEqual(sample.dr.size, 0)

        # Similar for CachedSample
        os.system("rm -rf test")
        sample.save("test")
        cached_sample = CachedSample("test")
        self.assertEqual(cached_sample.orb_eng.size, 0)
        self.assertEqual(cached_sample.orb_pos.size, 0)
        self.assertEqual(cached_sample.hop_ind.size, 0)
        self.assertEqual(cached_sample.hop_eng.size, 0)
        self.assertEqual(cached_sample.dr.size, 0)
        os.system("rm -rf test")

    def test_sample_private(self):
        """
        Test '__get_num_orb' and '__get_ind_start' methods of 'Sample' class.

        :return: None
        """
        sc1, sc2, sc3, inter_hop1, inter_hop2 = make_test_set()

        # Normal case with 1 supercell
        sample = Sample(sc1)
        num_orb = sample._get_num_orb_sc()
        ind_start = sample._get_ind_start()
        self.assertListEqual(num_orb, [18])
        self.assertListEqual(ind_start, [0])

        # Normal case with 2 supercells and 1 inter-hopping
        sample = Sample(sc1, sc2, inter_hop1)
        num_orb = sample._get_num_orb_sc()
        ind_start = sample._get_ind_start()
        self.assertListEqual(num_orb, [18, 18])
        self.assertListEqual(ind_start, [0, 18])

        # Normal case with 3 supercells and 2 inter-hopping
        sample = Sample(sc1, sc2, sc3, inter_hop1, inter_hop2)
        num_orb = sample._get_num_orb_sc()
        ind_start = sample._get_ind_start()
        self.assertListEqual(num_orb, [18, 18, 18])
        self.assertListEqual(ind_start, [0, 18, 36])

        # Abnormal case with 2 supercells and no inter-hopping
        sample = Sample(sc1, sc2)
        num_orb = sample._get_num_orb_sc()
        ind_start = sample._get_ind_start()
        self.assertListEqual(num_orb, [18, 18])
        self.assertListEqual(ind_start, [0, 18])

        # Abnormal case with 3 supercells and no inter-hopping
        sample = Sample(sc1, sc2, sc3)
        num_orb = sample._get_num_orb_sc()
        ind_start = sample._get_ind_start()
        self.assertListEqual(num_orb, [18, 18, 18])
        self.assertListEqual(ind_start, [0, 18, 36])

    def test_init_orb(self):
        """
        Test 'init_orb_eng' and 'init_orb_pos' methods of 'Sample' and
        'CachedSample' class.

        :return: None
        """
        sc1, sc2, sc3, inter_hop1, inter_hop2 = make_test_set()

        # Normal case with 1 supercell
        sample = Sample(sc1)
        sample.init_orb_eng()
        sample.init_orb_pos()
        self.assertTupleEqual(sample.orb_eng.shape, (18,))
        self.assertTupleEqual(sample.orb_pos.shape, (18, 3))
        self.assertTupleEqual(sample.hop_ind.shape, (27, 2))
        self.assertTupleEqual(sample.dr.shape, (27, 3))

        # Normal case with 2 supercells and 1 inter-hopping
        sample = Sample(sc1, sc2, inter_hop1)
        sample.init_orb_eng()
        sample.init_orb_pos()
        self.assertTupleEqual(sample.orb_eng.shape, (36,))
        self.assertTupleEqual(sample.orb_pos.shape, (36, 3))
        self.assertTupleEqual(sample.hop_ind.shape, (57, 2))
        self.assertTupleEqual(sample.dr.shape, (57, 3))

        # Normal case with 3 supercells and 2 inter-hopping
        sample = Sample(sc1, sc2, sc3, inter_hop1, inter_hop2)
        sample.init_orb_eng()
        sample.init_orb_pos()
        self.assertTupleEqual(sample.orb_eng.shape, (54,))
        self.assertTupleEqual(sample.orb_pos.shape, (54, 3))
        self.assertTupleEqual(sample.hop_ind.shape, (87, 2))
        self.assertTupleEqual(sample.dr.shape, (87, 3))

        # Abnormal case with 2 supercells and no inter-hopping
        sample = Sample(sc1, sc2)
        sample.init_orb_eng()
        sample.init_orb_pos()
        self.assertTupleEqual(sample.orb_eng.shape, (36,))
        self.assertTupleEqual(sample.orb_pos.shape, (36, 3))
        self.assertTupleEqual(sample.hop_ind.shape, (54, 2))
        self.assertTupleEqual(sample.dr.shape, (54, 3))

        # Abnormal case with 3 supercells and no inter-hopping
        sample = Sample(sc1, sc2, sc3)
        sample.init_orb_eng()
        sample.init_orb_pos()
        self.assertTupleEqual(sample.orb_eng.shape, (54,))
        self.assertTupleEqual(sample.orb_pos.shape, (54, 3))
        self.assertTupleEqual(sample.hop_ind.shape, (81, 2))
        self.assertTupleEqual(sample.dr.shape, (81, 3))

        # Cached sample
        os.system("rm -rf test")
        sample.save("test")
        sample = CachedSample("test")
        sample.init_orb_eng()
        sample.init_orb_pos()
        self.assertTupleEqual(sample.orb_eng.shape, (54,))
        self.assertTupleEqual(sample.orb_pos.shape, (54, 3))
        self.assertTupleEqual(sample.hop_ind.shape, (81, 2))
        self.assertTupleEqual(sample.dr.shape, (81, 3))
        os.system("rm -rf test")

    def test_init_hop(self):
        """
        Test 'init_hop' and 'init_dr' methods of 'Sample' and 'CachedSample'
        class.

        :return: None
        """
        sc1, sc2, sc3, inter_hop1, inter_hop2 = make_test_set()

        # Normal case with 1 supercell
        sample = Sample(sc1)
        sample.init_hop()
        self.assertTupleEqual(sample.orb_eng.shape, (18,))
        self.assertTupleEqual(sample.orb_pos.shape, (18, 3))
        self.assertTupleEqual(sample.hop_ind.shape, (27, 2))
        self.assertTupleEqual(sample.dr.shape, (27, 3))

        # Normal case with 2 supercells and 1 inter-hopping
        sample = Sample(sc1, sc2, inter_hop1)
        sample.init_hop()
        self.assertTupleEqual(sample.orb_eng.shape, (36,))
        self.assertTupleEqual(sample.orb_pos.shape, (36, 3))
        self.assertTupleEqual(sample.hop_ind.shape, (57, 2))
        self.assertTupleEqual(sample.dr.shape, (57, 3))

        # Normal case with 3 supercells and 2 inter-hopping
        sample = Sample(sc1, sc2, sc3, inter_hop1, inter_hop2)
        sample.init_hop()
        self.assertTupleEqual(sample.orb_eng.shape, (54,))
        self.assertTupleEqual(sample.orb_pos.shape, (54, 3))
        self.assertTupleEqual(sample.hop_ind.shape, (87, 2))
        self.assertTupleEqual(sample.dr.shape, (87, 3))

        # Abnormal case with 2 supercells and no inter-hopping
        sample = Sample(sc1, sc2)
        sample.init_hop()
        self.assertTupleEqual(sample.orb_eng.shape, (36,))
        self.assertTupleEqual(sample.orb_pos.shape, (36, 3))
        self.assertTupleEqual(sample.hop_ind.shape, (54, 2))
        self.assertTupleEqual(sample.dr.shape, (54, 3))

        # Abnormal case with 3 supercells and no inter-hopping
        sample = Sample(sc1, sc2, sc3)
        sample.init_hop()
        self.assertTupleEqual(sample.orb_eng.shape, (54,))
        self.assertTupleEqual(sample.orb_pos.shape, (54, 3))
        self.assertTupleEqual(sample.hop_ind.shape, (81, 2))
        self.assertTupleEqual(sample.dr.shape, (81, 3))

        # Cached sample
        os.system("rm -rf test")
        sample.save("test")
        sample = CachedSample("test")
        sample.init_hop()
        self.assertTupleEqual(sample.orb_eng.shape, (54,))
        self.assertTupleEqual(sample.orb_pos.shape, (54, 3))
        self.assertTupleEqual(sample.hop_ind.shape, (81, 2))
        self.assertTupleEqual(sample.dr.shape, (81, 3))
        os.system("rm -rf test")

    def test_save_load(self):
        """
        Test if models can be correctly saved and loaded.
        :return: None
        """
        th = TestHelper(self)
        sc1, sc2, sc3, inter_hop1, inter_hop2 = make_test_set()
        sample = Sample(sc1, sc2, sc3)

        os.system("rm -rf test")
        sample.save("test")
        cached_sample = CachedSample("test")
        self.assertEqual(sample.num_orb, cached_sample.num_orb)
        self.assertEqual(sample.num_hop, cached_sample.num_hop)
        th.test_equal_array(sample.sc0_lat_vec, cached_sample.sc0_lat_vec)
        os.system("rm -rf test")

    def test_speed(self):
        """
        Test the efficiency of time-consuming methods of 'Sample' class.

        :return: None.
        """
        timer = Timer()

        # Sample initialization
        timer.tic("sc_init")
        sc = SuperCell(make_cell(), dim=(500, 1000), pbc=(True, True))
        sample = Sample(sc)
        timer.toc("sc_init")

        # init_*
        timer.tic("init_hop")
        sample.init_hop()
        timer.toc("init_hop")
        timer.tic("init_orb_pos")
        sample.init_orb_pos()
        timer.toc("init_orb_pos")

        print()
        timer.report_total_time()

    def test_plot(self):
        """
        Test the 'plot' method of 'Sample' class.

        :return: None.
        """
        print("\n3x3 Graphene supercell with pbc")
        sc = SuperCell(make_cell(), dim=(3, 3, 1), pbc=(True, True, False))
        sample = Sample(sc)
        sample.plot()

        print("3x3 Graphene supercell with open bc")
        sc = SuperCell(make_cell(), dim=(3, 3, 1), pbc=(False, False, False))
        sample = Sample(sc)
        sample.plot()

        print("3x3 Graphene supercell with pbc and vacancies")
        vacancies = [(1, 1, 0, 0), (1, 1, 0, 1)]
        ref_map = RegularOrbitalMap((3, 3, 1), sc.num_orb_pc)
        vacancies = [ref_map.pc2sc(*v) for v in vacancies]
        sc = SuperCell(make_cell(), dim=(3, 3, 1), pbc=(True, True, False),
                       vacancies=vacancies)
        sample = Sample(sc)
        sample.plot()

        print("3x3 Graphene supercell with open bc and vacancies")
        vacancies = [(1, 1, 0, 0), (1, 1, 0, 1)]
        ref_map = RegularOrbitalMap((3, 3, 1), sc.num_orb_pc)
        vacancies = [ref_map.pc2sc(*v) for v in vacancies]
        sc = SuperCell(make_cell(), dim=(3, 3, 1), pbc=(False, False, False),
                       vacancies=vacancies)
        sample = Sample(sc)
        sample.plot()

        print("3x3 Graphene supercell with open bc and intra hopping")
        sc = SuperCell(make_cell(), dim=(3, 3, 1), pbc=(True, True, False))
        orb_map = sc.get_orb_map()
        id_pc_bra = np.array([(0, 0, 0, 1), (0, 2, 0, 0)], )
        id_pc_ket = np.array([(1, 2, 0, 0), (2, 0, 0, 1)])
        id_sc_bra = orb_map.pc2sc_array(id_pc_bra)
        id_sc_ket = orb_map.pc2sc_array(id_pc_ket)
        sc.add_hopping(rn=(0, 0, 0), orb_i=id_sc_bra.item(0), orb_j=id_sc_ket.item(0),
                       energy=1.5)
        sc.add_hopping(rn=(0, 0, 0), orb_i=id_sc_bra.item(1), orb_j=id_sc_ket.item(1),
                       energy=2.0)
        sample = Sample(sc)
        sample.plot()

        print("3x3 Graphene supercell with 2 layers")
        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))
        om1 = sc1.get_orb_map()
        om2 = sc2.get_orb_map()
        inter_hop = SCInterHopping(sc_bra=sc1, sc_ket=sc2)
        inter_hop.add_hopping(rn=(0, 0, 0), energy=-1.2,
                              orb_i=om1.pc2sc((0, 0, 0, 0)),
                              orb_j=om2.pc2sc((1, 2, 0, 1)))
        inter_hop.add_hopping(rn=(0, 0, 0), energy=1.5,
                              orb_i=om1.pc2sc((2, 1, 0, 1)),
                              orb_j=om2.pc2sc((1, 2, 0, 0)))
        inter_hop.add_hopping(rn=(0, 0, 0), energy=-0.7,
                              orb_i=om1.pc2sc((2, 0, 0, 0)),
                              orb_j=om2.pc2sc((0, 2, 0, 1)))
        sample = Sample(sc1, sc2, inter_hop)
        sample.plot(sc_hop_colors=["r", "b"], inter_hop_colors=["g"])

        print("User defined orbital colors")

        def orb_color1(orb_id_pc):
            scatter_color = ["r" if i % 2 == 0 else "b"
                             for i in range(len(orb_id_pc))]
            return scatter_color

        def orb_color2(orb_id_pc):
            scatter_color = ["m" if i % 2 == 0 else "g"
                             for i in range(len(orb_id_pc))]
            return scatter_color

        sample = Sample(sc1, sc2)
        sample.plot(sc_orb_colors=[orb_color1, orb_color2],
                    sc_hop_colors=["gray", "gray"])

        # Cached model
        os.system("rm -rf test")
        sample.save("test")
        sample = CachedSample("test")
        sample.plot()
        os.system("rm -rf test")

    def test_plot_advanced(self):
        """
        Test the 'plot' method of 'Sample' class with a complicated structure.

        :return: None.
        """
        sc = SuperCell(make_cell(), dim=(24, 24, 1))
        sample = Sample(sc)
        sample.init_orb_pos()
        positions = sample.orb_pos

        holes = []
        center = np.array([2.101, 1.361, 0.0])
        for i in range(sc.num_orb_sc):
            if np.linalg.norm(positions[i] - center) <= 0.5:
                holes.append(i)
        center = np.array([3.101, 3.361, 0.0])
        for i in range(sc.num_orb_sc):
            if np.linalg.norm(positions[i] - center) <= 0.5:
                holes.append(i)
        center = np.array([5.84, 3.51, 0.0])
        for i in range(sc.num_orb_sc):
            if np.linalg.norm(positions[i] - center) <= 0.5:
                holes.append(i)
        center = np.array([4.82, 1.11, 0.0])
        for i in range(sc.num_orb_sc):
            if np.linalg.norm(positions[i] - center) <= 0.5:
                holes.append(i)

        print("\n24*24 Graphene supercell with 4 holes without trimming"
              " dangling orbitals")
        sc = SuperCell(make_cell(), dim=(24, 24, 1), vacancies=holes)
        sample = Sample(sc)
        sample.plot(with_cells=False)

        print("\n24*24 Graphene supercell with 4 holes with dangling orbitals"
              " trimmed")
        sc.unlock()
        sc.trim()
        sample = Sample(sc)
        sample.plot(with_cells=False)

        print("\n24*24 Graphene supercell with Gaussian bump")

        def _make_pos_mod(c0):
            def _pos_mod(orb_pos):
                x, y = orb_pos[:, 0], orb_pos[:, 1]
                orb_pos[:, 2] += np.exp(-(x - c0[0])**2 - (y - c0[1])**2)
            return _pos_mod

        sc = SuperCell(make_cell(), dim=(24, 24, 1),
                       orb_pos_modifier=_make_pos_mod([4.17, 2.70]))
        sample = Sample(sc)
        sample.plot(with_cells=False, hop_as_arrows=False, view="ac")

        # Cached model
        os.system("rm -rf test")
        sample.save("test")
        sample = CachedSample("test")
        sample.plot(with_cells=False, hop_as_arrows=False, view="ac")
        os.system("rm -rf test")

    def test_pc_update(self):
        """
        Test data updating for primitive cell.

        :return: None
        """
        pc = tb.make_graphene_diamond()
        hash_old = hash(pc)

        # Add an orbital and compare
        pc.add_orbital((0.5, 0.5))
        hash_new = hash(pc)
        self.assertNotEqual(hash_old, hash_new)

        # Add a hopping term and compare
        hash_old = hash_new
        pc.add_hopping((0, 0), 0, 2, 1.0)
        hash_new = hash(pc)
        self.assertNotEqual(hash_old, hash_new)

    def test_sc_update(self):
        """
        Test data updating for supercell.

        :return: None
        """
        pc = tb.make_graphene_diamond()
        sc = tb.SuperCell(pc, dim=(3, 3, 1), pbc=(True, True, False))
        hash_old = hash(sc)

        # Modify the primitive cell
        pc.unlock()
        pc.add_orbital((0.5, 0.5))
        pc.add_hopping((0, 0), 0, 2, 1.0)
        hash_new = hash(sc)
        self.assertNotEqual(hash_old, hash_new)

        # Modify the vacancies
        hash_old = hash_new
        sc.add_vacancies([1])
        hash_new = hash(sc)
        self.assertNotEqual(hash_old, hash_new)

        # Add an hopping term
        sc.add_hopping((0, 0), 0, 9, 1.0)
        hash_new = hash(sc)
        self.assertNotEqual(hash_old, hash_new)

        # Set modifier
        def _test():
            pass
        hash_old = hash_new
        sc.set_orb_pos_modifier(_test)
        hash_new = hash(sc)
        self.assertNotEqual(hash_old, hash_new)

    def test_inter_hop_update(self):
        """
        Test data updating for supercell.

        :return: None
        """
        pc1 = tb.make_graphene_diamond()
        sc1 = tb.SuperCell(pc1, dim=(3, 3, 1), pbc=(True, True, False))
        pc2 = tb.make_graphene_diamond()
        sc2 = tb.SuperCell(pc2, dim=(3, 3, 1), pbc=(True, True, False))
        inter_hop = SCInterHopping(sc1, sc2)
        hash_old = hash(inter_hop)

        # Change primitive cell
        pc1.unlock()
        pc1.add_orbital((0.5, 0.5))
        hash_new = hash(inter_hop)
        self.assertNotEqual(hash_old, hash_new)
        pc2.unlock()
        pc1.add_orbital((0.5, 0.5))
        hash_new = hash(inter_hop)
        self.assertNotEqual(hash_old, hash_new)

    def test_sample_update_top_down(self):
        """
        Test data updating for sample.

        :return: None
        """
        prim_cell = tb.make_graphene_diamond()
        super_cell = tb.SuperCell(prim_cell, dim=(3, 3, 1),
                                  pbc=(True, True, False))
        sample = tb.Sample(super_cell)
        sample.init_hop()

        # Change primitive cell and update in top-down approach
        prim_cell.unlock()
        prim_cell.add_orbital((0.5, 0.5))
        sample.reset_array()
        self.assertEqual(prim_cell.num_orb, 3)
        self.assertEqual(prim_cell.orb_eng.shape[0], 3)
        self.assertEqual(prim_cell.orb_pos.shape[0], 3)
        self.assertEqual(super_cell.num_orb_pc, 3)
        self.assertEqual(super_cell.num_orb_sc, 27)
        self.assertEqual(super_cell._cached_data.orb_id_pc.shape[0], 27)
        self.assertEqual(sample.num_orb, 27)

        # Change supercell and update in top-down approach
        prim_cell.unlock()
        prim_cell.remove_orbitals([2])
        super_cell.unlock()
        super_cell.add_vacancies([1])
        sample.reset_array()
        self.assertEqual(prim_cell.num_orb, 2)
        self.assertEqual(prim_cell.orb_eng.shape[0], 2)
        self.assertEqual(prim_cell.orb_pos.shape[0], 2)
        self.assertEqual(super_cell.num_orb_pc, 2)
        self.assertEqual(super_cell.num_orb_sc, 17)
        self.assertEqual(super_cell._cached_data.orb_id_pc.shape[0], 17)
        self.assertEqual(sample.num_orb, 17)

        # Change supercell and update in top-down approach
        super_cell.unlock()
        super_cell.add_vacancies([2])
        sample.reset_array()
        self.assertEqual(prim_cell.num_orb, 2)
        self.assertEqual(prim_cell.orb_eng.shape[0], 2)
        self.assertEqual(prim_cell.orb_pos.shape[0], 2)
        self.assertEqual(super_cell.num_orb_pc, 2)
        self.assertEqual(super_cell.num_orb_sc, 16)
        self.assertEqual(super_cell._cached_data.orb_id_pc.shape[0], 16)
        self.assertEqual(sample.num_orb, 16)

    def test_sample_update_bottom_up(self):
        """
        Test data updating for sample.

        :return: None
        """
        prim_cell = tb.make_graphene_diamond()
        super_cell = tb.SuperCell(prim_cell, dim=(3, 3, 1),
                                  pbc=(True, True, False))
        sample = tb.Sample(super_cell)
        sample.init_hop()

        # Change primitive cell and update in top-down approach
        prim_cell.unlock()
        prim_cell.add_orbital((0.5, 0.5))
        prim_cell.update()
        self.assertEqual(prim_cell.num_orb, 3)
        self.assertEqual(prim_cell.orb_eng.shape[0], 3)
        self.assertEqual(prim_cell.orb_pos.shape[0], 3)
        self.assertEqual(super_cell.num_orb_pc, 3)
        self.assertEqual(super_cell.num_orb_sc, 27)
        self.assertEqual(super_cell._cached_data.orb_id_pc.shape[0], 27)
        self.assertEqual(sample.num_orb, 27)

        # Change supercell and update in top-down approach
        prim_cell.unlock()
        prim_cell.remove_orbitals([2])
        super_cell.unlock()
        super_cell.add_vacancies([1])
        prim_cell.update()
        self.assertEqual(prim_cell.num_orb, 2)
        self.assertEqual(prim_cell.orb_eng.shape[0], 2)
        self.assertEqual(prim_cell.orb_pos.shape[0], 2)
        self.assertEqual(super_cell.num_orb_pc, 2)
        self.assertEqual(super_cell.num_orb_sc, 17)
        self.assertEqual(super_cell._cached_data.orb_id_pc.shape[0], 17)
        self.assertEqual(sample.num_orb, 17)

        # Change supercell and update in top-down approach
        super_cell.unlock()
        super_cell.add_vacancies([2])
        super_cell.update()
        self.assertEqual(prim_cell.num_orb, 2)
        self.assertEqual(prim_cell.orb_eng.shape[0], 2)
        self.assertEqual(prim_cell.orb_pos.shape[0], 2)
        self.assertEqual(super_cell.num_orb_pc, 2)
        self.assertEqual(super_cell.num_orb_sc, 16)
        self.assertEqual(super_cell._cached_data.orb_id_pc.shape[0], 16)
        self.assertEqual(sample.num_orb, 16)


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