"""
Test Python interface for DIAG algorithms.

This module generates the same output files (*.dat or *.h5) as the C++ version, which can be
plotted by 'plot_tbpm.py' to produce *_cpp.dat. Meanwhile, it also calls Python Visualizer
to plot the results on-the-fly.
"""
import numpy as np

import tbplas as tb

from plot_diag import plot_sigma_band, plot_sigma_eng, plot_omega_xy
from model import make_graphene_pc, make_overlap, make_haldane
from utils import save_xy


__all__ = ["test_diag_bands", "test_diag_dos", "test_diag_ldos",
           "test_diag_states", "test_spin", "test_z2", "test_berry",
           "test_diag_dyn_pol", "test_diag_ac_cond"]


# Algorithm for diagonalization, should be in eigen or lapacke
DIAG_ALGO = "eigen"

# Algorithm for post-processing, currently only supports cpu
PP_ALGO = "cpu"

# Convention of Hamiltonian, should be 1 or 2
CONVENTION = 1

# Whether to test algorithms with overlaps (zhegv)
WITH_OVERLAP = False


def test_diag_bands():
    model = tb.make_graphene_diamond()
    if WITH_OVERLAP:
        overlap = make_overlap(model)
        solver = tb.DiagSolver(model, overlap)
    else:
        solver = tb.DiagSolver(model)
    k_points = np.array([
        [0.0, 0.0, 0.0],
        [2./3, 1./3, 0.0],
        [0.5, 0.0, 0.0],
        [0.0, 0.0, 0.0]
    ])
    k_path, k_idx = tb.gen_kpath(k_points, (1000, 1000, 1000))
    solver.config.prefix = "graphene"
    solver.config.k_points = k_path
    solver.config.diag_algo = DIAG_ALGO
    solver.config.convention = CONVENTION

    timer = tb.Timer()
    timer.tic("bands")
    k_len, bands = solver.calc_bands()
    timer.toc("bands")

    if solver.is_master:
        timer.report_total_time()
        vis = tb.Visualizer()
        vis.plot_bands(k_len, bands, k_idx, ["G", "K", "M", "G"])


def test_diag_dos():
    model = tb.make_graphene_diamond()
    if WITH_OVERLAP:
        overlap = make_overlap(model)
        solver = tb.DiagSolver(model, overlap)
    else:
        solver = tb.DiagSolver(model)
    solver.config.prefix = "graphene"
    solver.config.k_points = tb.gen_kmesh((1000, 1000, 1))
    solver.config.e_min = -10.0
    solver.config.e_max = 10.0
    solver.config.diag_algo = DIAG_ALGO
    solver.config.convention = CONVENTION

    timer = tb.Timer()
    timer.tic("dos")
    eng, dos = solver.calc_dos()
    timer.toc("dos")

    if solver.is_master:
        timer.report_total_time()
        vis = tb.Visualizer()
        vis.plot_dos(eng, dos)


def test_diag_ldos():
    model = make_graphene_pc(t=-2.7, a=0.142, with_onsite=True)
    if WITH_OVERLAP:
        overlap = make_overlap(model)
        solver = tb.DiagSolver(model, overlap)
    else:
        solver = tb.DiagSolver(model)
    solver.config.prefix = "graphene_local"
    solver.config.k_points = tb.gen_kmesh((1000, 1000, 1))
    solver.config.e_min = -10.0
    solver.config.e_max = 10.0
    solver.config.diag_algo = DIAG_ALGO
    solver.config.ldos = True
    solver.config.ldos_orbital_indices = (0,)
    solver.config.convention = CONVENTION

    timer = tb.Timer()
    timer.tic("ldos")
    eng, dos = solver.calc_dos()
    timer.toc("ldos")

    if solver.is_master:
        timer.report_total_time()
        vis = tb.Visualizer()
        vis.plot_dos(eng, dos)


def test_diag_states():
    model = tb.make_graphene_diamond()

    solver = tb.DiagSolver(model)
    solver.config.prefix = "graphene"
    solver.config.k_points = tb.gen_kmesh((3, 3, 1))
    solver.config.diag_algo = DIAG_ALGO
    solver.config.convention = CONVENTION

    timer = tb.Timer()
    timer.tic("states")
    solver.calc_states()
    timer.toc("states")

    if solver.is_master:
        timer.report_total_time()


def test_spin():
    model_graph = tb.make_graphene_soc()
    if WITH_OVERLAP:
        overlap = make_overlap(model_graph)
        solver_graph = tb.SpinTexture(model_graph, overlap)
    else:
        solver_graph = tb.SpinTexture(model_graph)
    # Params
    solver_graph.config.prefix = "kane_mele"
    solver_graph.config.k_points = 2 * (tb.gen_kmesh((640, 640, 1)) - 0.5)
    solver_graph.config.k_points[:, 2] = 0.0
    solver_graph.config.spin_major = False
    solver_graph.config.diag_algo = DIAG_ALGO
    solver_graph.config.convention = CONVENTION
    # Calculation
    timer = tb.Timer()
    timer.tic("spin_kane_mele")
    data = solver_graph.calc_spin_texture()
    timer.toc("spin_kane_mele")
    # Plot
    if solver_graph.is_master:
        timer.report_total_time()
        vis = tb.Visualizer()
        plot_sigma_band(data, vis)
        plot_sigma_eng(data, vis)

    model_bismuth = tb.make_bismuth()
    # Params
    solver_bismuth = tb.SpinTexture(model_bismuth)
    solver_bismuth.config.prefix = "bismuth"
    solver_bismuth.config.k_points = 2 * (tb.gen_kmesh((640, 640, 1)) - 0.5)
    solver_graph.config.k_points[:, 2] = 0.0
    solver_bismuth.config.spin_major = True
    solver_bismuth.config.diag_algo = DIAG_ALGO
    solver_bismuth.config.convention = CONVENTION;
    # Calculation
    timer = tb.Timer()
    timer.tic("spin_bismuth")
    data = solver_bismuth.calc_spin_texture()
    timer.toc("spin_bismuth")
    # Plot
    if solver_bismuth.is_master:
        timer.report_total_time()
        vis = tb.Visualizer()
        plot_sigma_band(data, vis)
        plot_sigma_eng(data, vis)


def test_z2():
    model = tb.make_graphene_soc()
    if WITH_OVERLAP:
        overlap = make_overlap(model)
        z2 = tb.Z2(model, overlap)
    else:
        z2 = tb.Z2(model)
    z2.config.prefix = "kane_mele"
    z2.config.num_occ = 2
    z2.config.ka_array = np.linspace(-0.5, 0.5, 2000)
    z2.config.kb_array = np.linspace(0.0, 0.5, 1000)
    z2.config.diag_algo = DIAG_ALGO
    z2.config.convention = CONVENTION

    timer = tb.Timer()
    timer.tic("z2")
    kb_array, phases = z2.calc_phases()
    timer.toc("z2")

    if z2.is_master:
        timer.report_total_time()
        vis = tb.Visualizer()
        vis.plot_phases(kb_array, phases)
        vis.plot_phases(kb_array, phases, polar=True)


def test_berry():
    model = make_haldane()
    if WITH_OVERLAP:
        overlap = make_overlap(model)
        solver = tb.Berry(model, overlap)
    else:
        solver = tb.Berry(model)
    solver.config.k_grid_size = (100, 100, 1)
    solver.config.num_occ = 1
    solver.config.diag_algo = DIAG_ALGO
    solver.config.ham_deriv_analytical = True
    solver.config.convention = CONVENTION

    timer = tb.Timer()
    timer.tic("berry")
    solver.config.prefix = "haldane_kubo"
    data_kubo = solver.calc_berry_curvature_kubo()
    solver.config.prefix = "haldane_wilson"
    data_wilson = solver.calc_berry_curvature_wilson()
    timer.toc("berry")

    if solver.is_master:
        timer.report_total_time()
        vis = tb.Visualizer()
        plot_omega_xy(data_kubo, vis)
        plot_omega_xy(data_wilson, vis)


def test_diag_dyn_pol():
    # Reference: Phys. Rev. B 84, 035439 (2011)
    # dyn_pol: fig. 3(e), eps: fig. 7(c).
    # a is C-C length, not lattice constant (0.246 in make_graphene_diamond).
    t = 3.0
    a = 0.142
    model = tb.make_graphene_diamond(t=t)
    if WITH_OVERLAP:
        overlap = make_overlap(model)
        lind = tb.Lindhard(model, overlap)
    else:
        lind = tb.Lindhard(model)
    lind.config.prefix = "graphene"
    lind.config.e_min = 0.0
    lind.config.e_max = 10.0
    lind.config.e_step = 10.0 / 2048
    lind.config.dimension = 2
    lind.config.k_grid_size = (512, 512, 1)
    lind.config.q_points = np.zeros((2, 3), dtype=np.float64)
    lind.config.q_points[0] = np.array([[0.86602540, 0.5, 0.0]]) / a
    lind.config.q_points[1] = np.array([[4.122280922013927, 2.38, 0.0]])
    lind.config.diag_algo = DIAG_ALGO
    lind.config.pp_algo = PP_ALGO
    lind.config.convention = CONVENTION

    timer = tb.Timer()
    timer.tic("dyn_pol")
    omegas, dyn_pol = lind.calc_dyn_pol()
    epsilon = lind.calc_epsilon(dyn_pol)
    timer.toc("dyn_pol")

    if lind.is_master:
        timer.report_total_time()
        vis = tb.Visualizer()
        vis.plot_xy(omegas/t, -dyn_pol[0].imag*t*a**2, color="b")
        vis.plot_xy(omegas, epsilon[1].real, color="r")

        # calc_epsilon DOES NOT save data as C++ counterpart. We have to save
        # the data manually.
        save_xy(f"{lind.config.prefix}_epsilon.dat", omegas, epsilon.T)


def test_diag_ac_cond():
    # a is C-C length, not lattice constant (0.246 in make_graphene_diamond).
    t = 3.0
    a = 0.142
    model = tb.make_graphene_diamond(t=t)
    if WITH_OVERLAP:
        overlap = make_overlap(model)
        lind = tb.Lindhard(model, overlap)
    else:
        lind = tb.Lindhard(model)
    lind.config.prefix = "graphene"
    lind.config.e_min = 0.0
    lind.config.e_max = 10.0
    lind.config.e_step = 10.0 / 2048
    lind.config.dimension = 2
    lind.config.num_spin = 2
    lind.config.k_grid_size = (512, 512, 1)
    lind.config.diag_algo = DIAG_ALGO
    lind.config.pp_algo = PP_ALGO

    # AC Cond
    timer = tb.Timer()
    timer.tic("ac_cond")
    omegas, ac_cond = lind.calc_ac_cond()
    timer.toc("ac_cond")

    if lind.is_master:
        timer.report_total_time()

        vis = tb.Visualizer()
        vis.plot_xy(omegas/t, ac_cond.real*4, color="r")

        # Epsilon_q0
        lind.config.dimension = 3
        eps0 = lind.calc_epsilon_q0(omegas, ac_cond)
        vis.plot_xy(omegas, eps0, color="r")

        # calc_epsilon_q0 DOES NOT save data as C++ counterpart. We have to save
        # the data manually.
        save_xy(f"{lind.config.prefix}_epsilon_q0.dat", omegas, eps0.T)
