#! /usr/bin/env python
"""
This testsuite tests mostly the assertions. Functionality tests are in other
tests and examples, e.g., test_sk_soc.py and test_hetero.py.
"""
import unittest

import numpy as np
import matplotlib.pyplot as plt

import tbplas as tb
from tbplas.base.utils import TestHelper


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


def make_cell_rect():
    cell = tb.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_extend_prim_cell(self):
        """
        Test function 'extend_prim_cell'.

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

        # Exception handling
        def _test():
            tb.extend_prim_cell(cell, dim=(0, 0, 1))
        th.test_raise(_test, AssertionError,
                      r"Dimension along direction [0-9]"
                      r" should not be smaller than 1")

        # Plot
        extend_cell = tb.extend_prim_cell(cell, dim=(3, 3, 1))
        extend_cell.plot()

        # Band structure
        extend_cell = tb.extend_prim_cell(cell, dim=(5, 5, 1))
        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 = tb.gen_kpath(k_points, [40, 40, 40])
        k_len, bands = tb.calc_bands(extend_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()

        # DOS
        k_points = tb.gen_kmesh((24, 24, 1))
        energies, dos = tb.calc_dos(extend_cell, k_points, e_min=-10, e_max=10)
        plt.plot(energies, dos)
        plt.show()

        # Orbital labels
        self.assertSetEqual(set([orb.label for orb in extend_cell.orbitals]),
                            {"C_pz"})

        # Void cases
        cell = make_cell_orb_only()
        extend_cell = tb.extend_prim_cell(cell, dim=(3, 3, 1))
        extend_cell.plot()

    def test_reshape_prim_cell(self):
        """
        Test function 'reshape_prim_cell'.

        :return: None
        """
        print("\nReference rectangular cell:")
        sqrt3 = 1.73205080757
        a = 2.46
        cc_bond = sqrt3 / 3 * a
        vectors = tb.gen_lattice_vectors(sqrt3 * cc_bond, 3 * cc_bond)
        cell = tb.PrimitiveCell(vectors)
        cell.add_orbital((0, 0))
        cell.add_orbital((0, 2./3))
        cell.add_orbital((1./2, 1./6))
        cell.add_orbital((1./2, 1./2))
        cell.add_hopping((0, 0), 0, 2, -2.7)
        cell.add_hopping((0, 0), 2, 3, -2.7)
        cell.add_hopping((0, 0), 3, 1, -2.7)
        cell.add_hopping((0, 1), 1, 0, -2.7)
        cell.add_hopping((1, 0), 3, 1, -2.7)
        cell.add_hopping((1, 0), 2, 0, -2.7)
        cell.plot()

        print("\nTest rectangular cell:")
        cell = make_cell()
        lat_frac = np.array([[1, 0, 0], [-1, 2, 0], [0, 0, 1]])
        cell = tb.reshape_prim_cell(cell, lat_frac)
        self.assertEqual(cell.num_orb, 4)
        cell.plot()

        # Check orbital labels
        self.assertSetEqual(set([orb.label for orb in cell.orbitals]), {"C_pz"})

        # Void cases
        cell = make_cell_orb_only()
        lat_frac = np.array([[1, 0, 0], [-1, 2, 0], [0, 0, 1]])
        cell = tb.reshape_prim_cell(cell, lat_frac)
        cell.plot()
        cell = make_cell_empty()
        cell = tb.reshape_prim_cell(cell, lat_frac)
        cell.plot()

        # Assertion
        def _test():
            _cell = make_cell()
            _lat_frac = np.array([[1, 0], [-1, 2]])
            tb.reshape_prim_cell(_cell, _lat_frac)
        th = tb.TestHelper(self)
        th.test_raise(_test, AssertionError, "Lat_vec")

        def _test():
            _cell = make_cell()
            _lat_frac = np.zeros((3, 3))
            tb.reshape_prim_cell(_cell, _lat_frac)
        th.test_raise(_test, AssertionError, ".*")

    def test_spiral_prim_cel(self):
        cell = make_cell_rect()
        cell.plot()
        tb.spiral_prim_cell(cell, np.pi * 0.25)
        cell.plot()

    def test_merge_prim_cell(self):
        th = tb.TestHelper(self)

        def _test():
            tb.merge_prim_cell()
        th.test_raise(_test, AssertionError, "No components to merge")

        def _test():
            tb.merge_prim_cell(1, 1)
        th.test_raise(_test, AssertionError, "Component")

        def _test():
            tb.merge_prim_cell(make_cell(), make_cell_rect())
        th.test_raise(_test, AssertionError, "Lattice vectors")

        def _test():
            cell1 = make_cell()
            cell2 = make_cell()
            cell3 = make_cell()
            inter_hop = tb.PCInterHopping(cell2, cell3)
            tb.merge_prim_cell(cell1, cell3, inter_hop)
        th.test_raise(_test, AssertionError, "pc_bra of inter_hop")

        def _test():
            cell1 = make_cell()
            cell2 = make_cell()
            cell3 = make_cell()
            inter_hop = tb.PCInterHopping(cell2, cell3)
            tb.merge_prim_cell(cell1, cell2, inter_hop)
        th.test_raise(_test, AssertionError, "pc_ket of inter_hop")

    def test_find_neighbours(self):
        th = tb.TestHelper(self)

        def _test():
            cell1 = make_cell_empty()
            cell2 = make_cell()
            tb.find_neighbors(cell1, cell2)
        th.test_raise(_test, AssertionError, "cell_bra should not be empty")

        def _test():
            cell1 = make_cell()
            cell2 = make_cell_empty()
            tb.find_neighbors(cell1, cell2)
        th.test_raise(_test, AssertionError, "cell_ket should not be empty")


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