#! /usr/bin/env python

import unittest

import numpy as np
import matplotlib.pyplot as plt

import tbplas as tb
from tbplas import TestHelper


class MyTest(unittest.TestCase):

    def test_cart_frac_conversion(self):
        """
        Test frac2cart and cart2frac.

        :return: None
        """
        lattice = tb.gen_lattice_vectors(a=10, b=8, c=6)
        prim_cell = tb.PrimitiveCell(lattice, unit=tb.ANG)
        lindhard = tb.Lindhard(prim_cell)

        th = TestHelper(self)
        q_frac = np.array([[0.15, -0.13, 0.35]])
        q_cart_ang_ref = 2 * np.pi * np.array([[0.15/10, -0.13/8, 0.35/6]])
        q_cart_nm_ref = q_cart_ang_ref * 10
        q_cart_ang_test = lindhard.k_frac2cart(q_frac, unit=tb.ANG)
        q_cart_nm_test = lindhard.k_frac2cart(q_frac, unit=tb.NM)
        th.test_equal_array(q_cart_ang_test, q_cart_ang_ref, almost=True)
        th.test_equal_array(q_cart_nm_test, q_cart_nm_ref, almost=True)

        q_cart_nm = np.array([[2.3, -1.5, 0.7]])
        q_frac_ref = 1 / (2 * np.pi) * np.array([[2.3, -1.5*0.8, 0.7*0.6]])
        q_frac_test = lindhard.k_cart2frac(q_cart_nm, unit=tb.NM)
        th.test_equal_array(q_frac_test, q_frac_ref, almost=True)

        q_cart_ang = np.array([[2.3, -1.5, 0.7]])
        q_frac_ref = 1 / (2 * np.pi) * np.array([2.3*10, -1.5*8, 0.7*6])
        q_frac_test = lindhard.k_cart2frac(q_cart_ang, unit=tb.ANG)
        th.test_equal_array(q_frac_test, q_frac_ref, almost=True)

    def test_dyn_pol_prb(self):
        """
        Reproducing Phys. Rev. B 84, 035439 (2011) with |q| = 1/a and theta = 30
        degrees. See examples/prim_cell/lindhard.py for more details.

        :return: None
        """
        # Construct primitive cell
        t = 3.0  # Absolute hopping energy in eV
        a = 0.142  # C-C distance in NM
        vectors = tb.gen_lattice_vectors(a=0.246, b=0.246, c=1.0, gamma=60)
        cell = tb.PrimitiveCell(vectors, unit=tb.NM)
        cell.add_orbital((0.0, 0.0), label="C_pz")
        cell.add_orbital((1 / 3., 1 / 3.), label="C_pz")
        cell.add_hopping((0, 0), 0, 1, t)
        cell.add_hopping((1, 0), 1, 0, t)
        cell.add_hopping((0, 1), 1, 0, t)

        # Set parameter for Lindhard function
        lind = tb.Lindhard(cell)
        lind.config.prefix = "graphene"
        lind.config.e_min = 0.0
        lind.config.e_max = 18.0
        lind.config.e_step = 0.005
        lind.config.dimension = 2
        lind.config.k_grid_size = (1200, 1200, 1)
        lind.config.q_points = np.zeros((1, 3), dtype=np.float64)
        lind.config.q_points[0] = np.array([[0.86602540, 0.5, 0.0]]) / a

        # Calculate and plot dyn_pol
        omegas, dyn_pol = lind.calc_dyn_pol()
        plt.plot(omegas/t, -dyn_pol.imag[0]*t*a**2)
        plt.grid()
        plt.savefig("dyn_pol.png")
        plt.close()

    def test_epsilon_prb(self):
        """
        Reproducing Phys. Rev. B 84, 035439 (2011) with |q| = 0.476 / Angstrom
        and theta = 30 degrees. See examples/prim_cell/lindhard.py for more details.

        :return: None
        """
        # Construct primitive cell
        t = 3.0  # Absolute hopping energy in eV
        vectors = tb.gen_lattice_vectors(a=0.246, b=0.246, c=1.0, gamma=60)
        cell = tb.PrimitiveCell(vectors, unit=tb.NM)
        cell.add_orbital((0.0, 0.0), label="C_pz")
        cell.add_orbital((1 / 3., 1 / 3.), label="C_pz")
        cell.add_hopping((0, 0), 0, 1, t)
        cell.add_hopping((1, 0), 1, 0, t)
        cell.add_hopping((0, 1), 1, 0, t)

        # Set parameter for Lindhard function
        lind = tb.Lindhard(cell)
        lind.config.prefix = "graphene"
        lind.config.e_min = 0.0
        lind.config.e_max = 18.0
        lind.config.e_step = 0.005
        lind.config.dimension = 2
        lind.config.k_grid_size = (1200, 1200, 1)
        lind.config.q_points = np.zeros((1, 3), dtype=np.float64)
        lind.config.q_points[0] = np.array([[4.122280922013927, 2.38, 0.0]])

        # Evaluate dielectric function
        omegas, dyn_pol = lind.calc_dyn_pol()
        epsilon = lind.calc_epsilon(dyn_pol)

        # Plot
        plt.plot(omegas, epsilon[0].real, color="r")
        plt.minorticks_on()
        plt.grid()
        plt.savefig("epsilon.png")
        plt.close()

    def test_ac_cond(self):
        """
        Calculate the AC conductivity of monolayer graphene.
        See examples/prim_cell/lindhard.py for more details.

        :return: None
        """
        t = 3.0
        model = tb.make_graphene_diamond(t=t)

        lind = tb.Lindhard(model)
        lind.config.prefix = "graphene"
        lind.config.e_min = 0.0
        lind.config.e_max = 18.0
        lind.config.e_step = 0.005
        lind.config.dimension = 2
        lind.config.num_spin = 2
        lind.config.k_grid_size = (2048, 2048, 1)

        omegas, ac_cond = lind.calc_ac_cond()

        plt.plot(omegas/t, ac_cond.real*4, color="r")
        plt.minorticks_on()
        plt.grid()
        plt.savefig("sigma_xx.png")
        plt.close()


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