#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
    // Reference: tbplas/tests/test_units/test_ac.py
    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
    // Reference: tbplas/tests/test_units/test_dp.py and test_eps.py
    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
    // Reference: tbplas/tests/test_units/test_hall.py
    // 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";
    }
}

void test_dos_tbg(int pool_size_param, int T_param, int seed)
{
    // Model
    model_t model = make_tbg(31);
    // dim_t dim = {7,7,1};
    // dim_t dim = {3,3,1};
    dim_t dim = {20,20,1};
    model = extend_prim_cell(model,dim);
    Timer timer;
    std::cout << "Model num_orbitals: " << model.get_num_orbitals()
            << " num_hopping: " << model.get_num_hopping() << std::endl;

    // 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 = T_param; //4096
    solver.config.algo = "cpu";
    solver.config.ldos = false;
    solver.config.ldos_orbital_indices = { 1 };

    solver.config.seed = seed;


    // std::vector<int> specific_ts_values = {64,128,256,512,1024,2048,4096};
    // // std::vector<int> specific_b_values = {16};
    // for (int ts : specific_ts_values) {
    //     solver.config.num_time_steps = ts;
    // Calculation
    std::cout << "--- Running with TBPM  ---" << std::endl;
    solver.config.algo = "cpu";
    solver.config.prefix = std::string("dosSP") + std::to_string(seed);
    if(pool_size_param==0){
        timer.tic("dos");
        solver.calc_corr_dos();
        timer.toc("dos");
    }


    // Parameters
    std::cout << "--- Running with TBPM-fast ---" << std::endl;
    solver.config.algo = "cpu_fast";
    solver.config.prefix = std::string("dosMoment") + std::to_string(seed);
    if(pool_size_param>0){
        timer.tic("dos_fast");
        solver.calc_corr_dos();
        timer.toc("dos_fast");
    }
    // }

    if (solver.is_master()) {
        timer.report_total_time();
    }

}

void test_qe_tbg(int pool_size_param, int E_param, int T_param, int seed)
{
    // Model setup 
    model_t model = make_tbg(31);
    dim_t dim = {20,20,1};
    // dim_t dim = {3,3,1};
    // dim_t dim = {1, 1, 1}; // Using dim = {1,1,1} as in your example
    model = extend_prim_cell(model, dim);

    Timer timer;
    

    std::cout << "Model num_orbitals: " << model.get_num_orbitals()
            << " num_hopping: " << model.get_num_hopping() << std::endl;
    // --- Common Configuration Setup ---
    // Create the solver once
    TBPMSolver<model_t> solver(model);

    // Set all common parameters that apply to both runs
    solver.config.dimension = 2;
    solver.config.num_random_samples = 1;
    solver.config.rescale = -0.1;
    solver.config.num_time_steps = T_param;
    Eigen::VectorXd qe_energies = Eigen::VectorXd::LinSpaced(E_param, -1.0, 1.0);
    solver.config.qe_energies = qe_energies;

    solver.config.seed = seed;
    solver.config.save_data = false;
    // std::vector<int> specific_ne_values = {2,16,64,128,256,512,1024};
    // for (int ne : specific_ne_values) {
    //     Eigen::VectorXd qe_energies = Eigen::VectorXd::LinSpaced(ne, -1.0, 1.0);
    //     solver.config.qe_energies = qe_energies;
    //     std::cout << "N_E: "<< ne << std::endl;
    // std::vector<int> specific_ts_values = {64,128,256,512,1024,2048,4096};
    // std::vector<int> specific_b_values = {16};
    // for (int ts : specific_ts_values) {
        // solver.config.num_time_steps = ts;
    solver.config.algo = "cpu"; // Assuming ALGO is defined elsewhere

    // --- First Run: Old Method ---
    std::cout << "--- Running with TBPM  ---" << std::endl;
    solver.config.prefix = std::string("qeSP") + std::to_string(seed);
    if(pool_size_param==0){
        timer.tic("qe");
        solver.calc_quasi_eigenstates(); // Call the calculation
        timer.toc("qe");
    }

    //--- Second Run: Chebyshev Polynomial Method ---
    std::cout << "--- Running with Chebyshev Polynomial Method ---" << std::endl;
    solver.config.algo = "cpu_fast";
    // solver.config.prefix = "fastpoly";          // Update prefix for output files
    solver.config.qe_method = "cheby_poly"; // Set the specific method
    solver.config.prefix = std::string("qeE") + std::to_string(seed);
    if(pool_size_param<0){
        timer.tic("qe_cheby_poly");
        solver.calc_quasi_eigenstates(); // Call the calculation
        timer.toc("qe_cheby_poly");
    }

    // --- Third Run: Block-Time Propagation Method ---
    std::cout << "\n--- Running with Block-Time Propagation Method ---" << std::endl;
    solver.config.algo = "cpu_fast";
    solver.config.qe_method = "block_time"; // Switch to the other method
    // solver.config.pool_size = 32;           // Set pool_size, specific to block_time
    // solver.config.prefix = "fastblock";          // Update prefix for output files
    solver.config.prefix = std::string("qeT") + std::to_string(seed);
    // std::vector<int> specific_b_values = {2, 4, 8, 10, 12, 14, 16, 32, 64};
    // std::vector<int> specific_b_values = {2, 8, 16, 24, 27, 30, 32, 64};
    // std::vector<int> specific_b_values = {2, 8, 16, 32, 42, 54, 60, 64, 80};
    // std::vector<int> specific_b_values = {2, 4, 8, 16, 19, 20, 24, 25, 32, 64};
    // std::vector<int> specific_b_values = {2, 8, 16, 24, 30, 32, 36, 45, 64};
    //std::vector<int> specific_b_values = {17, 18, 19, 20, 21, 22, 23, 32, 128};
    // std::vector<int> specific_b_values;
    // for (int i = 2; i <= 64; i += 2) {
    //     specific_b_values.push_back(i);
    // }
    // for (int b : specific_b_values) {
        // solver.config.pool_size = b;
    if(pool_size_param>0){
        solver.config.pool_size = pool_size_param;
        timer.tic("qe_block_time");
        solver.calc_quasi_eigenstates(); // Call the calculation again with modified config
        timer.toc("qe_block_time");
    }
    // } // End of loop over specific_ts_values

    if (solver.is_master()) {
        timer.report_total_time();
    }

}

void test_dc_cond_tbg(int pool_size_param, int E_param, int T_param, int seed)
{
    // Model
    model_t model = make_tbg(31);
    dim_t dim = {20,20,1};
    // dim_t dim = {3,3,1};
    // dim_t dim = {1,1,1};
    model = extend_prim_cell(model,dim);
    Timer timer;
    // Eigen::VectorXd qe_energies = Eigen::VectorXd::LinSpaced(21, -1.0, 1.0);

    std::cout << "Model num_orbitals: " << model.get_num_orbitals()
            << " num_hopping: " << model.get_num_hopping() << std::endl;

    // 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 = 1024; //4096
    double e_shift = 0.79; //if dc_ne, 0.78; others, 0.79+3-> NE=5
    double e_m = 0.01;
    solver.config.dc_cond_energy_limits = {
        {e_shift - e_m * static_cast<double>(E_param),
         e_shift + e_m * static_cast<double>(E_param)}};
    solver.config.dc_cond_num_time_steps = T_param;
    solver.config.seed = seed;


    solver.config.algo = "cpu";
    // --- First Run: Old Method ---
    std::cout << "--- Running with TBPM  ---" << std::endl;
    solver.config.prefix = std::string("dcSP") + std::to_string(seed);
    if(pool_size_param==0){
        timer.tic("dc_cond");
        solver.calc_corr_dc_cond();
        timer.toc("dc_cond");
    }
    //--- Second Run: Chebyshev Polynomial Method ---
    std::cout << "--- Running with Chebyshev Polynomial Method ---" << std::endl;
    solver.config.algo = "cpu_fast";
    // solver.config.prefix = "fastpoly";          // Update prefix for output files
    // solver.config.pool_size = 32;   
    solver.config.qe_method = "cheby_poly"; // Set the specific method

    // std::vector<int> specific_b_values = {2, 4, 8, 10, 12, 14, 16, 32, 64};
    // std::vector<int> specific_b_values = {2, 8, 16, 24, 27, 30, 32, 64};
    // std::vector<int> specific_b_values = {2, 8, 16, 32, 42, 54, 60, 64, 80};
    // std::vector<int> specific_b_values = {2, 4, 8, 16, 19, 20, 24, 25, 32, 64};
    // std::vector<int> specific_b_values = {2, 8, 16, 24, 30, 32, 36, 45, 64};
    // std::vector<int> specific_b_values = {17, 18, 19, 20, 21, 22, 23, 32, 128};
    // std::vector<int> specific_b_values;
    // for (int i = 2; i <= 64; i += 2) {
    //     specific_b_values.push_back(i);
    // }
    // // std::vector<int> specific_b_values = {16};
    // for (int b : specific_b_values) {
    //     solver.config.pool_size = b;
    if(pool_size_param>0){
        solver.config.dc_method = "block_time";
        solver.config.pool_size = pool_size_param;
        solver.config.prefix = std::string("dcT") + std::to_string(seed);
        timer.tic("dc_cond_cheby_poly");
        solver.calc_corr_dc_cond();
        timer.toc("dc_cond_cheby_poly");
    }

    if(pool_size_param<0){
        // solver.config.prefix = "fastcheby";
        solver.config.dc_method = "cheby_poly";
        solver.config.pool_size = pool_size_param;
        solver.config.prefix = std::string("dcE") + std::to_string(seed);
        timer.tic("dc_cond_cheby_poly_cheby");
        solver.calc_corr_dc_cond();
        timer.toc("dc_cond_cheby_poly_cheby");
    }


    // --- Third Run: Block-Time Propagation Method ---
    // std::cout << "\n--- Running with Block-Time Propagation Method ---" << std::endl;
    // solver.config.algo = "cpu_fast";
    // solver.config.qe_method = "block_time"; // Switch to the other method
    // solver.config.pool_size = 32;           // Set pool_size, specific to block_time
    // solver.config.prefix = "fastblock";          // Update prefix for output files

    // timer.tic("dc_cond_block_time");
    // // solver.calc_corr_dc_cond();
    // timer.toc("dc_cond_block_time");

    if (solver.is_master()) {
        timer.report_total_time();
    }
}


void test_ac_cond_tbg(int pool_size_param, int T_param, int seed)
{
    // Model
    model_t model = make_tbg(31);
     dim_t dim = {20,20,1};
    //dim_t dim = {3,3,1};
    // dim_t dim = {1,1,1};
    model = extend_prim_cell(model,dim);
    Timer timer;

    std::cout << "Model num_orbitals: " << model.get_num_orbitals()
            << " num_hopping: " << model.get_num_hopping() 
            << " nnz(ham)/num_orbitals: " << (2*model.get_num_hopping())/ model.get_num_orbitals() + 1
            << std::endl;

    // 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 = T_param; //8;//1024; //4096

    solver.config.seed = seed;

    solver.config.algo = "cpu";
    // solver.config.opt_pool_size = false;
    // solver.config.opt_pool_size = true;
    solver.config.num_internal_runs = 1000;
    solver.config.num_external_runs = 50;


    // --- First Run: Old Method ---
    std::cout << "--- Running with TBPM  ---" << std::endl;
    solver.config.prefix = std::string("ocSP") + std::to_string(seed);
    if(pool_size_param==0){
        timer.tic("ac_cond");
        solver.calc_corr_ac_cond();
        timer.toc("ac_cond");
    }

    //--- Second Run: Chebyshev Polynomial Method ---
    std::cout << "--- Running with TBPM-fast ---" << std::endl;
    solver.config.algo = "cpu_fast";
    // solver.config.prefix = "fast";
    solver.config.opt_pool_size = false;
    solver.config.prefix = std::string("ocT") + std::to_string(seed);
    // solver.config.pool_size = 32; 
    //
    // std::vector<int> specific_b_values = {2, 4, 8, 10, 12, 14, 16, 32, 64};
    // std::vector<int> specific_b_values = {2, 8, 16, 24, 27, 30, 32, 64};
    // std::vector<int> specific_b_values = {2, 8, 16, 32, 42, 54, 60, 64, 80};
    // std::vector<int> specific_b_values = {2, 4, 8, 16, 19, 20, 24, 25, 32, 64};
    // std::vector<int> specific_b_values = {2, 8, 16, 24, 30, 32, 36, 45, 64};
    // std::vector<int> specific_b_values = {17, 18, 19, 20, 21, 22, 23, 32, 128};
    // std::vector<int> specific_b_values = {25};
    // std::vector<int> specific_b_values;
    // for (int i = 2; i <= 64; i += 2) {
    //     specific_b_values.push_back(i);
    // }
    // for (int b : specific_b_values) {
    // // for(int b=2; b < 25; ++b) {
    //     solver.config.pool_size = b;
    if(pool_size_param>0){
        solver.config.pool_size = pool_size_param;
        timer.tic("ac_fast");
       solver.calc_corr_ac_cond();
        timer.toc("ac_fast");
    }

    if (solver.is_master()) {
        timer.report_total_time();
    }
}

void test_dyn_pol_tbg(int pool_size_param, int T_param, int seed)
{
    // Model
    model_t model = make_tbg(31);
    dim_t dim = {20,20,1};
    // dim_t dim = {3,3,1};
    // dim_t dim = {1,1,1};
    model = extend_prim_cell(model,dim);
    Timer timer;
    // Eigen::VectorXd qe_energies = Eigen::VectorXd::LinSpaced(21, -1.0, 1.0);
    Eigen::Matrix3d recip_vec = model.get_reciprocal_lattice();
    std::cout << "Model num_orbitals: " << model.get_num_orbitals()
            << " num_hopping: " << model.get_num_hopping() << std::endl;

    // 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 = T_param; //4096

    int num_q_points_for_test = 1;
    solver.config.dyn_pol_q_points = Eigen::Matrix3Xd::Zero(3, num_q_points_for_test);

    // Get the first reciprocal lattice vector
    Eigen::Vector3d b1 = recip_vec.row(0).transpose();
    // Generate points along the b1 direction, from 0 to just before b1
    for (int i = 0; i < num_q_points_for_test; ++i) {
        double fraction = static_cast<double>(i+1) / 2./num_q_points_for_test;
        solver.config.dyn_pol_q_points.col(i) = fraction * b1;
    }

    solver.config.seed = seed;
    // solver.config.qe_energies = qe_energies;
    solver.config.algo = "cpu";

    // --- First Run: Old Method ---
    std::cout << "--- Running with TBPM  ---" << std::endl;
    solver.config.prefix = std::string("dpSP") + std::to_string(seed);
    if(pool_size_param==0){
        timer.tic("dyn_pol");
        solver.calc_corr_dyn_pol();
        timer.toc("dyn_pol");
    }

    //--- Second Run: Chebyshev Polynomial Method ---
    std::cout << "--- Running with TBPM-fast ---" << std::endl;
    solver.config.algo = "cpu_fast";
    // solver.config.prefix = "fast";
    solver.config.prefix = std::string("dpT") + std::to_string(seed);
    // solver.config.pool_size = 32;   
    // std::vector<int> specific_b_values = {2, 4, 8, 10, 12, 14, 16, 32, 64};
    // std::vector<int> specific_b_values = {2, 8, 16, 24, 27, 30, 32, 64};
    // std::vector<int> specific_b_values = {2, 8, 16, 32, 42, 54, 60, 64, 80};
    // std::vector<int> specific_b_values = {2, 4, 8, 16, 19, 20, 24, 25, 32, 64};
    // std::vector<int> specific_b_values = {2, 8, 16, 24, 30, 32, 36, 45, 64};
   // std::vector<int> specific_b_values = {17, 18, 19, 20, 21, 22, 23, 32, 128};
    // std::vector<int> specific_b_values = {25};
    // std::vector<int> specific_b_values;
    // for (int i = 2; i <= 64; i += 2) {
    //     specific_b_values.push_back(i);
    // }
    // for (int b : specific_b_values) {
    //     solver.config.pool_size = b;
    if(pool_size_param>0){
        solver.config.pool_size = pool_size_param;
        timer.tic("dyn_pol_fast");
        solver.calc_corr_dyn_pol();
        timer.toc("dyn_pol_fast");
    }

    if (solver.is_master()) {
        timer.report_total_time();
    }
}

void test_cd_tbg(int pool_size_param, int T_param, int seed)
{
    // Model
    model_t model = make_tbg(31);
    dim_t dim = {20,20,1};
    // dim_t dim = {3,3,1};
    // dim_t dim = {1,1,1};
    model = extend_prim_cell(model,dim);
    Timer timer;
    // Eigen::VectorXd qe_energies = Eigen::VectorXd::LinSpaced(21, -1.0, 1.0);
    Eigen::Matrix3d recip_vec = model.get_reciprocal_lattice();
    std::cout << "Model num_orbitals: " << model.get_num_orbitals()
            << " num_hopping: " << model.get_num_hopping() << std::endl;

    // 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 = T_param; //4096
    solver.config.num_electron = 11908*20*20;
    solver.config.mu = 0.78;
    solver.config.occ_width = 12.;

    solver.config.seed = seed;

    // solver.config.qe_energies = qe_energies;
    solver.config.algo = "cpu";

    // --- First Run: Old Method ---
    std::cout << "--- Running with TBPM  ---" << std::endl;
    solver.config.prefix = std::string("cdSP") + std::to_string(seed);
    if(pool_size_param==0){
        timer.tic("cd");
        solver.calc_charge_density();
        timer.toc("cd");
    }

    //--- Second Run: Chebyshev Polynomial Method ---
    std::cout << "--- Running with TBPM-fast ---" << std::endl;
    solver.config.algo = "cpu_fast";
    solver.config.prefix = std::string("cdT") + std::to_string(seed);
    // solver.config.prefix = "fast";
    // solver.config.pool_size = 4;   

    // std::vector<int> specific_b_values = {2, 4, 8, 10, 12, 14, 16, 32, 64};
    // std::vector<int> specific_b_values = {2, 8, 16, 24, 27, 30, 32, 64};
    // std::vector<int> specific_b_values = {2, 8, 16, 32, 42, 54, 60, 64, 80};
    // std::vector<int> specific_b_values = {2, 4, 8, 16, 19, 20, 24, 25, 32, 64};
    // std::vector<int> specific_b_values = {2, 8, 16, 24, 30, 32, 36, 45, 64};
    // std::vector<int> specific_b_values = {17, 18, 19, 20, 21, 22, 23, 32, 128};
    // std::vector<int> specific_b_values = {128};
    // std::vector<int> specific_b_values = {16};
    // std::vector<int> specific_b_values;
    // for (int i = 2; i <= 64; i += 2) {
    //     specific_b_values.push_back(i);
    // }
    // for (int b : specific_b_values) {
    //     solver.config.pool_size = b;
    if(pool_size_param>0){
        solver.config.pool_size = pool_size_param;
        timer.tic("cd_fast");
        solver.calc_charge_density();
        timer.toc("cd_fast");
    }

    if (solver.is_master()) {
        timer.report_total_time();
    }
}


} // namespace tbplas::tests
