#include<iostream>
#include"Cvec.h"
#include"Ctridiag.h"
#include"Wavefunction.h"
#include"Potential.h"
#include"Constants.h"
#include"Pulse.h"
#include"TDSE.h"
#include"TrackManager.h"
#include"parse.h"
#include"mpimanager.h"

#define _USE_MATH_DEFINES
#include<math.h>
#include<time.h>
#include<stdio.h>
#include<stdlib.h>
#include<fstream>
#include<sstream>
#include<array>
#include<vector>
#include<filesystem>

namespace fs = std::filesystem;

int main(int argc, char* argv[]) {
	// Init MPI
	MpiManager mpi(argc, argv);
	//mpi.createGroup(1, 1); // Anzahl Gruppe, Anzahl Prozesse pro Gruppe
	mpi.createGroup();
	srand(mpi.rank_world);


	string parameter_file = "Parameter.toml";
	param_struct in = get_param_toml(parameter_file);
	if(mpi.rank_group == 0) std::cout << "Param input completed." << std::endl;

	// ----- Speicherort -----

	MPI_Barrier(mpi.comm);

	for (double loop_index = mpi.rank_group; loop_index <= in.loop_number; loop_index += mpi.nproc_group) {
		double loop = in.loop_start + loop_index * in.loop_step;

		string run = in.run;
		string input_field_filename = in.input;
		
		//Allgemeine Parameter
		if (in.make_loop == 1) {
			//if (mpi.rank_group == 0) std::cout << "Run " << loop << std::endl;
			std::cout << "Run " << loop << " by MPI Rank " << mpi.rank_group << std::endl;

			if (in.loop_var == "FWHM") {
				in.FWHM[0] = loop;
				double help = 2.5 * in.FWHM[0];
				if (help < 15) help = 15;
				in.t_prop = Cvec(-help, in.prop_step, help);
			}
			else if (in.loop_var == "wavelength") in.wavelength[0] = loop;
			else if (in.loop_var == "intensity") in.intensity[0] = pow(10, loop);
			else if (in.loop_var == "gamma0") in.gamma0 = loop;
			else if (in.loop_var == "phase") in.phase[0] = loop * M_PI;
			else if (in.loop_var == "input_field") {
				char filename_help[100];
				sprintf(filename_help, "%s%0.2f", input_field_filename.c_str(), loop);
				input_field_filename = filename_help;
				replace(input_field_filename.begin(), input_field_filename.end(), '.', '_');
			}
			else if (in.loop_var == "t0") in.t0[0] = loop;


			string trimmedString = to_string(loop).substr(0, to_string(loop).find(".") + 2 + 1);
			run = in.run + "/" + in.loop_var + "_" + trimmedString;	// File Name
		}

		string path_directory = "Output/" + run + "/";
		fs::create_directories(path_directory);
		fs::copy(parameter_file, path_directory, fs::copy_options::overwrite_existing);

		TDSE tdse(in.dx_au);
		tdse.init_xlist_au(in.x_list);
		tdse.init_tdse();
		tdse.init_tlist(in.t_prop);
		tdse.path = path_directory;
		tdse.x_right = in.x_right;
		tdse.write_spec_evol = in.write_spec_evol;
		tdse.mpi = &mpi;

		if (in.potential_type == "box") tdse.add_box(-in.V0, in.width, in.x0);
		else if (in.potential_type == "fermi") tdse.add_fermibox(-in.V0, in.smooth, in.width, in.x0);
		else if (in.potential_type == "personal") tdse.set_box_benutzer(-in.V0, in.smooth, in.width, in.x0);
		else if (in.potential_type == "coulomb") tdse.add_Coulomb(0.0, 2.0, 3.5);
		else if (in.potential_type == "find_box") tdse.find_box_pot(in.W, in.E_F);

		tdse.set_Hamiltonian();
		if (in.use_parallel == 1) tdse.imaginary_time_prop_parallel(in.maximal_states);
		else tdse.imaginary_time_prop_Vmin(in.maximal_states, in.V0); // use 1.5*V for many box states -> ensure that first state found is groundstate

		tdse.set_psi_start(in.state_numbers);

		if (in.laser == 1) {
			for (int i = 0; i < in.wavelength.size(); i++) tdse.add_pulse_vp(in.wavelength[i], in.intensity[i], in.t0[i], in.phase[i], in.FWHM[i], in.SC_laser[i]);
		}
		if (in.import_field == 1) tdse.input_field(input_field_filename);

		if (in.enhancenment == 1) tdse.set_gamma(in.gamma0, in.lambda_NF);
		if(in.propagation == 1) tdse.propagation(tdse.Psi_start, in.write_spec, in.write_step_x, in.write_step_t, input_field_filename);

		if (in.outprojection == 1) {
			for (auto& psi : tdse.Psi_states) {
				tdse.outprojection(psi);
			}
		}

		if (in.calc_spec == 1) tdse.calc_spec(tdse.Psi, in.Elist_spec, in.order, in.gamma, in.half_spec);

		if (mpi.rank_group == 0) std::cout << "Save data..." << std::endl;
		tdse.save_all();

	}

	MPI_Barrier(mpi.comm);
	return 0;

};
