#include "relax_driver.h"
#include "source_base/global_file.h"
#include "source_io/module_output/cif_io.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_output/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,
 std::ofstream& ofs_running)
{
 ModuleBase::TITLE("Relax_Driver", "relax_driver");
 ModuleBase::timer::start("Relax_Driver", "relax_driver");

 this->init_relax(ucell.nat, inp);

 // steps[0]: istep (main iteration step)
 // steps[1]: force_step
 // steps[2]: stress_step
 std::vector<int> steps = {0, 1, 1};

 // Main iteration loop for relaxation calculations
 // For scf/nscf calculations, relax_step returns true immediately,
 // so the loop exits after one iteration
 while (steps[0] < inp.relax_nmax)
 {
 ModuleBase::matrix force(ucell.nat, 3);
 ModuleBase::matrix stress(3, 3);
 double etot = 0.0;

 this->iter_info(steps, inp);
 this->esolve(steps[0], p_esolver, ucell, inp, force, stress, etot);
 bool converged = this->relax_step(steps, p_esolver, ucell, inp, force, stress, etot, ofs_running);
 this->json_out(p_esolver, ucell, inp, force, stress);

 // Check stop conditions
 if (converged)
 {
 // Relaxation converged, exit loop immediately
 break;
 }
 else if (ModuleIO::read_exit_file(GlobalV::MY_RANK, "EXIT", ofs_running))
 {
 // EXIT file detected, exit loop
 break;
 }

 ++steps[0];
 }

 this->final_out(steps[0], ucell, inp);

 ModuleBase::timer::end("Relax_Driver", "relax_driver");
 return;
}

void Relax_Driver::init_relax(const int nat, const Input_para& inp)
{
 if (inp.calculation == "relax" || inp.calculation == "cell-relax")
 {
 if (!inp.relax_new)
 {
 this->rl_old.init_relax(nat);
 }
 else
 {
 this->rl.init_relax(nat);
 }
 }
}

void Relax_Driver::iter_info(const std::vector<int>& steps, const Input_para& inp)
{
 if (inp.out_level == "ie"
 && (inp.calculation == "relax"
 || inp.calculation == "cell-relax"
 || inp.calculation == "scf"
 || inp.calculation == "nscf")
 && (inp.esolver_type != "lr"))
 {
 ModuleIO::print_screen(steps[2], steps[1], steps[0]+1);
 }

#ifdef __RAPIDJSON
 Json::init_output_array_obj();
#endif
}

void Relax_Driver::esolve(const int istep,
		ModuleESolver::ESolver* p_esolver,
		UnitCell& ucell,
		const Input_para& inp,
		ModuleBase::matrix& force,
		ModuleBase::matrix& stress,
		double& etot)
{
 p_esolver->runner(ucell, istep);

 etot = p_esolver->cal_energy();

 if (inp.cal_force)
 {
 p_esolver->cal_force(ucell, force);
 }

 if (inp.cal_stress)
 {
 p_esolver->cal_stress(ucell, stress);
 }
}

bool Relax_Driver::relax_step(std::vector<int>& steps,
		ModuleESolver::ESolver* p_esolver,
		UnitCell& ucell,
		const Input_para& inp,
		const ModuleBase::matrix& force,
		const ModuleBase::matrix& stress,
		const double etot,
		std::ofstream& ofs_running)
{
 // Guard: For non-relaxation calculations (scf, nscf, etc.), return true immediately
 // to ensure the main loop exits after one iteration. This provides robustness
 // even if relax_nmax is set to a large value.
 if (inp.calculation != "relax" && inp.calculation != "cell-relax")
 {
 return true;
 }

 bool converged = false;

 if (inp.relax_new)
 {
 converged = this->rl.relax_step(ucell, force, stress, etot, ofs_running);
	// stress step +1
 steps[2]++;
	// fix force step to 1
 steps[1] = 1;
 }
 else
 {
 converged = this->rl_old.relax_step(steps[0]+1, etot, ucell, force,
			stress, steps[1], steps[2], ofs_running);
 }

 this->stru_out(steps[0], ucell, inp);

 ModuleIO::output_after_relax(converged, p_esolver->conv_esolver, ofs_running);

 return converged;
}

void Relax_Driver::stru_out(const int istep, UnitCell& ucell, const Input_para& inp)
{
 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 << PARAM.globalv.global_out_dir << "STRU_ION_D";

 unitcell::print_stru_file(ucell,
 ucell.atoms,
 ucell.latvec,
 ss.str(),
 inp.nspin,
 true,
 inp.calculation == "md",
 inp.out_mul,
 need_orb,
 PARAM.globalv.deepks_setorb,
 GlobalV::MY_RANK);

 if (inp.out_stru)
 {
 if (inp.out_freq_ion == 0 || istep % inp.out_freq_ion == 0)
 {
 ss1 << PARAM.globalv.global_out_dir << "STRU_ION";
 ss1 << istep+1 << "_D";
 unitcell::print_stru_file(ucell,
 ucell.atoms,
 ucell.latvec,
 ss1.str(),
 inp.nspin,
 true,
 inp.calculation == "md",
 inp.out_mul,
 need_orb,
 PARAM.globalv.deepks_setorb,
 GlobalV::MY_RANK);

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

void Relax_Driver::json_out(ModuleESolver::ESolver* p_esolver, UnitCell& ucell, const Input_para& inp, const ModuleBase::matrix& force, const ModuleBase::matrix& stress)
{
#ifdef __RAPIDJSON
 Json::add_output_energy(p_esolver->cal_energy() * ModuleBase::Ry_to_eV);

 double unit_transform = ModuleBase::RYDBERG_SI / pow(ModuleBase::BOHR_RADIUS_SI, 3) * 1.0e-8;
 double fac = ModuleBase::Ry_to_eV / 0.529177;
 Json::add_output_cell_coo_stress_force(&ucell, force, fac, stress, unit_transform);
#endif
}

void Relax_Driver::final_out(const int istep, UnitCell& ucell, const Input_para& inp)
{
 if (inp.calculation != "relax" && inp.calculation != "cell-relax")
 {
 return;
 }

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

 if (istep == inp.relax_nmax)
 {
 std::cout << "\n ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~" << std::endl; 
 std::cout << " Geometry relaxation stops here due to reaching the maximum " << std::endl;
 std::cout << " relaxation steps. More steps are needed to converge the results " << std::endl;
 std::cout << " ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~" << std::endl; 
 }
 else
 {
 std::cout << "\n ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~" << std::endl; 
 std::cout << " Geometry relaxation thresholds are reached within " << istep << " steps." << std::endl; 
 std::cout << " ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~" << std::endl; 
 }

 if (inp.relax_nmax == 0)
 {
 std::cout << "-----------------------------------------------" << std::endl;
 std::cout << " relax_nmax = 0, DRY RUN TEST SUCCEEDS :)" << std::endl;
 std::cout << "-----------------------------------------------" << std::endl;
 }
}