#include "base/kpoints.h"
#include "base/utils.h"
#include "diag/base.h"
#include "diag/berry.h"
#include "diag/lindhard.h"
#include "diag/magnetic.h"
#include "diag/z2.h"
#include "materials/bismuth.h"
#include "materials/graphene.h"

#include "model.h"

namespace tbplas::tests {

using base::gen_kmesh;
using base::gen_kpath;
using base::Timer;
using diag::Berry;
using diag::DiagSolver;
using diag::Lindhard;
using diag::SpinTexture;
using diag::Z2;
using materials::make_bismuth;
using materials::make_graphene_diamond;
using materials::make_graphene_soc;

// Algorithm for diagonalization, should be in eigen or lapacke
constexpr char DIAG_ALGO[] = "eigen";

// Algorithm for post-processing, currently only supports cpu
constexpr char PP_ALGO[] = "cpu";

// Convention of Hamiltonian, should be 1 or 2
constexpr int CONVENTION = 1;

// Whether to test algorithms with overlaps (zhegv)
// #define WITH_OVERLAP

void test_diag_bands()
{
    model_t model = make_graphene_diamond();
#ifdef WITH_OVERLAP
    model_t overlap = make_overlap(model);
    DiagSolver<model_t> solver(model, overlap);
#else
    DiagSolver<model_t> solver(model);
#endif
    Eigen::MatrixX3d k_points { { 0.0, 0.0, 0.0 }, { 2. / 3, 1. / 3, 0.0 }, { 0.5, 0.0, 0.0 }, { 0.0, 0.0, 0.0 } };
    Eigen::Matrix3Xd k_path;
    Eigen::VectorXi k_idx;
    std::tie(k_path, k_idx) = gen_kpath(k_points.transpose(), { 1000, 1000, 1000 });
    solver.config.prefix = "graphene";
    solver.config.k_points = k_path;
    solver.config.diag_algo = DIAG_ALGO;
    solver.config.convention = CONVENTION;

    Timer timer;
    timer.tic("bands");
    auto data = solver.calc_bands();
    timer.toc("bands");
    if (solver.is_master()) {
        timer.report_total_time();
    }
}

void test_diag_dos()
{
    model_t model = make_graphene_diamond();
#ifdef WITH_OVERLAP
    model_t overlap = make_overlap(model);
    DiagSolver<model_t> solver(model, overlap);
#else
    DiagSolver<model_t> solver(model);
#endif
    solver.config.prefix = "graphene";
    solver.config.k_points = 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 timer;
    timer.tic("dos");
    auto data = solver.calc_dos();
    timer.toc("dos");
    if (solver.is_master()) {
        timer.report_total_time();
    }
}

void test_diag_ldos()
{
    model_t model = make_graphene_pc(-2.7, 0.142, true);
#ifdef WITH_OVERLAP
    model_t overlap = make_overlap(model);
    DiagSolver<model_t> solver(model, overlap);
#else
    DiagSolver<model_t> solver(model);
#endif
    solver.config.prefix = "graphene_local";
    solver.config.k_points = 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 timer;
    timer.tic("ldos");
    auto data = solver.calc_dos();
    timer.toc("ldos");
    if (solver.is_master()) {
        timer.report_total_time();
    }
}

void test_diag_states()
{
    model_t model = make_graphene_diamond();

    DiagSolver<model_t> solver(model);
    solver.config.prefix = "graphene";
    solver.config.k_points = gen_kmesh({ 3, 3, 1 });
    solver.config.diag_algo = DIAG_ALGO;
    solver.config.convention = CONVENTION;

    Timer timer;
    timer.tic("states");
    solver.calc_states();
    timer.toc("states");
    if (solver.is_master()) {
        timer.report_total_time();
    }
}

void test_spin()
{
    model_t model_graph = make_graphene_soc();
#ifdef WITH_OVERLAP
    model_t overlap = make_overlap(model_graph);
    SpinTexture<model_t> solver_graph(model_graph, overlap);
#else
    SpinTexture<model_t> solver_graph(model_graph);
#endif
    // Params
    solver_graph.config.prefix = "kane_mele";
    solver_graph.config.k_points = 2 * (gen_kmesh({ 640, 640, 1 }).array() - 0.5);
    solver_graph.config.k_points.row(2).array() = 0.0;
    solver_graph.config.spin_major = false;
    solver_graph.config.diag_algo = DIAG_ALGO;
    solver_graph.config.convention = CONVENTION;
    // Calculation
    Timer timer;
    timer.tic("spin_kane_mele");
    auto data = solver_graph.calc_spin_texture();
    timer.toc("spin_kane_mele");

    model_t model_bismuth = make_bismuth();
    // Params
    SpinTexture<model_t> solver_bismuth(model_bismuth);
    solver_bismuth.config.prefix = "bismuth";
    solver_bismuth.config.k_points = 2 * (gen_kmesh({ 640, 640, 1 }).array() - 0.5);
    solver_bismuth.config.k_points.row(2).array() = 0.0;
    solver_bismuth.config.spin_major = true;
    solver_bismuth.config.diag_algo = DIAG_ALGO;
    solver_bismuth.config.convention = CONVENTION;
    // Calculation
    timer.tic("spin_bismuth");
    data = solver_bismuth.calc_spin_texture();
    timer.toc("spin_bismuth");
    if (solver_bismuth.is_master()) {
        timer.report_total_time();
    }
}

void test_z2()
{
    model_t model = make_graphene_soc();
#ifdef WITH_OVERLAP
    model_t overlap = make_overlap(model);
    Z2<model_t> z2(model, overlap);
#else
    Z2<model_t> z2(model);
#endif
    z2.config.prefix = "kane_mele";
    z2.config.num_occ = 2;
    z2.config.ka_array = Eigen::VectorXd::LinSpaced(2000, -0.5, 0.5);
    z2.config.kb_array = Eigen::VectorXd::LinSpaced(1000, 0.0, 0.5);
    z2.config.diag_algo = DIAG_ALGO;
    z2.config.convention = CONVENTION;

    Timer timer;
    timer.tic("z2");
    auto data = z2.calc_phases();
    timer.toc("z2");
    if (z2.is_master()) {
        timer.report_total_time();
    }
}

void test_berry()
{
    // QAH model;
    // Berry<QAH> solver(model);
    // CheckerBoard model(20);
    // Berry<CheckerBoard> solver(model);
    // Haldane model;
    // Berry<Haldane> solver(model);
    model_t model = make_haldane();
#ifdef WITH_OVERLAP
    model_t overlap = make_overlap(model);
    Berry<model_t> solver(model, overlap);
#else
    Berry<model_t> solver(model);
#endif
    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 timer;
    timer.tic("berry");
    solver.config.prefix = "haldane_kubo";
    auto data_kubo = solver.calc_berry_curvature_kubo();
    solver.config.prefix = "haldane_wilson";
    auto data_wilson = solver.calc_berry_curvature_wilson();
    timer.toc("berry");
    if (solver.is_master()) {
        timer.report_total_time();
    }
}

void test_diag_dyn_pol()
{
    // Reference: tbplas/tests/test_units/test_lindhard.py
    double t = 3.0;
    double a = 0.142;
    model_t model = make_graphene_diamond(1.0, t);
#ifdef WITH_OVERLAP
    model_t overlap = make_overlap(model);
    Lindhard<model_t> lind(model, overlap);
#else
    Lindhard<model_t> lind(model);
#endif
    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 = Eigen::Matrix3Xd(3, 2);
    lind.config.q_points.col(0) = Eigen::Vector3d(0.86602540, 0.5, 0.0) / a;
    lind.config.q_points.col(1) = Eigen::Vector3d(4.122280922013927, 2.38, 0.0);
    lind.config.diag_algo = DIAG_ALGO;
    lind.config.pp_algo = PP_ALGO;
    lind.config.convention = CONVENTION;

    Timer timer;
    timer.tic("dyn_pol");
    auto dyn_pol = lind.calc_dyn_pol();
    auto epsilon = lind.calc_epsilon(dyn_pol);
    timer.toc("dyn_pol");
    if (lind.is_master()) {
        timer.report_total_time();
    }
}

void test_diag_ac_cond()
{
    // Reference: tbplas/tests/test_units/test_lindhard.py
    double t = 3.0;
    model_t model = make_graphene_diamond(1.0, t);
#ifdef WITH_OVERLAP
    model_t overlap = make_overlap(model);
    Lindhard<model_t> lind(model, overlap);
#else
    Lindhard<model_t> lind(model);
#endif
    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 timer;
    timer.tic("ac_cond");
    auto ac_cond = lind.calc_ac_cond();
    timer.toc("ac_cond");
    if (lind.is_master()) {
        timer.report_total_time();
    }

    // Epsilon_q0
    lind.config.dimension = 3;
    auto eps0 = lind.calc_epsilon_q0(ac_cond);
}

} // namespace tbplas::tests
