#include<iostream>
#include<fstream>
#include<vector>
#include<algorithm>
#include"TDSE.h"
#include"Save_Array.h"
#include"Constants.h"
#include"io.h"

using namespace std;

TDSE::TDSE(double& dxin) :Psi(dxin), V(dxin), V_las(dxin), V_imag_vec(dxin), Psi_start(dxin), help_wave(dxin), wave_interm(dxin) {
	
	/*Psi_time = new Creal* [20000];
	for (int i = 0; i < 20000 ; i++) {
		Psi_time[i] = new Creal[40000];
	}
	Population_time = new Creal* [2];
	for (int i = 0; i < 1; i++) {
		Population_time[i] = new Creal[40000];
	}*/
}
TDSE::~TDSE(){

	/*for (int i = 0; i < 1; i++) delete[] Psi_time[i];
	delete[] Psi_time;
	Psi_time = NULL;

	for (int i = 0; i < 1; i++) delete[] Population_time[i];
	delete[] Population_time;
	Population_time = NULL;*/
}


void TDSE::init_xlist(double start, double step, double end) {
	int n = (int)(end + step - start) / step;
	xlist.init(n);
	for (int i = 0; i < xlist.get_size(); i++) xlist[i] = start + i * step;
	xlist = xlist * 1e-10 / au_L; // au
}

void TDSE::init_xlist_au(double start, double step, double end) {
	int n = (int)(end + step - start) / step;
	xlist.init(n);
	for (int i = 0; i < xlist.get_size(); i++) xlist[i] = start + i * step;
}

void TDSE::init_xlist_au(Cvec x_list) {
	xlist.init(x_list.get_size());
	xlist = x_list;
}

void TDSE::init_tlist(double start, double step, double end) {
	double n_double = (end + step - start) / step;
	int n = (int)round(n_double);
	tlist.init(n);
	for (int i = 0; i < tlist.get_size(); i++) tlist[i] = start + i * step;
	tlist = tlist / (au_T * 1e15); // au
}

void TDSE::init_tlist(Cvec t_input) {
	int n = t_input.get_size();
	tlist.init(n);
	tlist = t_input;
	tlist = tlist / (au_T * 1e15); // au
}

void TDSE::init_tdse() {
	Psi.init(xlist.get_size());
	Psi.init_random();

	Psi_start.init(xlist.get_size());
	Psi_start = 0.0;

	V.init(xlist.get_size());

	Creal one(1, 0);
	Creal m_two(-2, 0);
	Creal zero(0, 0);
	Creal m_one(-1, 0);

	V_tri.init(xlist.get_size());
	H_0.init(xlist.get_size());
	H_0.sub = one;
	H_0.diag = m_two;
	H_0.sup = one;
	double help = -2 * Psi.get_dx() * Psi.get_dx();
	H_0 = H_0 / help;
	H_kin = H_0;
	Am.init(xlist.get_size());
	Ap.init(xlist.get_size());

	P_stencil.init(xlist.get_size());
	P_stencil.sub = m_one;
	P_stencil.diag = zero;
	P_stencil.sup = one;
	Creal factor = -Creal(0, 1) / (2 * Psi.get_dx());
	P_stencil = P_stencil * factor;

	// Numerov-Boost
	Delta = H_kin * (-2.0);
	M.init(xlist.get_size());
	M.sub = Creal(1, 0);
	M.diag = Creal(10, 0);
	M.sup = Creal(1, 0);
	help = -6.0;
	M = M / help;

}

void TDSE::init_tracker(string filename) {
	tracker.init(path, filename);

	tracker.add_variable_double("time", tracks.time);
	tracker.add_variable_double("expect_value_x", tracks.expectaionvalue_x);
	tracker.add_variable_double("expect_value_E", tracks.expectaionvalue_E);
	tracker.add_variable_double("expect_value_Ekin", tracks.expectaionvalue_Ekin);
	tracker.add_variable_double("expect_value_p", tracks.expectaionvalue_p);
	tracker.add_variable_double("expval_x_l", tracks.expt_x_left);
	tracker.add_variable_double("expval_x_r", tracks.expt_x_right);
	tracker.add_variable_double("norm_left", tracks.norm_left);
	tracker.add_variable_double("norm_right", tracks.norm_right);

	tracker.write_begin(tlist.get_size());

}

void TDSE::calculate_tracks_variables(Wavefunction psi_prop_in) {
	Wavefunction psi_left(psi_prop_in);
	Wavefunction psi_right(psi_prop_in);

	help_wave = psi_prop_in * xlist;
	tracks.expectaionvalue_x = real(psi_prop_in.scalar(help_wave));


	for (int i = 0; i < xlist.get_size(); i++) {
		if (xlist[i].real() > 0) psi_left[i] = Creal(0, 0);
		else psi_right[i] = Creal(0, 0);
	}

	tracks.norm_left = real(psi_left.scalar(psi_left));
	tracks.expt_x_left = real(psi_left.scalar(help_wave)) / tracks.norm_left;
	tracks.norm_right = real(psi_right.scalar(psi_right));
	tracks.expt_x_right = real(psi_right.scalar(help_wave)) / tracks.norm_right;


	H_0.matmul(help_wave, psi_prop_in);
	tracks.expectaionvalue_E = real(psi_prop_in.scalar(help_wave));


	H_kin.matmul(help_wave, psi_prop_in);
	tracks.expectaionvalue_Ekin = real(psi_prop_in.scalar(help_wave));

	P_stencil.matmul(help_wave, psi_prop_in);
	tracks.expectaionvalue_p = real(psi_prop_in.scalar(help_wave));
}

void TDSE::add_Coulomb(double x0, int Z, double a) {
	V.init_Coulomb(xlist, x0, Z, a);
}

void TDSE::add_box(double depth, double length, double x0) {
	V.init_box(xlist, depth, length, x0);
}

void TDSE::add_fermibox(double depth, double s, double length, double x0) {
	V.init_fermibox(xlist,depth, s, length, x0);
}

void TDSE::set_box_benutzer(double depth, double s, double length, double x0) {
	for (int i = 0; i < xlist.get_size(); i++) {
		if (xlist[i].real() <= x0) V[i] = depth / (exp(-(xlist[i] - (x0 - length / 2)) / s) + 1.0); 
		else V[i] = depth / (exp((xlist[i] - (x0 + length / 2)) / 0.1) + 1.0); ;
	}
}

void TDSE::set_fermibox(double depth, double s, double length, double x0) {
	for (int i = 0; i < V.get_size(); i++) V[i] = 0;
	V.init_fermibox(xlist, depth, s, length, x0);
}

void TDSE::add_Yukawa(double x0, int Z, double a) {
	V.init_Yukawa(xlist, x0, Z, a);
}

void TDSE::set_Hamiltonian() {
	V_tri = V;
	H_0 = H_kin + V_tri;
	if (mpi->rank_group == 0) std::cout << "Hamiltonian was set." << endl;
}

void TDSE::set_AmAp(Creal timestep) {
	// siehe Numerik-Skript Seite 89
	/*Am = M + (Delta + M * V_tri) * 0.5 * Creal(0,1) * dt[j];
	Ap = M - (Delta + M * V_tri) * 0.5 * Creal(0,1) * dt[j];*/

	double dz = xlist[1].real() - xlist[0].real();

	for (int i = 0; i < xlist.get_size() - 1; i++) {
		Am.sub[i] = -1.0 / 6.0 + 0.5 * Creal(0, 1) * timestep * (1.0 / pow(dz, 2) - V_tri.diag[i] / 6.0);
		Am.diag[i] = -10.0 / 6.0 + 0.5 * Creal(0, 1) * timestep * (-2.0 / pow(dz, 2) - 10.0 / 6.0 * V_tri.diag[i]);
		Am.sup[i] = -1.0 / 6.0 + 0.5 * Creal(0, 1) * timestep * (1.0 / pow(dz, 2) - V_tri.diag[i + 1] / 6.0);

		Ap.sub[i] = -1.0 / 6.0 - 0.5 * Creal(0, 1) * timestep * (1.0 / pow(dz, 2) - V_tri.diag[i] / 6.0);
		Ap.diag[i] = -10.0 / 6.0 - 0.5 * Creal(0, 1) * timestep * (-2.0 / pow(dz, 2) - 10.0 / 6.0 * V_tri.diag[i]);
		Ap.sup[i] = -1.0 / 6.0 - 0.5 * Creal(0, 1) * timestep * (1.0 / pow(dz, 2) - V_tri.diag[i + 1] / 6.0);
	}


	Am.diag[xlist.get_size() - 1] = -10.0 / 6.0 + 0.5 * Creal(0, 1) * timestep * (-2.0 / pow(dz, 2) - 10.0 / 6.0 * V_tri.diag[xlist.get_size() - 1]);
	Ap.diag[xlist.get_size() - 1] = -10.0 / 6.0 - 0.5 * Creal(0, 1) * timestep * (-2.0 / pow(dz, 2) - 10.0 / 6.0 * V_tri.diag[xlist.get_size() - 1]);

}

double TDSE::calc_Energy(Wavefunction psi) {
	// nach Numerov-Boost: H = M2^-1*D + V
	// also E = <psi|H|psi> = <psi|M2^-1 D|psi> + <psi|V|psi>

	help_wave = psi;
	Delta.matmul(help_wave, psi); // help_wave = D*Psi

	wave_interm = psi;

	M.solve_ref(wave_interm, help_wave, help_tridi); // M2*wave_interm = help_wave
	V_tri.matmul(help_wave, psi); // helpwave = V*psi;

	wave_interm += help_wave; // H*psi

	double E = real(psi.scalar(wave_interm));  // E = <psi|wave_interm>

	return E;
}

void TDSE::imaginary_time_prop(const int max_state, double time_step) {
	// set #states maximal calculated
	state_max = max_state;
	double timestep = time_step;

	// Init needed variables
	Creal zero(0, 0);
	Creal one(1, 0);
	Ctridiag One(xlist.get_size(), zero, one, zero);
	//Ctridiag Am(One + H_0 * (time_step * 0.5));
	//Ctridiag Ap(One - H_0 * (time_step * 0.5));
	set_AmAp(-Creal(0, 1) * timestep);
	Wavefunction b(xlist.get_size(), Psi.get_dx());
	Wavefunction psi_i(xlist.get_size(), Psi.get_dx());
	Wavefunction psi_h(xlist.get_size(), Psi.get_dx());
	Wavefunction psi_man(psi_h);
	Creal N = (0, 0);
	double Energy;
	double delta;

	// Loop of states 
	for (int n = 0; n < max_state; n++) {
		if (mpi->rank_group == 0) std::cout << " Calc state Nr. " << n + 1 << endl;
		psi_h.init_random();
		timestep = time_step;

		// Loop of steps
		for (int i = 1; i < 50; i++) {

			// Solve imag. propagation
			Ap.matmul(b, psi_h);	//b = Ap * psi_h; 
			Am.solve_ref(psi_h, b, help_tridi);	// psi_h = b / Am; 

			// Outprojection of calculated states
			psi_man = psi_h;
			for (int state = 0; state < n; state++) {
				for (int i = 0; i < xlist.get_size(); i++) psi_i[i] = Psi_states[state][i];
				N = psi_i.scalar(psi_h);
				psi_man = psi_man - psi_i * N;
			}
			// normalization
			psi_man.normalize();
			psi_h = psi_man;

			// Abort criterion: Calculate difference H*Psi - E*Psi (should be 0, TISE) 
			// Numerov-Boost -> H = M_2^-1 * D + V
			Energy = calc_Energy(psi_h);
			help_wave = wave_interm - psi_h * Energy;

			// old (without Numerov-Boost):
			//help_wave = H_0 * psi_h;
			//Energy = real(psi_h.scalar(help_wave));
			//help_wave -= psi_h * E;

			if (real(help_wave.norm()) < 1e-15) {
				Energylevels.push_back(Energy);
				break;
			}

			// Adjusted time step for improved calculations
			delta = Energy * timestep * 0.5;
			if (delta > 0) {}
			else if (delta < -1) { timestep = 0.9 * timestep; }
			else { timestep = 1.1 * timestep; }

			// Calculate new propagation matrices
			/*Am = One + H_0 * timestep * 0.5;
			Ap = One - H_0 * timestep * 0.5;*/
			set_AmAp(-Creal(0, 1) * timestep);

			// No state found within maximal steps
			if (i == 49) {
				Energy = 0;
			}

		}

		// Abort criterion for maximal state (if no more state is found)
		if (Energy == 0) {
			state_max = n;
			Energylevels[n] = 0;
			break;
		}
		// Save groundstates
		Psi_states.push_back(psi_h);
	}

	if (mpi->rank_group == 0) std::cout << "Energylevels:" << endl;
	for (int i = 0; i < state_max; i++) {
		if (mpi->rank_group == 0) std::cout << " E[" << i+1 << "] =" << Energylevels[i] << endl;
	}
}

void TDSE::imaginary_time_prop_Vmin(const int max_state, double Vmin) {
	// set #states maximal calculated
	state_max = max_state;

	// set timestep that delta = -1 for minimal V (depth)
	double timestep = real(2 / (1.1 * Vmin));

	// Init needed variables
	Creal zero(0, 0);
	Creal one(1, 0);
	Ctridiag One(xlist.get_size(), zero, one, zero);
	//Ctridiag Am(One + H_0 * (timestep * 0.5));
	//Ctridiag Ap(One - H_0 * (timestep * 0.5));
	set_AmAp(-Creal(0,1)*timestep);

	Wavefunction b(xlist.get_size(), Psi.get_dx());
	Wavefunction psi_i(xlist.get_size(), Psi.get_dx());
	Wavefunction psi_h(xlist.get_size(), Psi.get_dx());
	Wavefunction psi_man(psi_h);
	Creal N = (0, 0);
	double Energy = -Vmin;
	double delta;

	// Loop of states 
	for (int n = 0; n < max_state; n++) {
		if (mpi->rank_group == 0) std::cout << " Calc state Nr. " << n + 1 << endl;
		psi_h.init_random();
		timestep = real(-2 / (1.1 * Energy));

		// Loop of steps
		for (int i = 1; i < 100; i++) {

			// Solve imag. propagation
			Ap.matmul(b, psi_h);	//b = Ap * psi_h; 
			Am.solve_ref(psi_h, b, help_tridi);	// psi_h = b / Am; 

			// Outprojection of calculated states
			psi_man = psi_h;
			for (int state = 0; state < n; state++) {
				for (int i = 0; i < xlist.get_size(); i++) psi_i[i] = Psi_states[state][i];
				N = psi_i.scalar(psi_h);
				psi_man = psi_man - psi_i * N;
			}

			// normalization
			psi_man.normalize();
			psi_h = psi_man;

			// Abort criterion: Calculate difference H*Psi - E*Psi (should be 0, TISE) 
			// Numerov-Boost -> H = M_2^-1 * D + V
			Energy = calc_Energy(psi_h);
			help_wave = wave_interm - psi_h * Energy;

			// old (without Numerov-Boost):
			//help_wave = H_0 * psi_h;
			//Energy = real(psi_h.scalar(help_wave));
			//help_wave -= psi_h * E;
			
			if (real(help_wave.norm()) < 1e-15) {
				Energylevels.push_back(Energy);
				break;
			}

			// Adjusted time step for improved calculations
			delta = Energy * timestep * 0.5;
			if (delta > 0) {}
			else if (delta < -1) { timestep = 0.99 * timestep; }
			else { timestep = 1.01 * timestep; }

			// Calculate new propagation matrices
			/*Am = One + H_0 * timestep * 0.5;
			Ap = One - H_0 * timestep * 0.5;*/
			set_AmAp(-Creal(0, 1) * timestep);

			// No state found within maximal steps
			if (i == 99) {
				Energy = 0;
			}
		}

		// Abort criterion for maximal state (if no more state is found)
		if (Energy == 0) {
			state_max = n;
			Energylevels[n] = 0;
			break;
		}
		// Save groundstates
		Psi_states.push_back(psi_h);
	}

	if (mpi->rank_group == 0) std::cout << "Energylevels:" << endl;
	for (int i = 0; i < state_max; i++) {
		if (mpi->rank_group == 0) std::cout << " E[" << i+1 << "] =" << Energylevels[i] << endl;
	}
}

void TDSE::imaginary_time_prop_parallel(int all_states) {
	// set #states maximal calculated
	std::vector<Wavefunction> Psi_interm;
	Psi.init(xlist.get_size());
	for (int i = 0; i < all_states; i++) {
		Psi.init_random();
		Psi.set_energy(-V.minimum());
		Psi_interm.push_back(Psi);
	}

	// set timestep that delta = -1 for minimal V (depth)
	double timestep = -2 / V.minimum();

	vector<double> dt;
	dt.assign(all_states, timestep);

	vector<double> error;
	error.assign(all_states, 100);

	// Init needed variables
	Creal zero(0, 0);
	Creal one(1, 0);
	Ctridiag One(xlist.get_size(), zero, one, zero);
	//Ctridiag Am(One + H_0 * (timestep * 0.5));
	//Ctridiag Ap(One - H_0 * (timestep * 0.5));
	set_AmAp(-Creal(0, 1) * timestep);
	Wavefunction b(xlist.get_size(), Psi.get_dx());
	Wavefunction psi_i(xlist.get_size(), Psi.get_dx());
	Wavefunction psi_h(xlist.get_size(), Psi.get_dx());
	Wavefunction psi_man(psi_h);
	Creal N = (0, 0);

	double delta;
	double max;


	for (int ti = 0; ti < 800; ti++) {
		if(ti%50 == 0 && mpi->rank_group == 0) cout << "Imaginary time iteration " << ti << endl;

		// Loop of states 
		for (int n = 0; n < all_states; n++) {
				Psi = Psi_interm[n];

				set_AmAp(-Creal(0, 1) * dt[n]);
				/*double help = dt[n] * 0.5;
				Am.sub = H_0.sub * help;
				Am.sup = H_0.sup * help;
				Am.diag = 1 + H_0.diag * help;

				help *= -1;
				Ap.sub = H_0.sub * help;
				Ap.sup = H_0.sup * help;
				Ap.diag = 1 + H_0.diag * help;*/

				Ap.matmul(b, Psi);	//b = Ap * psi_h; 
				Am.solve_ref(Psi, b, help_tridi);	// psi_h = b / Am;

				Psi.normalize();

				Psi_interm[n] = Psi;
		}

		for (int n = 0; n < all_states; n++) {
			Psi = Psi_interm[n];

			if (ti % 2 == 0) {
				for (int k = 1; k <= n; k++) {
					help_wave = Psi_interm[(int)n - k];
					N = Psi.scalar(help_wave);
					Psi -= help_wave * N;
				}
				
				Psi.normalize();
			}

			// Numerov-Boost -> H = M_2^-1 * D + V
			Psi.set_energy(calc_Energy(Psi));
			help_wave = wave_interm - Psi * Psi.get_energy();

			// old (without Numerov-Boost):
			/*help_wave = H_0 * Psi;
			Psi.set_energy(real(Psi.scalar(help_wave)));
			help_wave -= Psi * Psi.get_energy();*/

			error[n] = real(help_wave.norm());


			// Adjusted time step for improved calculations
			delta = Psi.get_energy() * dt[n] * 0.5;
			if (delta > 0) {}
			else if (delta < -0.95)  dt[n] = dt[n] * (1 - 0.05 * exp(-ti / 50));
			else  dt[n] = dt[n] * (1 + 0.05 * exp(-ti / 50));

			Psi_interm[n] = Psi;
		}

		max = *max_element(error.begin(), error.end());

		if (max < 1e-7) {
			break;
		}
	}

	std::sort(Psi_interm.begin(), Psi_interm.end());
	
	int number_of_bound_states = 0;
	for (auto& psi : Psi_interm) {
		if (psi.get_energy() <= 0) {
			Psi_states_GS.push_back(psi);
			Energylevels.push_back(psi.get_energy());
			number_of_bound_states++;
		}
	}

	state_max = number_of_bound_states;

	Psi_states = Psi_states_GS;

	if (mpi->rank_group == 0) std::cout << "Energylevels:" << endl;
		for (int i = 0; i < number_of_bound_states; i++) {
			if (mpi->rank_group == 0) std::cout << " E[" << i << "] =" << Energylevels[i] << endl;
		}
}

double TDSE::find_state(double Energy) {
	Energy = -Energy / au_Ee;
	double test = 100;
	double state = 0;
	for (int i = 0; i < Energylevels.size(); i++) {
		if (abs(Energy - Energylevels[i]) < test) {
			state = i;
			test = abs(Energy - Energylevels[i]);
		}
}
	return state;
}

//void TDSE::find_box_pot(double W, double Ef) {
//	if (mpi->rank_group == 0) std::cout << "Search box ...." << endl;
//	double time_step = 10;
//	W = W / au_Ee;
//	Ef = Ef / au_Ee;
//
//	double V0 = -(W + Ef);
//	double x0 = 0;
//	double L_low = 0.1 / (au_L * 1e10);
//	double L_high = 5 / (au_L * 1e10);
//	double L = 0.5 * (L_high + L_low);
//	double Energy;
//	help_wave.init(xlist.get_size());
//	help_wave.init_random();
//	Potential Vhelp(xlist.get_size(), Psi.get_dx());
//	Wavefunction b(xlist.get_size(), Psi.get_dx());
//	Wavefunction psi_i(xlist.get_size(), Psi.get_dx());
//	Wavefunction psi_h(xlist.get_size(), Psi.get_dx());
//	Creal zero(0, 0);
//	Creal one(1, 0);
//	Ctridiag One(xlist.get_size(), zero, one, zero);
//
//	for (int i = 1; i < 100; i++) {
//		L = 0.5 * (L_high + L_low);
//
//		Vhelp.init_box(xlist, V0, L, x0);
//		V_tri = Vhelp;
//		H_0 = H_kin + V_tri;
//
//		// Init needed variables
//		Ctridiag Am(One + H_0 * (time_step * 0.5));
//		Ctridiag Ap(One - H_0 * (time_step * 0.5));
//		double delta;
//		psi_h.init_random();
//
//		// Loop of steps
//		for (int i = 1; i < 50; i++) {
//
//			// Solve imag. propagation
//			Ap.matmul(b, psi_h);	//b = Ap * psi_h; 
//			Am.solve_ref(psi_h, b, help_tridi);	// psi_h = b / Am; 
//
//
//			psi_h.normalize();
//
//			// Abort criterion: Calculate difference H*Psi - E*Psi (should be 0, TISE) 
//			help_wave = H_0 * psi_h;
//			Energy = real(psi_h.scalar(help_wave));
//			help_wave -= psi_h * Energy;
//
//			if (real(help_wave.norm()) < 1e-15) {
//				break;
//			}
//
//			// Adjusted time step for improved calculations
//			delta = Energy * time_step * 0.5;
//			if (delta > 0) {}
//			else if (delta < -1) { time_step = 0.9 * time_step; }
//			else { time_step = 1.1 * time_step; }
//
//			// Calculate new propagation matrices
//			Am = One + H_0 * time_step * 0.5;
//			Ap = One - H_0 * time_step * 0.5;
//
//			// No state found within maximal steps
//			if (i == 49) {
//				Energy = 0;
//			}
//		}
//
//		double error = W + Energy;
//
//		if (abs(error) < 1e-3) break;
//		else if (error < 0) L_high = L;
//		else L_low = L;
//
//		for (int i = 0; i < xlist.get_size(); i++) Vhelp[i] = 0;
//	}
//
//	if (mpi->rank_group == 0) std::cout << "Box was found: L = " << L * au_L * 1e10 << " Angstrom and E = " << Energy * au_Ee << " eV" << endl;
//	add_box(V0, L, x0);
//	Psi_states.push_back(psi_h);
//	Energylevels.push_back(Energy);
//	set_Hamiltonian();
//}

void TDSE::find_box_pot(double W, double Ef) {
	if (mpi->rank_group == 0) std::cout << "Search box ...." << endl;
	double V0 = -(W + Ef);

	// set timestep that delta = -1 for minimal V (depth)
	double timestep = real(2 / (1.1 * W));

	double L_low = 0.01 / (au_L * 1e10);
	double L_high = 3.0 / (au_L * 1e10);
	double L = 0.5 * (L_high + L_low);
	double x0 = -L / 2;

	Creal N = (0, 0);
	double Energy;
	double delta;

	help_wave.init(xlist.get_size());
	help_wave.init_random();
	Potential Vhelp(xlist.get_size(), Psi.get_dx());
	Wavefunction b(xlist.get_size(), Psi.get_dx());
	Wavefunction psi_i(xlist.get_size(), Psi.get_dx());
	Wavefunction psi_h(xlist.get_size(), Psi.get_dx());

	Creal zero(0, 0);
	Creal one(1, 0);
	Ctridiag One(xlist.get_size(), zero, one, zero);

	for (int i = 1; i < 1000; i++) {
		L = 0.5 * (L_high + L_low);
		x0 = -L / 2;

		Vhelp.init_box(xlist, V0, L, x0);
		V_tri = Vhelp;
		H_0 = H_kin + V_tri;

		// Init needed variables
		set_AmAp(-Creal(0, 1) * timestep);
		psi_h.init_random();

		// Loop of steps
		for (int i = 1; i < 100; i++) {

			// Solve imag. propagation
			Ap.matmul(b, psi_h);	//b = Ap * psi_h; 
			Am.solve_ref(psi_h, b, help_tridi);	// psi_h = b / Am; 

			psi_h.normalize();

			// Abort criterion: Calculate difference H*Psi - E*Psi (should be 0, TISE) 
			// Numerov-Boost -> H = M_2^-1 * D + V
			Energy = calc_Energy(psi_h);
			help_wave = wave_interm - psi_h * Energy;

			if (real(help_wave.norm()) < 1e-15) {
				break;
			}

			// Adjusted time step for improved calculations
			delta = Energy * timestep * 0.5;
			if (delta > 0) {}
			else if (delta < -1) { timestep = 0.99 * timestep; }
			else { timestep = 1.01 * timestep; }

			// Calculate new propagation matrices
			set_AmAp(-Creal(0, 1) * timestep);

			// No state found within maximal steps
			if (i == 99) {
				Energy = 0;
			}
		}

		double error = W + Energy;

		if (abs(error) < 1e-3) break;
		else if (error < 0) L_high = L;
		else L_low = L;

		for (int i = 0; i < xlist.get_size(); i++) Vhelp[i] = 0;
		if (mpi->rank_group == 0) std::cout << "Step " << i << " L = " << 0.5 * (L_high + L_low) << " W = " << W << " E = " << Energy << endl;
	}

	if (mpi->rank_group == 0) std::cout << "Box was found: L = " << L * au_L * 1e10 << " Angstrom and E = " << Energy * au_Ee << " eV" << endl;
	add_box(V0, L, x0);
	//Psi_states.push_back(psi_h);
	//Energylevels.push_back(Energy);
	//set_Hamiltonian();
}

void TDSE::read_param(string path) {
	std::ifstream ifs;
	ifs.open(path, std::ios::in | std::ios::binary);

	readHeader(ifs, "1.0");

	state_max = readInt(ifs);

	ifs.close();
}

void TDSE::import_GS(string in_path, string save) {
	// Parameter
	string file = in_path;
	file.append("param.dat");
	read_param(file);

	// z axis
	file = in_path;
	file.append("zlist.dat");
	xlist.read(file);


	init_tdse();
	path = save;

	// final potential v_int and v_x
	file = in_path;
	file.append("vint.dat");
	Potential v_int(Psi.get_dx());
	v_int.init(xlist.get_size());
	v_int.read(file);
	
	file = in_path;
	file.append("vx.dat");
	Potential v_x(Psi.get_dx());
	v_x.init(xlist.get_size());
	v_x.read(file);

	V = v_int + v_x;
	set_Hamiltonian();

	// final Orbitals
	for (int j = 0; j < state_max; j++) {
		file = in_path;
		file.append("Orbital");
		file.append(to_string(j));
		file.append(".dat");
		Psi.read(file);
		Psi_states.push_back(Psi);


		help_wave = H_0 * Psi;
		double Energy = real(Psi.scalar(help_wave));
		Energylevels.push_back(Energy);
	}


	if (mpi->rank_group == 0) std::cout << "GS import complete." << endl;
}

void TDSE::set_psi_start(vector<vector<double>> state_number) {
	Psi_start = 0.0;

	if (state_number.empty()) {

		for (int j = 0; j < state_max; j++) {
			for (int i = 0; i < xlist.get_size(); i++) {
				Psi_start[i] += sqrt(1.0/state_max) * Psi_states[j][i];
			}
		}
		return;
	}

	for (auto& states : state_number) {
		if (states[1] > state_max && mpi->rank_group == 0) {
			cout << "start wavefunction outside groundstate wavefunctions" << endl;

		}
		for (int i = 0; i < xlist.get_size(); i++) {
			Psi_start[i] += sqrt(states[0])*Psi_states[states[1]][i];
		}
	}
}

void TDSE::add_pulse(double lambda_SI, double I0_SI, double t0_SI, double phi_SI, double FWHM_SI) {
	Pulse new_pulse(lambda_SI, I0_SI, t0_SI, phi_SI, FWHM_SI);
	Up = new_pulse.get_Up();
	Pulse_prop.push_back(new_pulse);
}

void TDSE::add_pulse_vp(double lambda_SI, double I0_SI, double t0_SI, double phi_SI, double FWHM_in, int use_SC) {
	double FWHM_SI;
	if (use_SC == 1) FWHM_SI = lambda_SI * 1e6 / (2 * const_c);
	else FWHM_SI = FWHM_in;
	Pulse new_pulse(lambda_SI, I0_SI, t0_SI, phi_SI + M_PI/2, FWHM_SI);
	Up = new_pulse.get_Up();
	Pulse_prop_vp.push_back(new_pulse);
}

void TDSE::add_pulsetrain(double lambda_SI, double I0_SI, std::vector<double> t0_SI, double phi_SI, double FWHM_SI, double t0_env, double FWHM_env) {
	Pulsetrain_prop.set(lambda_SI, I0_SI, t0_SI, phi_SI, FWHM_SI, t0_env, FWHM_env);
}

void TDSE::input_field(string filename) {

	// ------ Anmerkungen: --------
	// Klappt erstmal nur, wenn die Input-Zeitliste und das Feld in [a.u.] gegeben ist. (t_list ist auch in a.u.)
	// Optimal wäre natürlich, wenn die Funktion Eingabeparameter bekommt, welche Einheiten der Input hat und das entsprechend umgerechnet wird... :D 
	// Die Input-Datei besitzt aktuell keinen Header, sondern nur die Größe der Listen sowie die Einträge (erst Zeit, dann Feld)

	string path = "./Input/" + filename + ".dat";
	if(mpi->rank_group == 0) std::cout << "Input field from " << path << std::endl;

	Efield_input.init(tlist.get_size());

	ifstream ifs;
	ifs.open(path, ios::in | ios::binary);
	
	int size = readInt(ifs);
	Cvec t_in(size);
	Cvec E_in(size);
	for (int i = 0; i < size; i++) {
		t_in[i] = readDouble(ifs); // a.u.
	}
	for (int i = 0; i < size; i++) {
		E_in[i] = readDouble(ifs); // a.u.
	}

	double Up_in = readDouble(ifs);
	ifs.close();

	Up = Up_in;

	for (int ti = 0; ti < tlist.get_size(); ti++) {
		double time = tlist[ti].real();
		double diff = 1000;
		double ind = 0.0;

		for (int j = 0; j < t_in.get_size(); j++) {
			if (abs(time - t_in[j]) < diff) {
				diff = abs(time - t_in[j]);
				ind = j;
			}
		}

		if (diff < 1e-12 || ind == 0 || ind == t_in.get_size() - 1.0)	Efield_input[ti] = E_in[ind];
		else if (t_in[ind].real() > time) Efield_input[ti] = E_in[ind - 1] + (E_in[ind] - E_in[ind - 1]) / (t_in[ind] - t_in[ind - 1]) * (time - t_in[ind - 1]);
		else if (t_in[ind].real() < time) Efield_input[ti] = E_in[ind] + (E_in[ind + 1] - E_in[ind]) / (t_in[ind + 1] - t_in[ind]) * (time - t_in[ind]);
	}
}

void TDSE::set_gamma(double gamma_0, double lambda_nf) {
	gamma.init(xlist.get_size());
	gamma = 0;
	lambda_nf /= 1e10 * au_L;
	for (int i = 0; i < xlist.get_size(); i++) {
		if (xlist[i].real() >= x_right) gamma[i] = 1.0 + (gamma_0 - 1.0) * exp(-xlist[i] / lambda_nf);
	}
	if (Pulse_prop_vp.size() > 0)	Up = Pulse_prop_vp[0].get_Up() * gamma_0 * gamma_0;
	if (Pulse_prop.size() > 0)	Up = Pulse_prop[0].get_Up() * gamma_0 * gamma_0;
	if (Efield_input.size > 0) Up = Up * gamma_0 * gamma_0;
}

Cvec TDSE::calculate_int_gamma() {
	Cvec res(xlist.get_size());
	double dx = xlist[1].real() - xlist[0].real();
	if (gamma.get_size() != xlist.get_size()) {
		res = xlist;
		return res;
	}
	res[0] = gamma[0]*dx;
	for (int i = 1; i < xlist.get_size(); i++) {
		res[i] = res[i - 1] + gamma[i]*dx;
	}
	return res;
}

void TDSE::propagation(Wavefunction psi_start, int write_timesteps, double write_x, int write_t, string input_path){
	if (mpi->rank_group == 0) std::cout << "Propagation" << endl;
	// Set time list for propagation
	double time_step = real(tlist[1] - tlist[0]);

	// Open Variables-Tracker if wanted
	if (track_var == 1) init_tracker("Tracker.dat");

	// Initialize needed Variables
	Wavefunction psi_prop(psi_start);
	Creal zero(0, 0);
	Creal one(1, 0);
	Ctridiag One(xlist.get_size(), zero, one, zero);
	Ctridiag Am_old(One + H_0 * Creal(0, 1) * (time_step * 0.5));
	Ctridiag Ap_old(One - H_0 * Creal(0, 1) * (time_step * 0.5));
	Wavefunction b(xlist.get_size(), Psi.get_dx());

	E_las.init(tlist.get_size());
	E_las = 0;

	// Load Input Field and interpolate (set 0 if not initialized)
	if (Efield_input.get_size() < tlist.get_size()) {
		Efield_input.init(tlist.get_size());
		Efield_input = 0;
	}

	// Set spectral field (set 0 if not initialized)
	if (Spectral.E_omega.get_size() == 0) {
		Spectral.E_time.init(tlist.get_size());
		Spectral.E_time = 0;
	}
	else Spectral.calculate_Etime(tlist);

	// Set field of SC_pulse zero if not existing
	E_field_vp.init(tlist.get_size());
	E_field_vp = 0;
	for (auto& p_ind : Pulse_prop_vp) {
			E_field_vp += p_ind.getField_vp(tlist);
	}

	// Init potentials and set Imaginary potential (BC)
	V_las.init(xlist.get_size());
	V_imag_vec.init(xlist.get_size());
	double x0 = 10;
	Creal V_imag_value = Creal(0, 1e4);
	V_imag_vec.init_imag(xlist, x0, V_imag_value, 2.0);

	// write t- and x-resolved WF and potential only if write_timestep = 1
	double write_size = write_x;
	const int size_save = (int)xlist.get_size() / write_size;
	double dx_save = Psi.get_dx() * write_size;
	Wavefunction wave_save(size_save, dx_save);
	Cvec pot_save(size_save);
	if (write_timesteps == 1) {
		Psitime.init(path, "Psi_evolution.dat");
		Psitime.write_begin(tlist.get_size(), wave_save.get_size());

		Potential_save.init(path, "Potential_time.dat");
		Potential_save.write_begin(tlist.get_size(), wave_save.get_size());
	}

	// Save Spectrum over time
	Cvec Elist_spec(0.0, 0.05, 25.0);
	Spec_evolution.init(path, "Spec_evol.dat");
	Spec_evolution.write_begin(tlist.get_size(), Elist_spec.get_size());

	// write population over time only if states were calculated via imag. time prop.
	if (state_max >= 0) {
		Popul.init(path, "Population_over_time.dat");
		Popul.write_begin(tlist.get_size(), state_max);
	}

	// Calculate integral of enhancement factor gamma (Int_gamma = x_list if gamma not initilized)
	Cvec Int_gamma(xlist.get_size());
	Int_gamma = calculate_int_gamma();

	// ------------------ Time loop --- PROPAGATION ------------------
	for (int ti = 0; ti < tlist.get_size(); ti++) {
		double t = tlist[ti].real();
		if (track_var == 1) tracks.time = t;
		if (ti % 500 == 0 && mpi->rank_group == 0) std::cout << "Calc Prop: t = " << t * au_T * 1e15 << endl;

		// Calculate laser fields (all together)
		for (auto& p_ind : Pulse_prop) {
			E_las[ti] = E_las[ti] + p_ind.getField(t);
		}
		E_las[ti] += Pulsetrain_prop.getField(t);
		E_las[ti] += Efield_input[ti];
		E_las[ti] += Spectral.E_time[ti];
		E_las[ti] += E_field_vp[ti];

		// Calculate laser field potential
		for (int xi = 0; xi < xlist.get_size(); xi++) {
			V_las[xi] = au_e * E_las[ti].real() * Int_gamma[xi];
		}

		// Set Propagation matrices
		for (int i = 0; i < xlist.get_size(); i++) {
			Am_old.diag[i] = 1.0 + (H_0.diag[i] + V_las[i] + V_imag_vec[i]) * Creal(0, 1) * time_step * 0.5;
			Ap_old.diag[i] = 1.0 - (H_0.diag[i] + V_las[i] + V_imag_vec[i]) * Creal(0, 1) * time_step * 0.5;
 		}

		//help_cvec = V + V_las + V_imag_vec;
		//V_tri = help_cvec;
		//set_AmAp(time_step);

		// Calculate Propagation step
		Ap_old.matmul(b, psi_prop);	//b = Ap * psi_prop; //
		Am_old.solve_ref(psi_prop, b, help_tridi);	// psi_prop = b / Am;

		// Save WF/potential/spectra/population (if active)
		if (write_timesteps == 1 && ti % write_t == 0) {
			for (int i = 0; i < wave_save.get_size(); i++) {
				wave_save[i] = psi_prop[write_size * i];
				pot_save[i] = V[write_size * i] + V_las[write_size * i];
			}
			Psitime.write(wave_save);
			Potential_save.write(pot_save);
		}
		if (write_spec_evol == 1 && ti % 500 == 0) {
			Psi = psi_prop;
			outprojection(Psi_start);
			calc_spec(Psi, Elist_spec, 3, 0.0025, 1);
			Spec_evolution.write(spec);
		}
		int index = 0;
		for (auto& psi:Psi_states){
			Popul.write(Creal(abs(psi.scalar(psi_prop)),0));
			index++;
		}
		// Save track variables
		if (track_var == 1) {
			calculate_tracks_variables(psi_prop);
			tracker.write();
		}

	}

	// set actual/last WF to psi
	Psi = psi_prop;

	// Close all files
	Psitime.close();
	Potential_save.close();
	Popul.close();
	tracker.close();
	Spec_evolution.close();
}

//void TDSE::propagation_test(double time_start, double time_step, double time_end) {
//	init_tlist(time_start, time_step, time_end);
//	time_step = time_step / (au_T * 1e15);
//
//	// Tracker
//	double t = tlist[0].real();
//	double expectaionvalue_x = 1;
//	tracker.add_variable_double("time", t);
//	tracker.add_variable_double("expect_value_x", t);
//
//	tracker.write_begin(tlist.get_size());
//
//	Psitime.init(path, "Psi_evolution.dat");
//	Psitime.write_begin(tlist.get_size(), xlist.get_size());
//
//	if (state_max != 0) {
//		Popul.init(path, "Population_over_time.dat");
//		Popul.write_begin(tlist.get_size(), state_max);
//	}
//
//	for (int ti = 0; ti < tlist.get_size(); ti++) {
//		t = tlist[ti].real();
//		if (mpi->rank_group == 0) std::cout << t << endl;
//		
//		tracker.write();
//		Psitime.write(Psi_start);
//	}
//	Psitime.close();
//	Popul.close();
//	tracker.close();
//}

void TDSE::propagation_au(Wavefunction psi_start, int write_timesteps, double write_x, int write_t, string input_path) {
	propagation(psi_start, write_timesteps, write_x, write_t, input_path);
}

void TDSE::outprojection(Wavefunction psi_out) {
	Creal N = psi_out.scalar(Psi);
	Psi = Psi - psi_out * N;
}

void TDSE::calc_spec(Wavefunction psi_start, double Energy_start, double Energy_step, double Energy_end, int order, double gamma, int half_spec) {
	// Init and set energy list
	Cvec Elist_help(Energy_start, Energy_step, Energy_end);
	Elist.init(Elist_help.get_size());
	Elist = Elist_help;
	spec.init(Elist.get_size());

	// Init needed variables
	Creal zero(0, 0);
	Creal one(1, 0);
	Ctridiag One(xlist.get_size(), zero, one, zero);
	Wavefunction psi_scale(xlist.get_size(), Psi.get_dx());
	Wavefunction psi_k(psi_scale);
	Wavefunction psi_interm(psi_scale);
	Wavefunction wavehelp(psi_k);
	Ctridiag P1(H_0);
	Ctridiag P2(H_0);
	Ctridiag help(One);
	help_tridi = H_0;

	// Calculate normalization constant
	double C = pow(2.0, order) * sin(M_PI / pow(2.0, order)) / (2.0 * M_PI * gamma);
	// Set scaled WF
	for (int i = 0; i < xlist.get_size(); i++) psi_scale[i] = C * pow(gamma, pow(2.0, order)) * psi_start[i];

	// Energy loop --- WINDOW FUNCTION ---
	for (int Eint = 0; Eint < Elist.get_size(); Eint++) {
		double E_test = Elist[Eint].real() / au_Ee;
		if(Eint%20 == 0 && mpi->rank_group == 0) std::cout << "Calc Spec: E = " << E_test*au_Ee << " eV" << endl;

		psi_k = psi_scale;

		for (int i = 0; i < xlist.get_size(); i++) {
			help_tridi.diag[i] -= E_test;
		}

		// Order loop
		for (double k = pow(2.0, order - 1.0); k >= 1.0; k--) {
			double nu_kn = (2.0 * k - 1.0) * M_PI / pow(2.0, order);

			// Set Window matrices
			for (int i = 0; i < xlist.get_size(); i++) {
				P1.diag[i] = help_tridi.diag[i] + gamma * exp(Creal(0, nu_kn));
				P2.diag[i] = help_tridi.diag[i] - gamma * exp(Creal(0, nu_kn));
			}

			// Solve equation system
			P1.solve_ref(psi_interm, psi_k, help);
			P2.solve_ref(psi_k, psi_interm, help);
		}

		// If only x > 0 should be used for calculating the spectrum (half_spec = 1), set wavehelp to 0 for x < 0
		wavehelp = psi_k;

		//if (Eint == Elist.get_size() / 2) wavehelp.save("test_vor.dat");
		if (half_spec == 1) {
			for (int i = 0; i < xlist.get_size(); i++) {
				/*if (xlist[i].real() < x_right) wavehelp[i] = 0.0;*/
				wavehelp[i] = wavehelp[i] * 1.0 / (exp(-(xlist[i].real() - x_right) / 5.0) + 1.0);
			}
		}
		//if (Eint == Elist.get_size() / 2) wavehelp.save("test_nach.dat");

		// Calculate spectrum entry
		spec[Eint] = psi_start.scalar(wavehelp) / au_Ee;
		// reset help function
		help_tridi = H_0;
	}
}

void TDSE::calc_spec(Wavefunction psi_start, Cvec Energy, int order, double gamma, int half_spec) {
	calc_spec(psi_start, Energy[0].real(), real(Energy[2] - Energy[1]), Energy[Energy.get_size() - 1].real(), order, gamma, half_spec);
}

void save_array(Creal** arr, int rows, int columns, std::string filename) {
	std::ofstream ofs;
	ofs.open(filename, std::ios::out | std::ios::binary);
	saveHeader(ofs, "1.0");
	ofs.write((char*)&rows, sizeof(rows));
	ofs.write((char*)&columns, sizeof(columns));
	for (int n = 0; n < rows; n++) {
		for (int m = 0; m < columns; m++) {
			double entry_real = real(arr[n][m]);
			ofs.write((char*)&entry_real, sizeof(entry_real));
		}
	}
	for (int n = 0; n < rows; n++) {
		for (int m = 0; m < columns; m++) {
			double entry_imag = imag(arr[n][m]);
			ofs.write((char*)&entry_imag, sizeof(entry_imag));
		}
	}
	ofs.close();
}

void save_vector(vector<double> vec, std::string filename) {
	std::ofstream ofs;
	ofs.open(filename, std::ios::out | std::ios::binary);
	saveHeader(ofs, "1.0");
	int size = vec.size();
	ofs.write((char*)&size, sizeof(size));
	for (auto& entry: vec) {
			ofs.write((char*)&entry, sizeof(entry));
	}
	ofs.close();
}

void TDSE::save_param(string input_path) {
	std::ofstream ofs;
	ofs.open(input_path, std::ios::out | std::ios::binary);

	saveHeader(ofs, "1.0");
	saveDouble(ofs, Up);

	ofs.close();
}

void TDSE::save_all() {
	string paramfile = path;
	paramfile.append("paramfile.dat");
	save_param(paramfile);

	string x_dat = path;
	x_dat.append("x_list.dat");
	xlist.save(x_dat);

	string t_dat = path;
	t_dat.append("t_list.dat");
	tlist.save(t_dat);

	string Elist_dat = path;
	Elist_dat.append("energy_list.dat");
	Elist.save(Elist_dat);

	string spec_dat = path;
	spec_dat.append("spectrum.dat");
	spec.save(spec_dat);

	string V_dat = path;
	V_dat.append("potential.dat");
	V.save(V_dat);

	string Psistart_dat = path;
	Psistart_dat.append("Psi_start.dat");
	Psi_start.save(Psistart_dat);

	string Psi_dat = path;
	Psi_dat.append("Psi.dat");
	Psi.save(Psi_dat);

	string vimag = path;
	vimag.append("V_imag.dat");
	V_imag_vec.save(vimag);

	string Vlasdat = path;
	Vlasdat.append("V_las.dat");
	V_las.save(Vlasdat);

	string E_las_dat = path;
	E_las_dat.append("E_las.dat");
	E_las.save(E_las_dat);


	int  index = 0;
	for (auto& psi : Psi_states) {
		string filename = path;
		filename.append("Psi");
		filename += to_string(index);
		filename.append(".dat");
		psi.save(filename);
		index++;
	}

	string E_dat = path;
	E_dat.append("energylevels.dat");
	save_vector(Energylevels, E_dat);

	string gamma_dat = path;
	gamma_dat.append("gamma.dat");
	gamma.save(gamma_dat);

	/*string Psi_dat = path;
	Psi_dat.append("Psi_evolution.dat");
	save_array(Psi_time, tlist.get_size(), xlist.get_size(), Psi_dat);

	string Pop_dat = path;
	Pop_dat.append("Population_over_time.dat");
	save_array(Population_time, tlist.get_size(), state_max, Pop_dat);*/
}

