#include "esolver_lj.h"
#include "source_io/module_parameter/parameter.h"
#include "source_cell/module_neighbor/sltk_atom_arrange.h"
#include "source_cell/module_neighbor/sltk_grid_driver.h"
#include "source_io/output_log.h"
#include "source_io/cif_io.h"


namespace ModuleESolver
{

void ESolver_LJ::before_all_runners(UnitCell& ucell, const Input_para& inp)
{
    lj_potential = 0;
    lj_force.create(ucell.nat, 3);
    lj_virial.create(3, 3);

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

    // determine the maximum rcut and lj_rcut
    rcut_search_radius(ucell.ntype, inp.mdp.lj_rcut);

    // determine the LJ parameters
    set_c6_c12(ucell.ntype, inp.mdp.lj_rule, inp.mdp.lj_epsilon, inp.mdp.lj_sigma);

    // calculate the energy shift so that LJ energy is zero at rcut
    cal_en_shift(ucell.ntype, inp.mdp.lj_eshift);
}

void ESolver_LJ::runner(UnitCell& ucell, const int istep)
{
    Grid_Driver grid_neigh(PARAM.inp.test_deconstructor, PARAM.inp.test_grid);
    atom_arrange::search(PARAM.globalv.search_pbc,
                         GlobalV::ofs_running,
                         grid_neigh,
                         ucell,
                         search_radius,
                         PARAM.inp.test_atom_input);

    double distance = 0.0;
    int index = 0;

    // Important! potential, force, virial must be zero per step
    lj_potential = 0;
    lj_force.zero_out();
    lj_virial.zero_out();

    ModuleBase::Vector3 tau1, tau2, dtau;
    for (int it = 0; it < ucell.ntype; ++it)
    {
        Atom* atom1 = &ucell.atoms[it];
        for (int ia = 0; ia < atom1->na; ++ia)
        {
            tau1 = atom1->tau[ia];
            grid_neigh.Find_atom(ucell, tau1, it, ia);
            for (int ad = 0; ad < grid_neigh.getAdjacentNum(); ++ad)
            {
                tau2 = grid_neigh.getAdjacentTau(ad);
                int it2 = grid_neigh.getType(ad);
                dtau = (tau1 - tau2) * ucell.lat0;
                distance = dtau.norm();
                if (distance < lj_rcut(it, it2))
                {
                    lj_potential += LJ_energy(distance, it, it2) - en_shift(it, it2);
                    ModuleBase::Vector3 f_ij = LJ_force(dtau, it, it2);
                    lj_force(index, 0) += f_ij.x;
                    lj_force(index, 1) += f_ij.y;
                    lj_force(index, 2) += f_ij.z;
                    LJ_virial(f_ij, dtau);
                }
            }
            index++;
        }
    }

    lj_potential /= 2.0;
    GlobalV::ofs_running