"""
Test Python interface for DIAG algorithms based on FEAST library.

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
from model import make_graphene_pc, make_overlap


__all__ = ["test_diag_bands_feast", "test_diag_dos_feast", "test_diag_ldos_feast",
           "test_spin_feast", "test_diag_bands_tbg"]


# Algorithm for diagonalization, should be in feast
FEAST_DIAG_ALGO = "feast_dense"

# Whether to reuse results of last calculation
FEAST_REUSE_RESULTS = False

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

# Whether to test algorithms with overlaps (zfeast_hegv/zfeast_hcsrgv)
WITH_OVERLAP = False


def test_diag_bands_feast():
    model = tb.make_graphene_diamond()
    model = tb.extend_prim_cell(model, dim=(3, 3, 1))
    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, (50, 50, 50))
    solver.config.prefix = "graphene"
    solver.config.k_points = k_path

    # Use FEAST solver to evaluate all the states within [-10, 10] eV.
    solver.config.diag_algo = FEAST_DIAG_ALGO
    solver.config.feast_e_min = -10
    solver.config.feast_e_max = 10
    solver.config.feast_num_states = model.num_orb
    # Not working on n02 since it calls the MKL version of FEAST, which do not
    # support negative fpm[0]. Don't know why.
    # solver.config.feast_fpm[0] = -1
    # feast sparse conflicts with reuse_results at some k-point.
    solver.config.feast_reuse_results = False

    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_feast():
    model = tb.make_graphene_diamond()
    model = tb.extend_prim_cell(model, dim=(3, 3, 1))
    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((32, 32, 1))
    solver.config.e_min = -10.0
    solver.config.e_max = 10.0

    # Use FEAST solver to evaluate all the states within [-10, 10] eV.
    solver.config.diag_algo = FEAST_DIAG_ALGO
    solver.config.feast_e_min = -10
    solver.config.feast_e_max = 10
    solver.config.feast_num_states = model.num_orb
    # solver.config.feast_fpm[0] = -2
    solver.config.feast_reuse_results = FEAST_REUSE_RESULTS

    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_feast():
    model = make_graphene_pc(t=-2.7, a=0.142, with_onsite=True)
    model = tb.extend_prim_cell(model, dim=(3, 3, 1))
    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((32, 32, 1))
    solver.config.e_min = -10.0
    solver.config.e_max = 10.0

    # Use FEAST solver to evaluate all the states within [-10, 10] eV.
    solver.config.diag_algo = FEAST_DIAG_ALGO
    solver.config.feast_e_min = -10
    solver.config.feast_e_max = 10
    solver.config.feast_num_states = model.num_orb
    # solver.config.feast_fpm[0] = -2
    solver.config.feast_reuse_results = FEAST_REUSE_RESULTS
    solver.config.ldos = True
    solver.config.ldos_orbital_indices = (0,)

    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_spin_feast():
    model_graph = tb.make_graphene_soc()
    model_graph = tb.extend_prim_cell(model_graph, dim=(3, 3, 1))
    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((32, 32, 1)) - 0.5)
    solver_graph.config.k_points[:, 2] = 0.0
    solver_graph.config.spin_major = False
    # Use SpinTexture for testing since it involves both energies and wave functions.
    # Use FEAST solver to evaluate all the states within [-6, 6] eV.
    solver_graph.config.diag_algo = FEAST_DIAG_ALGO
    solver_graph.config.feast_e_min = -6
    solver_graph.config.feast_e_max = 6
    solver_graph.config.feast_num_states = model_graph.num_orb
    # solver_graph.config.feast_fpm[0] = -3
    solver_graph.config.feast_reuse_results = FEAST_REUSE_RESULTS
    # 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)


def test_diag_bands_tbg():
    raise NotImplementedError("test_diag_bands_tbg not implemented yet!")
