#include "base/kpoints.h"
#include "base/utils.h"
#include "builder/advanced.h"
#include "builder/sample.h"
#include "tbpm/solver.h"

#include "model.h"
#include "model_sample.h"

namespace tbplas::tests {

using base::gen_kpath;
using base::Timer;
using builder::dim_t;
using builder::MagneticField;
using builder::Sample;
using tbpm::TBPMSolver;

// Algorithm of TBPM, should be in cpu, cpu_fast and gpu
constexpr char ALGO[] = "cpu";

// Whether to use smaller models for quick tests
constexpr bool 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
constexpr bool WITH_ONSITE = false;

// Whether to build a sample or a large primitive cell
#define BUILD_SAMPLE

void test_haydock()
{
    // Model
    double t = -2.7;
    double a = 0.142;
    dim_t dim;
    if (QUICK_TEST) {
        dim = { 1024, 1024, 1 };
    } else {
        dim = { 4096, 4096, 1 };
    }
#ifdef BUILD_SAMPLE
    Sample<complex_t> model = make_graphene_sample(t, a, WITH_ONSITE, dim);
    model.init_array();
    TBPMSolver<Sample<complex_t>> solver(model);
#else
    model_t model = make_graphene(t, a, WITH_ONSITE, dim);
    TBPMSolver<model_t> solver(model);
#endif

    // Parameters
    solver.config.rescale = 9.0;
    solver.config.ldos_recursion_depth = 2000;
    solver.config.dimension = 2;
    solver.config.algo = ALGO;

    // Calculation
    Timer timer;
    timer.tic("ldos_haydock");
    solver.calc_ldos_haydock();
    timer.toc("ldos_haydock");
    if (solver.is_master()) {
        timer.report_total_time();
    }
}

void test_dos()
{
    // Model
    double t = -2.7;
    double a = 0.142;
    dim_t dim;
    if (QUICK_TEST) {
        dim = { 1024, 1024, 1 };
    } else {
        dim = { 4096, 4096, 1 };
    }
#ifdef BUILD_SAMPLE
    Sample<complex_t> model = make_graphene_sample(t, a, WITH_ONSITE, dim);
    model.init_array();
    TBPMSolver<Sample<complex_t>> solver(model);
#else
    model_t model = make_graphene(t, a, WITH_ONSITE, dim);
    TBPMSolver<model_t> solver(model);
#endif

    // Parameters
    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 timer;
    timer.tic("dos");
    solver.calc_corr_dos();
    timer.toc("dos");
    if (solver.is_master()) {
        timer.report_total_time();
    }
}

void test_ac_cond()
{
    // Model
    double t = 3.0;
    double a = 0.142;
    dim_t dim;
    if (QUICK_TEST) {
        dim = { 512, 512, 1 };
    } else {
        dim = { 4096, 4096, 1 };
    }
#ifdef BUILD_SAMPLE
    Sample<complex_t> model = make_graphene_sample(t, a, WITH_ONSITE, dim);
    model.init_array();
    TBPMSolver<Sample<complex_t>> solver(model);
#else
    model_t model = make_graphene(t, a, WITH_ONSITE, dim);
    TBPMSolver<model_t> solver(model);
#endif

    // Parameters
    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 timer;
    timer.tic("ac_cond");
    solver.calc_corr_ac_cond();
    timer.toc("ac_cond");
    if (solver.is_master()) {
        timer.report_total_time();
    }
}

void test_dyn_pol()
{
    // Model
    double t = 3.0;
    double a = 0.142;
    dim_t dim;
    if (QUICK_TEST) {
        dim = { 512, 512, 1 };
    } else {
        dim = { 4096, 4096, 1 };
    }
#ifdef BUILD_SAMPLE
    Sample<complex_t> model = make_graphene_sample(t, a, WITH_ONSITE, dim);
    model.init_array();
    TBPMSolver<Sample<complex_t>> solver(model);
#else
    model_t model = make_graphene(t, a, WITH_ONSITE, dim);
    TBPMSolver<model_t> solver(model);
#endif

    // Parameters
    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 = Eigen::Matrix3Xd(3, 2);
    solver.config.dyn_pol_q_points.col(0) = Eigen::Vector3d(0.86602540, 0.5, 0.0) / a;
    solver.config.dyn_pol_q_points.col(1) = Eigen::Vector3d(4.122280922013927, 2.38, 0.0);
    solver.config.dimension = 2;
    solver.config.algo = ALGO;

    // Calculation
    Timer timer;
    timer.tic("dyn_pol");
    solver.calc_corr_dyn_pol();
    timer.toc("dyn_pol");
    if (solver.is_master()) {
        timer.report_total_time();
    }
}

void test_dc_cond()
{
    // Model
    double t = -2.7;
    double a = 0.142;
    dim_t dim;
    if (QUICK_TEST) {
        dim = { 256, 256, 1 };
    } else {
        dim = { 1024, 1024, 1 };
    }
#ifdef BUILD_SAMPLE
    Sample<complex_t> model = make_graphene_sample(t, a, WITH_ONSITE, dim);
    model.init_array();
    TBPMSolver<Sample<complex_t>> solver(model);
#else
    model_t model = make_graphene(t, a, WITH_ONSITE, dim);
    TBPMSolver<model_t> solver(model);
#endif

    // Parameters
    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 timer;
    timer.tic("dc_cond");
    solver.calc_corr_dc_cond();
    timer.toc("dc_cond");
    if (solver.is_master()) {
        timer.report_total_time();
    }
}

void test_hall_cond()
{
    // Model
    // For GPU test on 1050ti, use 64*64*1 model and 100 dckb_energies.
    double t = 3.0;
    double a = 0.142;
    dim_t dim;
    if (QUICK_TEST) {
        dim = { 256, 256, 1 };
    } else {
        dim = { 1024, 1024, 1 };
    }
#ifdef BUILD_SAMPLE
    int algo = 0;
    Sample<complex_t> model = make_graphene_sample(t, a, WITH_ONSITE, dim);
    if (algo == 0) {
        MagneticField mag(100.0);
        auto hop = model.get_sc_hopping(0);
        hop->add_modifier(mag);
        // Use force_init = true since We already called init_array in making sample.
        model.init_array(true);
    } else {
        model.apply_magnetic_field(100.0);
    }

    TBPMSolver<Sample<complex_t>> solver(model);
#else
    model_t model = make_graphene(t, a, WITH_ONSITE, dim);
    model.apply_magnetic_field(100.0);
    TBPMSolver<model_t> solver(model);
#endif

    // Parameters
    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 = Eigen::VectorXd::LinSpaced(1000, -1.0, 1.0);
    solver.config.dimension = 2;
    solver.config.dckb_unit = "h";
    solver.config.algo = ALGO;

    // Calculation
    Timer timer;
    timer.tic("hall_cond");
    solver.calc_hall_cond();
    timer.toc("hall_cond");
    if (solver.is_master()) {
        timer.report_total_time();
    }
}

void test_qe()
{
    // Model
    double t = -2.7;
    double a = 0.142;
    model_t model = make_graphene(t, a, WITH_ONSITE, { 17, 17, 1 });
    model.remove_orbitals({ 288 });

    // Parameters
    TBPMSolver<model_t> solver(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 = Eigen::Vector3d(-1.0, 0.0, 1.0);
    solver.config.algo = ALGO;

    // Calculation
    Timer timer;
    timer.tic("qe");
    solver.calc_quasi_eigenstates();
    timer.toc("qe");
    if (solver.is_master()) {
        timer.report_total_time();
    }
}

void test_wft()
{
    // Model
    double t = -2.7;
    double a = 0.142;
    dim_t dim = { 17, 17, 1 };
#ifdef BUILD_SAMPLE
    Sample<complex_t> model = make_graphene_sample(t, a, WITH_ONSITE, dim);
    model.init_array();
    TBPMSolver<Sample<complex_t>> solver(model);
#else
    model_t model = make_graphene(t, a, WITH_ONSITE, dim);
    TBPMSolver<model_t> solver(model);
#endif

    // Parameters
    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 = { 0, 1, 1024 };
    solver.config.wft_wf0 = Eigen::VectorXcd::Zero(model.get_num_orbitals());
    for (int i = 0; i < model.get_num_orbitals(); ++i) {
        if (i % 2 == 0) {
            solver.config.wft_wf0[i] = 1.0;
        }
    }
    solver.config.algo = ALGO;

    // Calculation
    Timer timer;
    timer.tic("wft");
    solver.calc_wft();
    timer.toc("wft");
    if (solver.is_master()) {
        timer.report_total_time();
    }
}

void test_qe_speed()
{
    // Model
    double t = -2.7;
    double a = 0.142;
    dim_t dim = { 4096, 4096, 1 };
#ifdef BUILD_SAMPLE
    Sample<complex_t> model = make_graphene_sample(t, a, WITH_ONSITE, dim);
    model.init_array();
    TBPMSolver<Sample<complex_t>> solver(model);
#else
    model_t model = make_graphene(t, a, WITH_ONSITE, dim);
    TBPMSolver<model_t> solver(model);
#endif

    // Parameters
    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 = Eigen::Vector3d(-1.0, 0.0, 1.0);
    solver.config.algo = ALGO;
    solver.config.save_data = false;

    // Calculation
    if (!QUICK_TEST) {
        Timer timer;
        timer.tic("qs_speed");
        solver.calc_quasi_eigenstates();
        timer.toc("qs_speed");
        if (solver.is_master()) {
            timer.report_total_time();
        }
    }
}

void test_wft_speed()
{
    // Model
    double t = -2.7;
    double a = 0.142;
    dim_t dim = { 4096, 4096, 1 };
#ifdef BUILD_SAMPLE
    Sample<complex_t> model = make_graphene_sample(t, a, WITH_ONSITE, dim);
    model.init_array();
    TBPMSolver<Sample<complex_t>> solver(model);
#else
    model_t model = make_graphene(t, a, WITH_ONSITE, dim);
    TBPMSolver<model_t> solver(model);
#endif

    // Parameters
    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 = { 0, 1, 1024 };
    solver.config.wft_wf0 = Eigen::VectorXcd::Zero(model.get_num_orbitals());
    for (int i = 0; i < model.get_num_orbitals(); ++i) {
        if (i % 2 == 0) {
            solver.config.wft_wf0[i] = 1.0;
        }
    }
    solver.config.algo = ALGO;
    solver.config.save_data = false;

    // Calculation
    if (!QUICK_TEST) {
        Timer timer;
        timer.tic("wft_speed");
        solver.calc_wft();
        timer.toc("wft_speed");
        if (solver.is_master()) {
            timer.report_total_time();
        }
    }
}

void test_bands()
{
    // Model
    double t = -2.7;
    double a = 0.142;
    dim_t dim = { 3, 3, 1 };
#ifdef BUILD_SAMPLE
    Sample<complex_t> model = make_graphene_sample(t, a, WITH_ONSITE, dim);
    model.init_array();
    TBPMSolver<Sample<complex_t>> solver(model);
#else
    model_t model = make_graphene(t, a, WITH_ONSITE, dim);
    TBPMSolver<model_t> solver(model);
#endif

    // Parameters
    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
    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(), { 40, 40, 40 });
    solver.config.k_points = k_path;

    // Calculation
    Timer timer;
    timer.tic("bands");
    solver.calc_corr_bands();
    timer.toc("bands");
    if (solver.is_master()) {
        timer.report_total_time();
        std::cout << k_idx << "\n";
    }
}

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

    // Parameters
    TBPMSolver<model_t> solver(model);
    solver.config.dimension = 2;
    solver.config.num_random_samples = 1;
    solver.config.rescale = -0.1;
    solver.config.num_time_steps = 4096;
    solver.config.algo = ALGO;

    // 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");
    solver.calc_corr_bands();
    timer.toc("bands_tbg");
    if (solver.is_master()) {
        timer.report_total_time();
        std::cout << k_idx << "\n";
    }
}
} // namespace tbplas::tests
