#include "base/kpoints.h"
#include "base/utils.h"
#include "builder/advanced.h"
#include "diag/base.h"
#include "diag/magnetic.h"
#include "materials/graphene.h"

#include "model.h"

namespace tbplas::tests {

using base::gen_kmesh;
using base::gen_kpath;
using base::Timer;
using diag::DiagSolver;
using diag::SpinTexture;
using materials::make_graphene_diamond;
using materials::make_graphene_soc;

// Algorithm for diagonalization, should be in feast
constexpr char FEAST_DIAG_ALGO[] = "feast_dense";

// Whether to reuse results of last calculation
constexpr bool FEAST_REUSE_RESULTS = false;

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

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

void test_diag_bands_feast()
{
    model_t model = make_graphene_diamond();
    model = builder::extend_prim_cell(model, { 3, 3, 1 });
#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(), { 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.get_num_orbitals();
    solver.config.feast_fpm[0] = -1;
    // feast sparse conflicts with reuse_results at some k-point.
    solver.config.feast_reuse_results = false;

    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_feast()
{
    model_t model = make_graphene_diamond();
    model = builder::extend_prim_cell(model, { 3, 3, 1 });
#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({ 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.get_num_orbitals();
    solver.config.feast_fpm[0] = -2;
    solver.config.feast_reuse_results = FEAST_REUSE_RESULTS;

    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_feast()
{
    model_t model = make_graphene_pc(-2.7, 0.142, true);
    model = builder::extend_prim_cell(model, { 3, 3, 1 });
#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({ 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.get_num_orbitals();
    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 timer;
    timer.tic("ldos");
    auto data = solver.calc_dos();
    timer.toc("ldos");
    if (solver.is_master()) {
        timer.report_total_time();
    }
}

void test_spin_feast()
{
    model_t model_graph = make_graphene_soc();
    model_graph = extend_prim_cell(model_graph, { 3, 3, 1 });
#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({ 32, 32, 1 }).array() - 0.5);
    solver_graph.config.k_points.row(2).array() = 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.get_num_orbitals();
    solver_graph.config.feast_fpm[0] = -3;
    solver_graph.config.feast_reuse_results = FEAST_REUSE_RESULTS;
    // Calculation
    Timer timer;
    timer.tic("spin_kane_mele");
    auto data = solver_graph.calc_spin_texture();
    timer.toc("spin_kane_mele");
    if (solver_graph.is_master()) {
        timer.report_total_time();
    }
}

void test_diag_bands_tbg()
{
    // Model
    double t = -2.7;
    double a = 0.142;
    model_t model = make_tbg(31);

    // Parameters
    DiagSolver<model_t> solver(model);
    solver.config.prefix = "tbg";

    // Use FEAST solver to evaluate the 15 within [-0.2, 0.2] + 0.78 eV.
    solver.config.diag_algo = FEAST_DIAG_ALGO;
    solver.config.feast_e_min = -0.2 + 0.78;
    solver.config.feast_e_max = 0.2 + 0.78;
    solver.config.feast_num_states = 15;
    solver.config.feast_m0_ratio = 2.0;
    solver.config.feast_fpm[0] = -1;
    // feast sparse conflicts with reuse_results at some k-point.
    solver.config.feast_reuse_results = false;

    // K-points
    Eigen::MatrixX3d k_points {
        { 2. / 3, 1. / 3, 0.0 },
        { 0.0, 0.0, 0.0 },
        { 0.5, 0.0, 0.0 },
        { 2. / 3, 1. / 3, 0.0 },
    };
    Eigen::Matrix3Xd k_path;
    Eigen::VectorXi k_idx;
    std::tie(k_path, k_idx) = gen_kpath(k_points.transpose(), { 40, 40, 40 });
    solver.config.k_points = k_path;

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

} // namespace tbplas::tests
