#ifndef ESOLVER_OF_H
#define ESOLVER_OF_H

#include "esolver_fp.h"
#include "source_base/opt_DCsrch.h"
#include "source_base/opt_TN.hpp"
#include "source_pw/module_ofdft/kedf_manager.h"
#include "source_psi/psi.h"

namespace ModuleESolver
{
class ESolver_OF : public ESolver_FP
{
  public:
    ESolver_OF();
    ~ESolver_OF();

    virtual void before_all_runners(UnitCell& ucell, const Input_para& inp) override;

    virtual void runner(UnitCell& ucell, const int istep) override;

    virtual void after_all_runners(UnitCell& ucell) override;

    virtual double cal_energy() override;

    virtual void cal_force(UnitCell& ucell, ModuleBase::matrix& force) override;

    virtual void cal_stress(UnitCell& ucell, ModuleBase::matrix& stress) override;

  private:
    // ======================= variables ==========================
    // ---------- the kinetic energy density functionals ----------
    KEDF_Manager* kedf_manager_ = nullptr; // KEDF manager, which will be initialized in before_all_runners

    // ----------------- the optimization methods ------------------
    ModuleBase::Opt_CG* opt_cg_ = nullptr;
    ModuleBase::Opt_TN* opt_tn_ = nullptr;
    ModuleBase::Opt_DCsrch* opt_dcsrch_ = nullptr;
    ModuleBase::Opt_CG* opt_cg_mag_ = nullptr; // for spin2 case, under testing

    // ----------------- necessary parameters from INPUT ------------
    std::string of_kinetic_ = "wt";  // Kinetic energy functional, such as TF, VW, WT
    std::string of_method_ = "tn";   // optimization method, include cg1, cg2, tn (default), bfgs
    std::string of_conv_ = "energy"; // select the convergence criterion, potential, energy (default), or both
    double of_tole_ = 2e-6; // tolerance of the energy change (in Ry) for determining the convergence, default=2e-6 Ry
    double of_tolp_ = 1e-5; // tolerance of potential for determining the convergence, default=1e-5 in a.u.
    int max_iter_ = 50;     // scf_nmax

    // ------------------ parameters from other module --------------
    double dV_ = 0;           // volume of one grid point in real space
    double* nelec_ = nullptr; // number of electrons with each spin

    // ----- parameters and arrays used in density optimization -----
    int iter_ = 0;                                // iteration number
    double** pdirect_ = nullptr;                  // optimization direction of phi, which is sqrt(rho)
    std::complex<double>** precip_dir_ = nullptr; // direction in reciprocal space, used when of_full_pw=false.
    double* theta_ = nullptr;                     // step length
    double** pdEdphi_ = nullptr;                  // dE/dphi
    double** pdLdphi_ = nullptr;                  // dL/dphi
    double** pphi_ = nullptr;                     // pphi[i] = ppsi.get_pointer(i), which will be freed in ~Psi().
    char* task_ = nullptr;                        // used in line search
    int tn_spin_flag_ = -1;                       // spin flag used in cal_potential, which will be called by opt_tn
    int max_dcsrch_ = 200;                        // max no. of line search
    int flag_ = -1;                               // flag of TN
    Charge* ptemp_rho_ = nullptr;                 // used in line search
    psi::Psi<double>* psi_ = nullptr;             // sqrt(rho)

    // ----------------- used for convergence check -------------------
    double energy_llast_ = 0;
    double energy_last_ = 0;
    double energy_current_ = 0;
    double normdLdphi_llast_ = 100;
    double normdLdphi_last_ = 100;
    double normdLdphi_ = 100.;

    // ==================== main process of OFDFT ======================
    void before_opt(const int istep, UnitCell& ucell);
    void update_potential(UnitCell& ucell);
    void optimize(UnitCell& ucell);
    void update_rho();
    bool check_exit(bool& conv_esolver);
    void after_opt(const int istep, UnitCell& ucell, const bool conv_esolver);

    // ============================ tools ===============================
    // --------------------- initialize ---------------------------------
    void init_elecstate(UnitCell& ucell);
    void allocate_array();

    // --------------------- calculate physical qualities ---------------
    std::function<void(double*, double*)> bound_cal_potential_;
    void cal_potential_wrapper(double* ptemp_phi, double* rdLdphi);
    void cal_potential(double* ptemp_phi, double* rdLdphi, UnitCell& ucell);
    void cal_dEdtheta(double** ptemp_phi, Charge* temp_rho, UnitCell& ucell, double* ptheta, double* rdEdtheta);
    double cal_mu(double* pphi, double* pdEdphi, double nelec);

    // --------------------- determine the optimization direction -------
    void adjust_direction();
    void check_direction(double* dEdtheta, double** ptemp_phi, UnitCell& ucell);
    void test_direction(double* dEdtheta, double** ptemp_phi, UnitCell& ucell);

    // --------------------- output the necessary information -----------
    void print_info(const bool conv_esolver);

    // --------------------- interface to blas --------------------------
    double inner_product(double* pa, double* pb, int length, double dV = 1)
    {
        double innerproduct = BlasConnector::dot(length, pa, 1, pb, 1);
        innerproduct *= dV;
        return innerproduct;
    }

    // ---------------------- interfaces to optimization methods --------
    void init_opt();
    void get_direction(UnitCell& ucell);
    void get_step_length(double* dEdtheta, double** ptemp_phi, UnitCell& ucell);
};
} // namespace ModuleESolver

#endif
