#! /usr/bin/env python
"""
Example for calculating response functions using the 'Lindhard' class.

References:
[1] https://journals.aps.org/prb/abstract/10.1103/PhysRevB.84.035439
[2] https://journals.aps.org/prb/abstract/10.1103/PhysRevB.82.115448
"""

import math

import numpy as np
import matplotlib.pyplot as plt

import tbplas as tb


def plot_dyn_pol(omegas: np.ndarray,
                 dyn_pol: np.ndarray,
                 iq: int = 0,
                 t: float = 3.0,
                 a: float = 0.142) -> None:
    """
    Plot dynamic polarizability.
    :param omegas: (num_eng,) float64 array, energy grid in eV
    :param dyn_pol: (num_qpt, num_eng) complex128 array, dynamic polarizability
    :param iq: index of q-point of dyn_pol to plot
    :param t: hopping integral in eV
    :param a: C-C bond length in NM
    :return: None
    """
    # Reference: fig. 3(e) of [1] for |q| = 1 / a and theta = 30 degrees
    plt.plot(omegas / t, -dyn_pol[iq].imag * t * a ** 2, color="red")
    plt.xticks(np.linspace(0.0, 6.0, 7))
    plt.yticks(np.linspace(0.0, 1.4, 8))
    plt.xlim(0.0, 6.0)
    plt.ylim(0.0, 1.4)
    plt.minorticks_on()
    plt.tick_params(which="major", direction="in", length=8, width=0.6)
    plt.tick_params(which="minor", direction="in", length=4, width=0.6)
    plt.xlabel(r"$\omega\ (t)$")
    plt.ylabel(r"-Im$\Pi\ (t^{-1}a^{-2})$")
    plt.tight_layout()
    plt.savefig("dyn_pol.pdf")
    plt.close()


def plot_epsilon(omegas: np.ndarray,
                 epsilon: np.ndarray,
                 iq: int = 0) -> None:
    """
    Plot dynamic polarization.
    :param omegas: (num_eng,) float64 array, energy grid in eV
    :param epsilon: (num_qpt, num_eng) complex128 array, dielectric function
    :param iq: index of q-point of epsilon to plot
    :return: None
    """
    # Reference: fig. 7(c) of [1] for |q| = 4.76 nm^-1 and theta = 30 degrees
    plt.plot(omegas, epsilon[iq].real, color="red")
    plt.xticks(np.linspace(0.0, 18.0, 10))
    plt.yticks(np.linspace(0.0, 15.0, 4))
    plt.xlim(0.0, 18.0)
    plt.ylim(-1.5, 15)
    plt.minorticks_on()
    plt.tick_params(which="major", direction="in", length=8, width=0.6)
    plt.tick_params(which="minor", direction="in", length=4, width=0.6)
    plt.axhline(y=0.0, xmin=0.0, xmax=18.0, linestyle="--", color="k", linewidth=1.0)
    plt.xlabel(r"$\omega\ (eV)$")
    plt.ylabel(r"Re $\epsilon$")
    plt.tight_layout()
    plt.savefig("epsilon.pdf")
    plt.close()


def calc_dyn_pol_epsilon() -> None:
    """
    Calculate dynamic polarizability and dielectric function of monolayer graphene.
    """
    # Model parameters
    # t: hopping integral in eV
    # a: C-C bond length in NM, NOT lattice constant (0.246 NM in make_graphene_diamond)
    t = 3.0 
    a = 0.142

    # q-points:
    # q0: |q| = 1 / a and theta = 30 degrees
    # q1: |q| = 4.76 nm^-1 and theta = 30 degrees
    theta = 30.0 / 180.0 * math.pi
    unit_vec = np.array([math.cos(theta), math.sin(theta), 0.0])
    q0 = 1.0 / a * unit_vec
    q1 = 4.76 * unit_vec

    # Create model, solver and set parameters
    cell = tb.make_graphene_diamond(t=t)
    lind = tb.Lindhard(cell)
    lind.config.prefix = "dyn_pol_epsilon"
    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.array([q0, q1], dtype=np.float64)

    # Calculate dyn_pol and epsilon
    from_scratch = True
    if from_scratch:
        omegas, dyn_pol = lind.calc_dyn_pol()
    else:
        omegas, dyn_pol = tb.load_dyn_pol(prefix=lind.config.prefix)
    epsilon = lind.calc_epsilon(dyn_pol)

    # Plot
    plot_dyn_pol(omegas, dyn_pol, iq=0, t=t, a=a)
    plot_epsilon(omegas, epsilon, iq=1)


def plot_ac_cond(omegas: np.ndarray,
                 ac_cond: np.ndarray,
                 t: float = 3.0) -> None:
    """
    Plot ac conductivity.
    :param omegas: (num_eng,) float64 array, energy grid in eV
    :param ac_cond: (num_eng,) complex128 array, ac conductivity
    :param t: hopping integral in eV
    :return: None
    """
    # Reference: fig. 13 of [2]
    plt.plot(omegas / t, ac_cond.real * 4, color="red")
    plt.xticks(np.linspace(0.0, 6.0, 7))
    plt.yticks(np.linspace(0.0, 4.0, 9))
    plt.xlim(0.0, 6.0)
    plt.ylim(0.0, 4.0)
    plt.minorticks_on()
    plt.tick_params(which="major", direction="out", length=8, width=0.6)
    plt.tick_params(which="minor", direction="out", length=4, width=0.6)
    plt.xlabel(r"$\omega\ (t)$")
    plt.ylabel(r"$\sigma_{xx}\ (\sigma_0)$")
    plt.tight_layout()
    plt.savefig("ac_cond.pdf")
    plt.close()


def calc_ac_cond() -> None:
    """
    Calculate ac conductivity of monolayer graphene.
    """
    # Model parameters
    # t: hopping integral in eV
    t = 3.0 

    # Create model, solver and set parameters
    cell = tb.make_graphene_diamond(t=t)
    lind = tb.Lindhard(cell)
    lind.config.prefix = "ac_cond"
    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)
    from_scratch = True
    if from_scratch:
        omegas, ac_cond = lind.calc_ac_cond()
    else:
        omegas, ac_cond = tb.load_ac_cond(prefix=lind.config.prefix)

    # Plot
    plot_ac_cond(omegas, ac_cond, t=t)


if __name__ == "__main__":
    calc_dyn_pol_epsilon()
    calc_ac_cond()
