"""
Test Python interface for TBPM algorithms.

This module generates the same output files (*.dat or *.h5) as the C++ version, which can be
post-processed by 'plot_tbpm.py' to produce *_cpp.dat. Meanwhile, it also calls Python
TBPMSolver and Analyzer on-the-fly to produce *_py.dat.

*_py.dat and *_cpp.dat can be plotted for comparison with the 'plot_tbpm_compare.py' script.
"""

import numpy as np

import tbplas as tb

from model import make_graphene
from model_sample import make_graphene_sample
from utils import save_xy


__all__ = ["test_haydock", "test_dos", "test_ac_cond", "test_dyn_pol",
           "test_dc_cond", "test_hall_cond", "test_qe", "test_wft",
           "test_qe_speed", "test_wft_speed", "test_bands", "test_bands_tbg"]


# Algorithm of TBPM, should be in cpu, cpu_fast and gpu
ALGO = "cpu"

# Whether to use smaller models for quick tests
QUICK_TEST = True

# Whether to use different orbital on-site energies for testing LDOS and
# for a fair benchmark with TBPLaS v1.x which drops zero on-site terms
WITH_ONSITE = False

# Whether to build a sample or a large primitive cell
BUILD_SAMPLE = True


def test_haydock():
    # Model
    t = -2.7
    a = 0.142
    if QUICK_TEST:
        dim = (1024, 1024, 1)
    else:
        dim = (4096, 4096, 1)
    if BUILD_SAMPLE:
        model = make_graphene_sample(t, a, with_onsite=WITH_ONSITE, dim=dim)
    else:
        model = make_graphene(t, a, with_onsite=WITH_ONSITE, dim=dim)

    # Parameters
    solver = tb.TBPMSolver(model)
    solver.config.rescale = 9.0
    solver.config.ldos_recursion_depth = 2000
    solver.config.dimension = 2
    solver.config.algo = ALGO

    # Calculation
    timer = tb.Timer()
    timer.tic("ldos_haydock")
    eng, ldos = solver.calc_ldos_haydock()
    timer.toc("ldos_haydock")

    # Output
    if solver.is_master:
        timer.report_total_time()
        save_xy("ldos_py.dat", eng, ldos)


def test_dos():
    # Model
    t = -2.7
    a = 0.142
    if QUICK_TEST:
        dim = (1024, 1024, 1)
    else:
        dim = (4096, 4096, 1)
    if BUILD_SAMPLE:
        model = make_graphene_sample(t, a, with_onsite=WITH_ONSITE, dim=dim)
    else:
        model = make_graphene(t, a, with_onsite=WITH_ONSITE, dim=dim)

    # Parameters
    solver = tb.TBPMSolver(model)
    solver.config.num_random_samples = 1
    solver.config.rescale = 9.0
    solver.config.ldos = False
    solver.config.ldos_orbital_indices = { 1 }
    solver.config.num_time_steps = 1024
    solver.config.dimension = 2
    solver.config.algo = ALGO

    # Calculation
    timer = tb.Timer()
    timer.tic("corr_dos")
    corr_dos = solver.calc_corr_dos()
    timer.toc("corr_dos")

    # Output
    if solver.is_master:
        timer.report_total_time()
        analyzer = tb.Analyzer(f"{solver.config.prefix}_info.dat")
        eng, dos = analyzer.calc_dos(corr_dos)
        save_xy("dos_py.dat", eng, dos)


def test_ac_cond():
    # Model
    # Reference: tbplas/tests/test_units/test_ac.py
    t = 3.0
    a = 0.142
    if QUICK_TEST:
        dim = (512, 512, 1)
    else:
        dim = (4096, 4096, 1)
    if BUILD_SAMPLE:
        model = make_graphene_sample(t, a, with_onsite=WITH_ONSITE, dim=dim)
    else:
        model = make_graphene(t, a, with_onsite=WITH_ONSITE, dim=dim)

    # Parameters
    solver = tb.TBPMSolver(model)
    solver.config.rescale = 9.0
    solver.config.num_random_samples = 1
    solver.config.num_time_steps = 1024
    solver.config.num_spin = 2
    solver.config.dimension = 2
    solver.config.algo = ALGO

    # Calculation
    timer = tb.Timer()
    timer.tic("corr_ac_cond")
    corr = solver.calc_corr_ac_cond()
    timer.toc("corr_ac_cond")

    # Output
    if solver.is_master:
        timer.report_total_time()
        analyzer = tb.Analyzer(f"{solver.config.prefix}_info.dat")
        energies, ac_cond = analyzer.calc_ac_cond(corr)
        save_xy("ac_cond_py.dat", energies/t, ac_cond[0].real*4)


def test_dyn_pol():
    # Model
    # Reference: tbplas/tests/test_units/test_dp.py and test_eps.py
    t = 3.0
    a = 0.142
    if QUICK_TEST:
        dim = (512, 512, 1)
    else:
        dim = (4096, 4096, 1)
    if BUILD_SAMPLE:
        model = make_graphene_sample(t, a, with_onsite=WITH_ONSITE, dim=dim)
    else:
        model = make_graphene(t, a, with_onsite=WITH_ONSITE, dim=dim)

    # Parameters
    solver = tb.TBPMSolver(model)
    solver.config.rescale = 9.0
    solver.config.num_random_samples = 1
    solver.config.num_time_steps = 1024
    solver.config.num_spin = 1
    solver.config.dyn_pol_q_points = np.zeros((2, 3), dtype=np.float64)
    solver.config.dyn_pol_q_points[0] = np.array([[0.86602540, 0.5, 0.0]]) / a
    solver.config.dyn_pol_q_points[1] = np.array([[4.122280922013927, 2.38, 0.0]])
    solver.config.dimension = 2
    solver.config.algo = ALGO

    # Calculation
    timer = tb.Timer()
    timer.tic("corr_dyn_pol")
    corr = solver.calc_corr_dyn_pol()
    timer.toc("corr_dyn_pol")

    # Output
    if solver.is_master:
        timer.report_total_time()
        analyzer = tb.Analyzer(f"{solver.config.prefix}_info.dat")
        q, omegas, dp = analyzer.calc_dyn_pol(corr)
        epsilon = analyzer.calc_epsilon(dp)
        save_xy("dyn_pol_py.dat", omegas/t, -dp[0].imag*t*a**2)
        save_xy("epsilon_py.dat", omegas, epsilon[1].real)


def test_dc_cond():
    # Model
    # Update to C++ version
    t = -2.7
    a = 0.142
    if QUICK_TEST:
        dim = (256, 256, 1)
    else:
        dim = (1024, 1024, 1)
    if BUILD_SAMPLE:
        model = make_graphene_sample(t, a, with_onsite=WITH_ONSITE, dim=dim)
    else:
        model = make_graphene(t, a, with_onsite=WITH_ONSITE, dim=dim)

    # Parameters
    solver = tb.TBPMSolver(model)
    solver.config.rescale = 9.0
    solver.config.num_random_samples = 1
    solver.config.num_time_steps = 1024
    solver.config.num_spin = 1
    solver.config.dc_cond_energy_limits = [(-2.5, 2.5)]
    solver.config.dc_cond_num_time_steps = 1024
    solver.config.dimension = 2
    solver.config.algo = ALGO

    # Calculation
    timer = tb.Timer()
    timer.tic("corr_dc_cond")
    corr_dos, corr_dc = solver.calc_corr_dc_cond()
    timer.toc("corr_dc_cond")

    # Output
    if solver.is_master:
        timer.report_total_time()
        analyzer = tb.Analyzer(f"{solver.config.prefix}_info.dat")
        energy, dc = analyzer.calc_dc_cond(corr_dos, corr_dc)
        save_xy("dc_cond_xx_py.dat", energy, dc[0])


def test_hall_cond():
    # Model
    # Reference: tbplas/tests/test_units/test_hall.py
    t = 3.0
    a = 0.142
    if QUICK_TEST:
        dim = (256, 256, 1)
    else:
        dim = (1024, 1024, 1)
    if BUILD_SAMPLE:
        model = make_graphene_sample(t, a, with_onsite=WITH_ONSITE, dim=dim)
    else:
        model = make_graphene(t, a, with_onsite=WITH_ONSITE, dim=dim)
    model.apply_magnetic_field(100.0)

    # Parameters
    solver = tb.TBPMSolver(model)
    solver.config.rescale = 9.0
    solver.config.num_random_samples = 1
    solver.config.num_time_steps = 1024
    solver.config.num_spin = 2
    solver.config.temperature = 10.0
    solver.config.dckb_component = 2
    solver.config.dckb_energies = np.linspace(-1.0, 1.0, 1000)
    solver.config.dimension = 2
    solver.config.dckb_unit = "h"
    solver.config.algo = ALGO

    # Calculation
    timer = tb.Timer()
    timer.tic("hall_cond")
    omegas, hall = solver.calc_hall_cond()
    timer.toc("hall_cond")

    # Output
    if solver.is_master:
        timer.report_total_time()
        save_xy("hall_cond_py.dat", omegas, hall)


def test_qe():
    # Model
    t = -2.7
    a = 0.142
    model = make_graphene(t, a, with_onsite=WITH_ONSITE, dim=(17, 17, 1))
    model.remove_orbitals([288])

    # Parameters
    solver = tb.TBPMSolver(model)
    solver.config.dimension = 2
    solver.config.num_random_samples = 1
    solver.config.rescale = 9.0
    solver.config.num_time_steps = 1024
    solver.config.qe_energies = np.array([-1.0, 0.0, 1.0])
    solver.config.algo = ALGO

    # Calculation
    timer = tb.Timer()
    timer.tic("qe")
    qs = solver.calc_quasi_eigenstates()
    timer.toc("qe")

    # Output
    if solver.is_master:
        timer.report_total_time()
        np.savetxt("qe_py.dat", qs.T)


def test_wft():
    # Model
    t = -2.7
    a = 0.142
    dim = (17, 17, 1)
    if BUILD_SAMPLE:
        model = make_graphene_sample(t, a, with_onsite=WITH_ONSITE, dim=dim)
    else:
        model = make_graphene(t, a, with_onsite=WITH_ONSITE, dim=dim)

    # Parameters
    solver = tb.TBPMSolver(model)
    solver.config.dimension = 2
    solver.config.num_random_samples = 1
    solver.config.rescale = 9.0
    solver.config.num_time_steps = 1024
    solver.config.wft_time_save = np.array([0, 1, 1024])
    solver.config.wft_wf0 = np.zeros(model.num_orb, dtype=np.complex128)
    for i in range(model.num_orb):
        if i % 2 == 0:
            solver.config.wft_wf0[i] = 1.0
    solver.config.algo = ALGO

    # Calculation
    timer = tb.Timer()
    timer.tic("wft")
    wft = solver.calc_wft()
    timer.toc("wft")

    # Output
    if solver.is_master:
        timer.report_total_time()
        wft = wft.T
        num_row, num_col = wft.shape
        data = np.zeros((num_row, 2*num_col))
        data[:, 0::2] = wft.real
        data[:, 1::2] = wft.imag
        np.savetxt("wft_py.dat", data)


def test_qe_speed():
    # Model
    t = -2.7
    a = 0.142
    dim = (4096, 4096, 1)
    if BUILD_SAMPLE:
        model = make_graphene_sample(t, a, with_onsite=WITH_ONSITE, dim=dim)
    else:
        model = make_graphene(t, a, with_onsite=WITH_ONSITE, dim=dim)

    # Parameters
    solver = tb.TBPMSolver(model)
    solver.config.dimension = 2
    solver.config.num_random_samples = 1
    solver.config.rescale = 9.0
    solver.config.num_time_steps = 1024
    solver.config.qe_energies = np.array([-1.0, 0.0, 1.0])
    solver.config.algo = ALGO
    solver.config.save_data = False

    # Calculation
    if not QUICK_TEST:
        timer = tb.Timer()
        timer.tic("qe_speed")
        solver.calc_quasi_eigenstates()
        timer.toc("qe_speed")
        timer.report_total_time()


def test_wft_speed():
    # Model
    t = -2.7
    a = 0.142
    dim = (4096, 4096, 1)
    if BUILD_SAMPLE:
        model = make_graphene_sample(t, a, with_onsite=WITH_ONSITE, dim=dim)
    else:
        model = make_graphene(t, a, with_onsite=WITH_ONSITE, dim=dim)

    # Parameters
    solver = tb.TBPMSolver(model)
    solver.config.dimension = 2
    solver.config.num_random_samples = 1
    solver.config.rescale = 9.0
    solver.config.num_time_steps = 1024
    solver.config.wft_time_save = np.array([0, 1, 1024])
    solver.config.wft_wf0 = np.zeros(model.num_orb, dtype=np.complex128)
    for i in range(model.num_orb):
        if i % 2 == 0:
            solver.config.wft_wf0[i] = 1.0
    solver.config.algo = ALGO
    solver.config.save_data = False

    # Calculation
    if not QUICK_TEST:
        timer = tb.Timer()
        timer.tic("wft_speed")
        solver.calc_wft()
        timer.toc("wft_speed")
        timer.report_total_time()


def test_bands():
    # Model
    t = -2.7
    a = 0.142
    dim = (3, 3, 1)
    if BUILD_SAMPLE:
        model = make_graphene_sample(t, a, with_onsite=WITH_ONSITE, dim=dim)
    else:
        model = make_graphene(t, a, with_onsite=WITH_ONSITE, dim=dim)

    # Parameters
    solver = tb.TBPMSolver(model)
    solver.config.dimension = 2
    solver.config.num_random_samples = 1
    solver.config.rescale = 9.0
    solver.config.num_time_steps = 1024
    solver.config.algo = ALGO

    # K-points
    k_points = np.array([
        [0.0, 0.0, 0.0],
        [2. / 3, 1. / 3, 0.0],
        [1. / 2, 0.0, 0.0],
        [0.0, 0.0, 0.0],
    ])
    k_path, k_idx = tb.gen_kpath(k_points, [40, 40, 40])
    solver.config.k_points = k_path

    # Calculation
    k_len, corr_bands = solver.calc_corr_bands()

    # Output
    if solver.is_master:
        analyzer = tb.Analyzer(f"{solver.config.prefix}_info.dat")
        energy, bands = analyzer.calc_bands(corr_bands)
        np.savetxt("energy", energy)
        np.savetxt("bands", bands)
        np.savetxt("k_len", k_len)
        np.savetxt("k_idx", k_idx)


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