#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_all") {
        tbplas::tests::test_diag_bands();
        tbplas::tests::test_diag_dos();
        tbplas::tests::test_diag_ldos();
        tbplas::tests::test_diag_states();
        tbplas::tests::test_spin();
        tbplas::tests::test_z2();
        tbplas::tests::test_berry();
        tbplas::tests::test_diag_dyn_pol();
        tbplas::tests::test_diag_ac_cond();
    } else if (test == "feast") {
        tbplas::tests::test_diag_bands_feast();
        tbplas::tests::test_diag_dos_feast();
        tbplas::tests::test_diag_ldos_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 if (test == "bands_tbg") {
//         tbplas::tests::test_bands_tbg();
//     } else if (test == "dos_tbg") {
//         tbplas::tests::test_dos_tbg();
//     } else if (test == "qe_tbg") {
//         tbplas::tests::test_qe_tbg();
//     } else if (test == "dc_cond_tbg") {
//         tbplas::tests::test_dc_cond_tbg();
//     } else if (test == "ac_cond_tbg") {
//         tbplas::tests::test_ac_cond_tbg();
//     } else if (test == "dyn_pol_tbg") {
//         tbplas::tests::test_dyn_pol_tbg();
//     } else if (test == "cd_tbg") {
//         tbplas::tests::test_cd_tbg();
//     } else {
//         matched = 0;
//     }
//     return matched;
// }

int run_tbpm(const std::string& test, int b_value, int E_value = 5, int T_value = 1024, int seed=1337)
{
    int matched = 1;
    if (test == "dos_tbg") {
        tbplas::tests::test_dos_tbg(b_value, T_value, seed);
    } else if (test == "qe_tbg") {
        tbplas::tests::test_qe_tbg(b_value, E_value, T_value, seed);
    } else if (test == "dc_cond_tbg") {
        tbplas::tests::test_dc_cond_tbg(b_value, E_value, T_value, seed);
    } else if (test == "ac_cond_tbg") {
        tbplas::tests::test_ac_cond_tbg(b_value, T_value, seed);
    } else if (test == "dyn_pol_tbg") {
        tbplas::tests::test_dyn_pol_tbg(b_value, T_value, seed);
    } else if (test == "cd_tbg") {
        tbplas::tests::test_cd_tbg(b_value, T_value, seed);
    } 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);
    if (argc == 3) {
        int b_value = std::atoi(argv[2]); // Convert string argument to integer
        matched += run_tbpm(test, b_value);
    } else if(argc == 4) {
        int b_value = std::atoi(argv[2]); // Convert string argument to integer
        int E_value = std::atoi(argv[3]);
        matched += run_tbpm(test, b_value, E_value);
    } else if(argc == 5) {
        int b_value = std::atoi(argv[2]); // Convert string argument to integer
        int E_value = std::atoi(argv[3]);
        int T_value = std::atoi(argv[4]);
        matched += run_tbpm(test, b_value, E_value, T_value);
    } 
     else if(argc == 6) {
        int b_value = std::atoi(argv[2]); // Convert string argument to integer
        int E_value = std::atoi(argv[3]);
        int T_value = std::atoi(argv[4]);
        int seed = std::atoi(argv[5]);
        std::cout << "Received parameters -> "
              << "test: "   << test
              << ", b_value: " << b_value
              << ", E_value: " << E_value
              << ", T_value: " << T_value
              << ", seed: "    << seed
              << std::endl;
        matched += run_tbpm(test, b_value, E_value, T_value, seed);
    } 

    
    
    matched += run_model(test);

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