#include "esolver_fp.h"

#include "source_estate/cal_ux.h"

#include "source_estate/module_charge/symmetry_rho.h"

#include "source_estate/read_pseudo.h"

#include "source_hamilt/module_ewald/H_Ewald_pw.h"

#include "source_hamilt/module_vdw/vdw.h"

#include "source_io/module_output/cif_io.h"

#include "source_io/module_output/cube_io.h" // use write_vdata_palgrid

#include "source_io/module_json/init_info.h"

#include "source_io/module_json/output_info.h"

#include "source_io/module_output/output_log.h"

#include "source_io/module_output/print_info.h"

#include "source_io/module_chgpot/rhog_io.h"

#include "source_io/module_parameter/parameter.h"

#include "source_pw/module_pwdft/setup_pwrho.h" // mohan 20251005

#include "source_hamilt/module_xc/xc_functional.h" // mohan 20251005

#include "source_io/module_ctrl/ctrl_output_fp.h"

#include "source_io/module_chgpot/write_init.h" // write_chg_init, write_pot_init

namespace ModuleESolver

{

ESolver_FP::ESolver_FP()

{

}

ESolver_FP::~ESolver_FP()

{

//****************************************************

// do not add any codes in this deconstructor funcion

//****************************************************

// mohan add 20251005

pw::teardown_pwrho(this->pw_rho_flag, PARAM.globalv.double_grid, this->pw_rho, this->pw_rhod);

delete this->pelec;

}

void ESolver_FP::before_all_runners(UnitCell& ucell, const Input_para& inp)

{

ModuleBase::TITLE("ESolver_FP", "before_all_runners");

//! 1) read pseudopotentials

elecstate::read_pseudo(GlobalV::ofs_running, ucell);

//! 2) setup pw_rho, pw_rhod, pw_big, sf, and read_pseudopotentials

pw::setup_pwrho(ucell, PARAM.globalv.double_grid, this->pw_rho_flag,

this->pw_rho, this->pw_rhod, this->pw_big, this->classname, inp);

//! 3) setup structure factors

this->sf.set(this->pw_rhod, inp.nbspline);

//! 4) write geometry file

ModuleIO::CifParser::write(PARAM.globalv.global_out_dir + "STRU.cif",

ucell, "# Generated by ABACUS ModuleIO::CifParser", "data_?");

//! 5) init charge extrapolation

this->CE.Init_CE(inp.nspin, ucell.nat, this->pw_rhod->nrxx, inp.chg_extrap);

//! 6) symmetry analysis should be performed every time the cell is changed

if (ModuleSymmetry::Symmetry::symm_flag == 1)

{

ucell.symm.analy_sys(ucell.lat, ucell.st, ucell.atoms, GlobalV::ofs_running);

ModuleBase::GlobalFunc::DONE(GlobalV::ofs_running, "SYMMETRY");

}

ModuleBase::GlobalFunc::DONE(GlobalV::ofs_running, "SETUP UNITCELL");

//! 7) setup k points in the Brillouin zone according to symmetry.

this->kv.set(ucell,ucell.symm, inp.kpoint_file, inp.nspin, ucell.G, ucell.latvec, GlobalV::ofs_running);

ModuleBase::GlobalFunc::DONE(GlobalV::ofs_running, "INIT K-POINTS");

//! 8) print information

ModuleIO::print_parameters(ucell, this->kv, inp);

//! 9) parallel of FFT grid

this->Pgrid.init(this->pw_rhod->nx, this->pw_rhod->ny, this->pw_rhod->nz,

this->pw_rhod->nplane, this->pw_rhod->nrxx, pw_big->nbz, pw_big->bz);

//! 10) calculate the structure factor

this->sf.setup(&ucell, Pgrid, this->pw_rhod);

//! 11) setup the xc functional

XC_Functional::set_xc_type(ucell.atoms[0].ncpp.xc_func);

GlobalV::ofs_running<<XC_Functional::output_info()<<std::endl;

//! 11) initialize the charge density, we need to first set xc_type,

// then we can call chr.allocate()

this->chr.set_rhopw(this->pw_rhod); // mohan add 20251130

const bool kin_den = this->chr.kin_density(); // mohan add 20251202

this->chr.allocate(inp.nspin, kin_den); // mohan move this from setup_estate_pw, 20251128

return;

}

void ESolver_FP::after_scf(UnitCell& ucell, const int istep, const bool conv_esolver)

{

ModuleBase::TITLE("ESolver_FP", "after_scf");

//! Output convergence information

ModuleIO::output_convergence_after_scf(conv_esolver, this->pelec->f_en.etot);

//! Write Fermi energy

ModuleIO::output_efermi(conv_esolver, this->pelec->eferm.ef);

//! Update delta_rho for charge extrapolation

CE.update_delta_rho(ucell, &(this->chr), &(this->sf));

//! print out charge density, potential, elf, etc.

ModuleIO::ctrl_output_fp(ucell, this->pelec, this->pw_big, this->pw_rhod,

this->chr, this->solvent, this->Pgrid, istep);

}

void ESolver_FP::before_scf(UnitCell& ucell, const int istep)

{

ModuleBase::TITLE("ESolver_FP", "before_scf");

// if the cell has changed

if (ucell.cell_parameter_updated)

{

// only G-vector and K-vector are changed due to the change of lattice

// vector FFT grids do not change!!

this->pw_rho->initgrids(ucell.lat0, ucell.latvec, pw_rho->nx, pw_rho->ny, pw_rho->nz);

this->pw_rho->collect_local_pw();

this->pw_rho->collect_uniqgg();

// if double grid used in USPP, update related quantities in dense grid

if (PARAM.globalv.double_grid)

{

this->pw_rhod->initgrids(ucell.lat0, ucell.latvec, pw_rhod->nx, pw_rhod->ny, pw_rhod->nz);

this->pw_rhod->collect_local_pw();

this->pw_rhod->collect_uniqgg();

}

// reset local pseudopotentials

this->locpp.init_vloc(ucell, this->pw_rhod);

ModuleBase::GlobalFunc::DONE(GlobalV::ofs_running, "LOCAL POTENTIAL");

// perform symmetry analysis

if (ModuleSymmetry::Symmetry::symm_flag == 1)

{

ucell.symm.analy_sys(ucell.lat, ucell.st, ucell.atoms, GlobalV::ofs_running);

ModuleBase::GlobalFunc::DONE(GlobalV::ofs_running, "SYMMETRY");

}

// reset k-points

KVectorUtils::set_after_vc(kv, PARAM.inp.nspin, ucell.G);

ModuleBase::GlobalFunc::DONE(GlobalV::ofs_running, "INIT K-POINTS");

}

// charge extrapolation

if (ucell.ionic_position_updated)

{

this->CE.update_all_dis(ucell);

this->CE.extrapolate_charge(&this->Pgrid, ucell, &this->chr, &this->sf,

GlobalV::ofs_running, GlobalV::ofs_warning);

}

//! calculate D2 or D3 vdW

auto vdw_solver = vdw::make_vdw(ucell, PARAM.inp, &(GlobalV::ofs_running));

if (vdw_solver != nullptr)

{

this->pelec->f_en.evdw = vdw_solver->get_energy();

}

//! calculate ewald energy

if (!PARAM.inp.test_skip_ewald)

{

this->pelec->f_en.ewald_energy = H_Ewald_pw::compute_ewald(ucell, this->pw_rhod, this->sf.strucFac);

}

//! set direction of magnetism, used in non-collinear case

elecstate::cal_ux(ucell);

//! output the initial charge density and potential

ModuleIO::write_chg_init(ucell, this->Pgrid, this->chr, this->pelec->eferm, istep, PARAM.inp);

// ModuleIO::write_pot_init(ucell, this->Pgrid, this->pelec, istep, PARAM.inp);

return;

}

void ESolver_FP::iter_finish(UnitCell& ucell, const int istep, int& iter, bool& conv_esolver)

{

//! output charge density in G-space, or if available, kinetic energy density in G-space

if (PARAM.inp.out_chg[0] != -1)

{

if (iter % PARAM.inp.out_freq_elec == 0 || iter == PARAM.inp.scf_nmax || conv_esolver)

{

for (int is = 0; is < PARAM.inp.nspin; is++)

{

this->pw_rhod->real2recip(this->chr.rho_save[is], this->chr.rhog_save[is]);

}

ModuleIO::write_rhog(PARAM.globalv.global_out_dir + PARAM.inp.suffix + "-CHARGE-DENSITY.restart",

PARAM.globalv.gamma_only_pw,

this->pw_rhod,

PARAM.inp.nspin,

ucell.GT,

this->chr.rhog_save,

GlobalV::MY_POOL,

GlobalV::RANK_IN_POOL,

GlobalV::NPROC_IN_POOL);

if (XC_Functional::get_ked_flag())

{

std::vector<std::complex<double>> kin_g_space(PARAM.inp.nspin * this->chr.ngmc, {0.0, 0.0});

std::vector<std::complex<double>*> kin_g;

for (int is = 0; is < PARAM.inp.nspin; is++)

{

kin_g.push_back(kin_g_space.data() + is * this->chr.ngmc);

this->pw_rhod->real2recip(this->chr.kin_r_save[is], kin_g[is]);

}

ModuleIO::write_rhog(PARAM.globalv.global_out_dir + PARAM.inp.suffix + "-TAU-DENSITY.restart",

PARAM.globalv.gamma_only_pw,

this->pw_rhod,

PARAM.inp.nspin,

ucell.GT,

kin_g.data(),

GlobalV::MY_POOL,

GlobalV::RANK_IN_POOL,

GlobalV::NPROC_IN_POOL);

}

}

}

}

void ESolver_FP::after_all_runners(UnitCell& ucell)

{

// print out the final total energy

GlobalV::ofs_running << "\n --------------------------------------------" << std::endl;

GlobalV::ofs_running << std::setprecision(16);

GlobalV::ofs_running << " !FINAL_ETOT_IS " << this->pelec->f_en.etot * ModuleBase::Ry_to_eV << " eV" << std::endl;

GlobalV::ofs_running << " --------------------------------------------\n\n" << std::endl;

}

} // namespace ModuleESolver