#include "relax_driver.h"

#include "source_base/global_file.h"
#include "source_pw/module_pwdft/global.h" // use chr.
#include "source_io/cif_io.h"
#include "source_io/json_output/output_info.h"
#include "source_io/output_log.h"
#include "source_io/print_info.h"
#include "source_io/read_exit_file.h"
#include "source_io/module_parameter/parameter.h"
#include "source_cell/print_cell.h"

void Relax_Driver::relax_driver(
		ModuleESolver::ESolver* p_esolver, 
		UnitCell& ucell,
		const Input_para& inp)
{ 
    ModuleBase::TITLE("Relax_Driver", "relax_driver");
    ModuleBase::timer::tick("Relax_Driver", "relax_driver");

    if (inp.calculation == "relax" || inp.calculation == "cell-relax" )
    {
        if (!inp.relax_new) // traditional relax
        {
            rl_old.init_relax(ucell.nat);
        }
        else // relax new
        {
            rl.init_relax(ucell.nat);
        }
    }

    this->istep = 1;
    int force_step = 1;
    int stress_step = 1;
    bool stop = false;

    while (istep runner(ucell, istep - 1);

        time_t eend = time(nullptr);
        time_t fstart = time(nullptr);
        ModuleBase::matrix force;
        ModuleBase::matrix stress;

        // I'm considering putting force and stress
        // as part of ucell and use ucell to pass information
        // back and forth between esolver and relaxation
        // but I'll use force and stress explicitly here for now

        // calculate the total energy
        this->etot = p_esolver->cal_energy();

        // calculate and gather all parts of total ionic forces
        if (inp.cal_force)
        {
            p_esolver->cal_force(ucell, force);
        }
		else
		{
            // do nothing
		}


        // calculate and gather all parts of stress
        if (inp.cal_stress)
        {
            p_esolver->cal_stress(ucell, stress);
        }
		else
		{
            // do nothing
		}

        if (inp.calculation == "relax" || inp.calculation == "cell-relax")
        {
            if (inp.relax_new)
            {
                stop = rl.relax_step(ucell, force, stress, this->etot);
                // mohan added 2025-07-14
                stress_step = istep+1;
                force_step = 1;
            }
            else
            {
                stop = rl_old.relax_step(istep,
                                         this->etot,
                                         ucell,
                                         force,
                                         stress,
                                         force_step,
                                         stress_step);
            }

            bool need_orb = inp.basis_type == "pw";
            need_orb = need_orb && inp.init_wfc.substr(0, 3) == "nao";
            need_orb = need_orb || inp.basis_type == "lcao";
            need_orb = need_orb || inp.basis_type == "lcao_in_pw";

            std::stringstream ss, ss1;
            ss