#include <iostream>

#include "base/parallel.h"

#include "diag.h"
#include "diag_feast.h"
#include "model.h"
#include "model_sample.h"
#include "tbpm.h"

#define LOOP_TEST(n, cmd)         \
    for (int i = 0; i < n; ++i) { \
        cmd;                      \
    }

int run_diag(const std::string& test)
{
    int matched = 1;
    if (test == "diag_bands") {
        tbplas::tests::test_diag_bands();
    } else if (test == "diag_dos") {
        tbplas::tests::test_diag_dos();
    } else if (test == "diag_ldos") {
        tbplas::tests::test_diag_ldos();
    } else if (test == "diag_states") {
        tbplas::tests::test_diag_states();
    } else if (test == "spin") {
        tbplas::tests::test_spin();
    } else if (test == "z2") {
        tbplas::tests::test_z2();
    } else if (test == "berry") {
        tbplas::tests::test_berry();
    } else if (test == "diag_dyn_pol") {
        tbplas::tests::test_diag_dyn_pol();
    } else if (test == "diag_ac_cond") {
        tbplas::tests::test_diag_ac_cond();
    } else if (test == "diag_bands_feast") {
        tbplas::tests::test_diag_bands_feast();
    } else if (test == "diag_dos_feast") {
        tbplas::tests::test_diag_dos_feast();
    } else if (test == "diag_ldos_feast") {
        tbplas::tests::test_diag_ldos_feast();
    } else if (test == "spin_feast") {
        tbplas::tests::test_spin_feast();
    } else {
        matched = 0;
    }
    return matched;
}

int run_tbpm(const std::string& test)
{
    int matched = 1;
    if (test == "haydock") {
        tbplas::tests::test_haydock();
    } else if (test == "dos") {
        tbplas::tests::test_dos();
    } else if (test == "ac_cond") {
        tbplas::tests::test_ac_cond();
    } else if (test == "dyn_pol") {
        tbplas::tests::test_dyn_pol();
    } else if (test == "dc_cond") {
        tbplas::tests::test_dc_cond();
    } else if (test == "hall_cond") {
        tbplas::tests::test_hall_cond();
    } else if (test == "qe") {
        tbplas::tests::test_qe();
    } else if (test == "wft") {
        tbplas::tests::test_wft();
    } else if (test == "qe_speed") {
        tbplas::tests::test_qe_speed();
    } else if (test == "wft_speed") {
        tbplas::tests::test_wft_speed();
    } else if (test == "bands") {
        tbplas::tests::test_bands();
    } else {
        matched = 0;
    }
    return matched;
}

int run_model(const std::string& test)
{
    int tbg_idx = 5;
    double radius = 3.0;
    int iteration = 3;

    int matched = 1;
    if (test == "tbg") {
        LOOP_TEST(5, tbplas::tests::make_tbg(tbg_idx));
    } else if (test == "tbg_sample") {
        LOOP_TEST(5, tbplas::tests::make_tbg_sample(tbg_idx));
    } else if (test == "qc") {
        LOOP_TEST(5, tbplas::tests::make_quasi_crystal(radius));
    } else if (test == "qc_sample") {
        LOOP_TEST(5, tbplas::tests::make_quasi_crystal_sample(radius));
    } else if (test == "fractal") {
        LOOP_TEST(5, tbplas::tests::make_fractal(iteration));
    } else if (test == "fractal_sample") {
        LOOP_TEST(5, tbplas::tests::make_fractal_sample(iteration));
    } else {
        matched = 0;
    }
    return matched;
}

int main(int argc, char* argv[])
{
    tbplas::base::MPIENV_INIT(argc, argv);
    Eigen::initParallel();
    std::string test = argc >= 2 ? argv[1] : "dos";

    // New tests should be add here.
    int matched = run_diag(test);
    matched += run_tbpm(test);
    matched += run_model(test);

    if (matched == 0) {
        std::cout << "Unknown test: " << test << "\n";
    }
    tbplas::base::MPIENV_FINALIZE();
}
