#! /usr/bin/env python

import unittest

import numpy as np

from tbplas import (gen_lattice_vectors, PrimitiveCell, Timer, TestHelper,
                    SuperCell, RegularOrbitalMap)
from tbplas.cython import orbital_map as core


class MyTest(unittest.TestCase):

    def setUp(self) -> None:
        vectors = gen_lattice_vectors(a=2.46, b=2.46, gamma=60)
        self.cell = PrimitiveCell(vectors)
        self.cell.add_orbital((0.0, 0.0), 0.0)
        self.cell.add_orbital((1./3, 1./3), 0.0)
        self.cell.add_hopping((0, 0), 0, 1, -2.7)
        self.cell.add_hopping((1, 0), 1, 0, -2.7)
        self.cell.add_hopping((0, 1), 1, 0, -2.7)
        vac = [(0, 0, 0, 1), (4, 6, 0, 0), (2, 0, 0, 0), (3, 7, 0, 0),
               (0, 6, 0, 0), (9, 0, 0, 0), (9, 9, 0, 1), (5, 7, 0, 1)]
        self.vac = vac[:8]
        self.dim_small = (17, 29, 1)
        self.dim_large = (5000, 1000, 1)
        self.timer = Timer()

    def tearDown(self) -> None:
        self.timer.report_time()

    def test_acc_sps(self):
        """
        Test the accuracy of core functions id_sc -> id_pc -> id_sc.

        :return: None
        """
        sc = SuperCell(self.cell, dim=self.dim_small)
        orb_map = sc.get_orb_map()
        record = "acc_sps"
        self.timer.tic(record)
        result = core.test_acc_sps(orb_map._dim,
                                   orb_map._num_orb_pc,
                                   orb_map._orb_id_pc)
        self.timer.toc(record)
        self.assertEqual(result, 0)

    def test_acc_psp(self):
        """
        Test the accuracy of core functions id_pc -> id_sc -> id_pc.

        :return: None
        """
        sc = SuperCell(self.cell, dim=self.dim_small)
        orb_map = sc.get_orb_map()
        record = "acc_psp"
        self.timer.tic(record)
        result = core.test_acc_psp(orb_map._dim,
                                   orb_map._num_orb_pc,
                                   orb_map._orb_id_pc)
        self.timer.toc(record)
        self.assertEqual(result, 0)

    def test_acc_sps_vac(self):
        """
        Test the accuracy of core functions id_sc -> id_pc -> id_sc
        in presence of vacancies.

        :return: None
        """
        sc = SuperCell(self.cell, dim=self.dim_small)
        ref_map = RegularOrbitalMap(self.dim_small, self.cell.num_orb)
        vac = [ref_map.pc2sc(*v) for v in self.vac]
        sc.set_vacancies(vac)
        sc._sync_array()
        orb_map = sc.get_orb_map()

        record = "acc_sps_vac"
        self.timer.tic(record)
        result = core.test_acc_sps_vac(orb_map._dim,
                                       orb_map._num_orb_pc,
                                       orb_map._orb_id_pc,
                                       orb_map._vac_id_sc)
        self.timer.toc(record)
        self.assertEqual(result, 0)

    def test_acc_psp_vac(self):
        """
        Test the accuracy of core functions id_pc -> id_sc -> id_pc
        in presence of vacancies.

        :return: None
        """
        sc = SuperCell(self.cell, dim=self.dim_small)
        ref_map = RegularOrbitalMap(self.dim_small, self.cell.num_orb)
        vac = [ref_map.pc2sc(*v) for v in self.vac]
        sc.set_vacancies(vac)
        sc._sync_array()
        orb_map = sc.get_orb_map()

        record = "acc_psp_vac"
        self.timer.tic(record)
        result = core.test_acc_psp_vac(orb_map._dim,
                                       orb_map._num_orb_pc,
                                       orb_map._orb_id_pc,
                                       orb_map._vac_id_sc)
        self.timer.toc(record)
        self.assertEqual(result, 0)

    def test_speed_pc2sc(self):
        """
        Test the efficiency of core function id_pc2sc.

        :return: None
        """
        sc = SuperCell(self.cell, dim=self.dim_large)
        orb_map = sc.get_orb_map()
        record = "speed_pc2sc"
        self.timer.tic(record)
        core.test_speed_pc2sc(orb_map._dim,
                              orb_map._num_orb_pc,
                              orb_map._orb_id_pc)
        self.timer.toc(record)

    def test_speed_pc2sc_vac(self):
        """
        Test the efficiency of core function id_pc2sc in presence of
        vacancies.

        :return: None
        """
        sc = SuperCell(self.cell, dim=self.dim_large)
        ref_map = RegularOrbitalMap(self.dim_large, self.cell.num_orb)
        vac = [ref_map.pc2sc(*v) for v in self.vac]
        sc.set_vacancies(vac)
        sc._sync_array()
        orb_map = sc.get_orb_map()

        record = "speed_pc2sc_vac"
        self.timer.tic(record)
        core.test_speed_pc2sc_vac(orb_map._dim,
                                  orb_map._num_orb_pc,
                                  orb_map._orb_id_pc,
                                  orb_map._vac_id_sc)
        self.timer.toc(record)

    def test_speed_sc2pc(self):
        """
        Test the efficiency of core function id_sc2pc.

        :return: None
        """
        sc = SuperCell(self.cell, dim=self.dim_large)
        orb_map = sc.get_orb_map()
        record = "speed_sc2pc"
        self.timer.tic(record)
        core.test_speed_sc2pc(orb_map._orb_id_pc)
        self.timer.toc(record)

    def test_acc_py(self):
        """
        Test the accuracy of Python interfaces to core functions.

        :return: None
        """
        sc = SuperCell(self.cell, dim=self.dim_small)
        orb_map = sc.get_orb_map()

        result = 0
        record = "acc_sps_py"
        self.timer.tic(record)
        for i in range(orb_map._num_orb_pc):
            pc = orb_map.sc2pc(i)
            sc = orb_map.pc2sc(pc)
            result += abs(i-sc)
        self.timer.toc(record)
        self.assertEqual(result, 0)

        result = 0
        record = "acc_psp_py"
        self.timer.tic(record)
        for i in orb_map._orb_id_pc:
            sc = orb_map.pc2sc(i)
            pc = orb_map.sc2pc(sc)
            for j in range(4):
                result += abs(i.item(j)-pc[j])
        self.timer.toc(record)
        self.assertEqual(result, 0)

    def test_acc_vac_py(self):
        """
        Test the accuracy of Python interfaces to core functions in presence
        of vacancies.

        :return: None
        """
        sc = SuperCell(self.cell, dim=self.dim_small)
        ref_map = RegularOrbitalMap(self.dim_small, self.cell.num_orb)
        vac = [ref_map.pc2sc(*v) for v in self.vac]
        sc.set_vacancies(vac)
        sc._sync_array()
        orb_map = sc.get_orb_map()

        result = 0
        record = "acc_sps_vac_py"
        self.timer.tic(record)
        for i in range(sc.num_orb_sc):
            pc = orb_map.sc2pc(i)
            sc = orb_map.pc2sc(pc)
            result += abs(i-sc)
        self.timer.toc(record)
        self.assertEqual(result, 0)

        result = 0
        record = "acc_psp_vac_py"
        self.timer.tic(record)
        for i in orb_map._orb_id_pc:
            sc = orb_map.pc2sc(i)
            pc = orb_map.sc2pc(sc)
            for j in range(4):
                result += abs(i.item(j)-pc[j])
        self.timer.toc(record)
        self.assertEqual(result, 0)

    def test_acc_pc2sc_array(self):
        """
        Test the accuracy of Python interface pc2sc_array.

        :return: None
        """
        sc = SuperCell(self.cell, dim=self.dim_small)
        orb_map = sc.get_orb_map()
        th = TestHelper(self)

        record = "acc_pc2sc_array"
        self.timer.tic(record)
        sc_test = orb_map.pc2sc_array(orb_map._orb_id_pc)
        sc_ref = np.linspace(0, sc.num_orb_sc-1, sc.num_orb_sc,
                             dtype=np.int64)
        self.timer.toc(record)
        th.test_equal_array(sc_test, sc_ref)

        record = "acc_pc2sc_array_vac"
        sc = SuperCell(self.cell, dim=self.dim_small)
        ref_map = RegularOrbitalMap(self.dim_small, self.cell.num_orb)
        vac = [ref_map.pc2sc(*v) for v in self.vac]
        sc.set_vacancies(vac)
        sc._sync_array()
        orb_map = sc.get_orb_map()
        self.timer.tic(record)
        sc_test = orb_map.pc2sc_array(orb_map._orb_id_pc)
        sc_ref = np.linspace(0, sc.num_orb_sc-1, sc.num_orb_sc,
                             dtype=np.int64)
        self.timer.toc(record)
        th.test_equal_array(sc_test, sc_ref)

    def test_acc_sc2pc_array(self):
        """
        Test the accuracy of Python interface sc2pc_array.

        :return: None
        """
        sc = SuperCell(self.cell, dim=self.dim_small)
        orb_map = sc.get_orb_map()
        th = TestHelper(self)

        record = "acc_sc2pc_array"
        sc = np.linspace(0, sc.num_orb_sc-1, sc.num_orb_sc, dtype=np.int64)
        self.timer.tic(record)
        pc = orb_map.sc2pc_array(sc)
        self.timer.toc(record)
        th.test_equal_array(pc, orb_map._orb_id_pc)

        record = "acc_sc2pc_array_vac"
        sc = SuperCell(self.cell, dim=self.dim_small)
        ref_map = RegularOrbitalMap(self.dim_small, self.cell.num_orb)
        vac = [ref_map.pc2sc(*v) for v in self.vac]
        sc.set_vacancies(vac)
        sc._sync_array()
        orb_map = sc.get_orb_map()
        sc = np.linspace(0, sc.num_orb_sc-1, sc.num_orb_sc, dtype=np.int64)
        self.timer.tic(record)
        pc = orb_map.sc2pc_array(sc)
        self.timer.toc(record)
        th.test_equal_array(pc, orb_map._orb_id_pc)

    def test_speed_pc2sc_array(self):
        """
        Test the efficiency of Python interface pc2sc_array.

        :return: None
        """
        sc = SuperCell(self.cell, dim=self.dim_large)
        orb_map = sc.get_orb_map()

        record = "speed_pc2sc_array"
        self.timer.tic(record)
        orb_map.pc2sc_array(orb_map._orb_id_pc)
        self.timer.toc(record)

        record = "speed_pc2sc_array_vac"
        sc = SuperCell(self.cell, dim=self.dim_large)
        ref_map = RegularOrbitalMap(self.dim_large, self.cell.num_orb)
        vac = [ref_map.pc2sc(*v) for v in self.vac]
        sc.set_vacancies(vac)
        sc._sync_array()
        orb_map = sc.get_orb_map()
        self.timer.tic(record)
        orb_map.pc2sc_array(orb_map._orb_id_pc)
        self.timer.toc(record)

    def test_speed_sc2pc_array(self):
        """
        Test the efficiency of Python interface sc2pc_array.

        :return: None
        """
        sc = SuperCell(self.cell, dim=self.dim_large)
        orb_map = sc.get_orb_map()

        record = "speed_sc2pc_array"
        sc = np.linspace(0, sc.num_orb_sc-1, sc.num_orb_sc, dtype=np.int64)
        self.timer.tic(record)
        orb_map.sc2pc_array(sc)
        self.timer.toc(record)

        record = "speed_sc2pc_array_vac"
        sc = SuperCell(self.cell, dim=self.dim_large)
        ref_map = RegularOrbitalMap(self.dim_large, self.cell.num_orb)
        vac = [ref_map.pc2sc(*v) for v in self.vac]
        sc.set_vacancies(vac)
        sc._sync_array()
        orb_map = sc.get_orb_map()
        sc = np.linspace(0, sc.num_orb_sc-1, sc.num_orb_sc, dtype=np.int64)
        self.timer.tic(record)
        orb_map.sc2pc_array(sc)
        self.timer.toc(record)


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