#include "diag/config.h"
#include "parser.h"

#ifndef TBPLAS_PYTHON_INTERFACE_DIAG_H
#define TBPLAS_PYTHON_INTERFACE_DIAG_H

namespace tbplas::python {
using tbplas::diag::DiagConfig;

template <typename solver_t>
void set_config(
    solver_t* solver,
    const std::string& config_name)
{
    DiagConfig& config = solver->config;
    InputParser input(config_name);

    // Set config mode to python.
    config.python_mode = true;

    input.read_single("prefix", config.prefix);
    // k-points should be sent as function arguments
    Eigen::Vector3i k_grid_size;
    input.read_eigen("k_grid_size", k_grid_size, false);
    config.k_grid_size = { k_grid_size[0], k_grid_size[1], k_grid_size[2] };
    input.read_single("convention", config.convention);
    input.read_single("dimension", config.dimension);
    input.read_single("num_occ", config.num_occ);
    input.read_single("num_spin", config.num_spin);
    input.read_single("spin_major", config.spin_major);
    input.read_single("e_min", config.e_min);
    input.read_single("e_max", config.e_max);
    input.read_single("e_step", config.e_step);
    input.read_single("sigma", config.sigma);
    input.read_single("smear_func", config.smear_func);
    input.read_single("ldos", config.ldos);
    input.read_set("ldos_orbital_indices", config.ldos_orbital_indices);
    // ka_array, kb_array and k_c are sent as function arguments
    Eigen::Vector3i bz_size;
    input.read_eigen("bz_size", bz_size, false);
    config.bz_size = { bz_size[0], bz_size[1], bz_size[2] };
    input.read_eigen("bz_shift", config.bz_shift, false);
    input.read_single("ham_deriv_analytical", config.ham_deriv_analytical);
    input.read_single("ham_deriv_delta", config.ham_deriv_delta);
    input.read_single("delta", config.delta);
    input.read_single("mu", config.mu);
    input.read_single("temperature", config.temperature);
    input.read_single("back_epsilon", config.back_epsilon);
    input.read_eigen("q_points", config.q_points, true);
    input.read_single("ac_cond_component", config.ac_cond_component);

    // Algorithm parameters
    input.read_single("diag_algo", config.diag_algo);
    input.read_single("pp_algo", config.pp_algo);
    input.read_single("feast_num_states", config.feast_num_states);
    input.read_single("feast_e_min", config.feast_e_min);
    input.read_single("feast_e_max", config.feast_e_max);
    input.read_single("feast_eps", config.feast_eps);
    input.read_single("feast_m0_ratio", config.feast_m0_ratio);
    input.read_single("feast_reuse_results", config.feast_reuse_results);
    Eigen::Matrix2Xi feast_fpm;
    input.read_eigen("feast_fpm", feast_fpm, true);
    for (int j = 0; j < feast_fpm.cols(); ++j) {
        int key = feast_fpm(0, j);
        int val = feast_fpm(1, j);
        config.feast_fpm[key] = val;
    }
}

template <typename solver_t>
void set_k_points(
    solver_t* solver,
    const double* k_points,
    size_t num_kpt)
{
    solver->config.set_k_points(k_points, num_kpt);
}

template <typename solver_t>
void unset_k_points(solver_t* solver)
{
    solver->config.set_k_points(nullptr, 0);
}

template <typename solver_t>
void set_k_array(
    solver_t* solver,
    const double* ka_array,
    size_t num_ka,
    const double* kb_array,
    size_t num_kb,
    double kc)
{
    solver->config.ka_array = Eigen::VectorXd(num_ka);
    vec_copy(ka_array, solver->config.ka_array, num_ka);
    solver->config.kb_array = Eigen::VectorXd(num_kb);
    vec_copy(kb_array, solver->config.kb_array, num_kb);
    solver->config.kc = kc;
}

} // namespace tbplas::python

#endif