/**
 * @file esolver_dp.cpp
#include "source_io/module_parameter/parameter.h"
 * @brief Implementation of ESolver_DP class for DeePMD method.
 *
 * This file contains the implementation of the ESolver_DP class, which is used for solving the energy and forces in a
 * Deep Potential Molecular Dynamics (DeePMD) simulation.
 * DeePMD is a method for training deep neural networks to accurately predict the potential energy surface of a
 * molecular system.
 *
 * For more information about DeePMD, see the following reference:
 *
 * Han Wang, Linfeng Zhang, Jiequn Han, and Roberto Car.
 * "DeePMD-kit: A deep learning package for many-body potential energy representation and molecular dynamics,"
 * Computer Physics Communications 228, 178-184 (2018). https://doi.org/10.1016/j.cpc.2018.03.016
 *
 * @author YuLiu98
 * @date 2023-05-15
 */
#include "esolver_dp.h"

#include "source_base/parallel_common.h"
#include "source_base/timer.h"
#include "source_io/output_log.h"
#include "source_io/cif_io.h"

#include 
#include 
#include 

using namespace ModuleESolver;

void ESolver_DP::before_all_runners(UnitCell& ucell, const Input_para& inp)
{
    dp_potential = 0;
    dp_force.create(ucell.nat, 3);
    dp_virial.create(3, 3);

    ModuleIO::CifParser::write(PARAM.globalv.global_out_dir + "STRU.cif", 
                               ucell, 
                               "# Generated by ABACUS ModuleIO::CifParser",
                               "data_?");

    atype.resize(ucell.nat);

    rescaling = inp.mdp.dp_rescaling;
    fparam = inp.mdp.dp_fparam;
    aparam = inp.mdp.dp_aparam;

#ifdef __DPMD
    /// determine the type map from STRU to DP model
    type_map(ucell);
#endif
}

void ESolver_DP::runner(UnitCell& ucell, const int istep)
{
    ModuleBase::TITLE("ESolver_DP", "runner");
    ModuleBase::timer::tick("ESolver_DP", "runner");

    std::vector cell(9, 0.0);
    cell[0] = ucell.latvec.e11 * ucell.lat0_angstrom;
    cell[1] = ucell.latvec.e12 * ucell.lat0_angstrom;
    cell[2] = ucell.latvec.e13 * ucell.lat0_angstrom;
    cell[3] = ucell.latvec.e21 * ucell.lat0_angstrom;
    cell[4] = ucell.latvec.e22 * ucell.lat0_angstrom;
    cell[5] = ucell.latvec.e23 * ucell.lat0_angstrom;
    cell[6] = ucell.latvec.e31 * ucell.lat0_angstrom;
    cell[7] = ucell.latvec.e32 * ucell.lat0_angstrom;
    cell[8] = ucell.latvec.e33 * ucell.lat0_angstrom;

    std::vector coord(3 * ucell.nat, 0.0);
    int iat = 0;
    for (int it = 0; it < ucell.ntype; ++it)
    {
        for (int ia = 0; ia < ucell.atoms[it].na; ++ia)
        {
            coord[3 * iat] = ucell.atoms[it].tau[ia].x * ucell.lat0_angstrom;
            coord[3 * iat + 1] = ucell.atoms[it].tau[ia].y * ucell.lat0_angstrom;
            coord[3 * iat + 2] = ucell.atoms[it].tau[ia].z * ucell.lat0_angstrom;
            iat++;
        }
    }
    assert(ucell.nat == iat);

#ifdef __DPMD
    std::vector f, v;
    dp_potential = 0;
    dp_force.zero_out();
    dp_virial.zero_out();

    dp.compute(dp_potential, f, v, coord, atype, cell, fparam, aparam);

    // rescale the energy, force, and stress
    const double fact_e = rescaling / ModuleBase::Ry_to_eV;
    const double fact_f = rescaling / (ModuleBase::Ry_to_eV * ModuleBase::ANGSTROM_AU);
    const double fact_v = rescaling / (ucell.omega * ModuleBase::Ry_to_eV);

    dp_potential *= fact_e;
    GlobalV::ofs_running