#include<iostream>
#include"Cvec.h"
#include"Ctridiag.h"
#include<time.h>
#include<stdio.h>
#include<stdlib.h>
#include<fstream>

using namespace std;

int main() {
	time_t start, end;
	time(&start);


	// Ini lists
	Cvec x_list(-300.0, 0.01, 300.0); // SI
	x_list = x_list * 1e-10 / au_L; // au
	x_list.save("x_list.dat");
	int x_size = x_list.get_size();
	double dx = x_list[2].real() - x_list[1].real();
	
	Cvec t_list(0.0, 0.1, 50.0); // SI
	t_list = t_list / (au_T * 1e15); // au
	double dt = t_list[2].real() - t_list[1].real();

	// How many states?
	const int max_state = 1;

	
	
	// Ini wavefunction
	Wavefunction psi_0(x_size, dx);
	psi_0.init_random();

	// Init Potential
	Potential v_plot(x_size, dx);
	v_plot.init_Coulomb(x_list, 0.0, 1, 2.0);
	//v_plot.init_box(x_list, -2.0, 4.0, 0.0);
	v_plot.save("v_plot.dat");


	// Find states
	Creal zero(0, 0);
	Creal one(1, 0);
	Creal m_two(-2, 0);
	Ctridiag One(x_size, zero, one, zero);


	Ctridiag V(x_size, zero, zero, zero);
	for (int i = 1; i < V.diag.get_size(); i++) V.diag[i] = v_plot[i];
	Ctridiag H_0(x_size, one, m_two, one);
	double help = -2 * dx * dx;
	H_0 = H_0 / help + V;
	Ctridiag Am(One + H_0 * (dt / 2));
	Ctridiag Ap(One - H_0 * (dt / 2));
	Ctridiag help_am = Am;


	int end_state = max_state;
	Creal** psi = new Creal * [max_state];
	for (int i = 0; i < max_state; i++) {
		psi[i] = new Creal[x_size];
	}
	double Energy[max_state] = {};

	Wavefunction b(x_size, dx);
	Wavefunction psi_i(x_size, dx);
	Wavefunction psi_h(psi_0);
	Wavefunction psi_man(psi_h);
	Creal N = (0, 0);


	for (int n = 0; n < max_state; n++) {
		cout << n << endl;
		psi_h = psi_0;

		for (int ti = 0; ti < t_list.get_size(); ti++) {
			
			double t = real(t_list[ti]);
			cout << "t = " << t * 1e15 * au_T << endl;

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

			psi_man = psi_h;
			for (int i = 0; i < n; i++) {
				for (int j = 0; j < x_size; j++) psi_i[j] = psi[i][j];
				N = psi_i.scalar(psi_h);
				psi_man = psi_man - psi_i * N;
			}
			psi_man.normalize();
			psi_h = psi_man;

		}
		Wavefunction H_psi = H_0 * psi_h;
		Energy[n] = real(psi_h.scalar(H_psi));

		if (Energy[n] > 0) {
			end_state = n - 1;
			Energy[n] = 0;
			break;
		}
		for (int i = 0; i < x_size; i++) {
			psi[n][i] = psi_h[i];
		}
		/*string name = to_string(n);
		psi_h.save(name);*/
	}
	save(psi, end_state, x_size, "psi_states.dat");

	for (int i = 0; i < end_state; i++) {
		cout << " E[" << i << "] =" << Energy[i] << endl;
	}

	int energy_size = sizeof Energy / sizeof(Energy[0]);
	save(Energy, energy_size, "energylevels.dat");
	


	//----------------- Propagation ---------------------
	Wavefunction psi_start(x_size, dx);
	Wavefunction psi_prop(x_size, dx);
	Ctridiag Ht(H_0);
	Ctridiag V_las(x_size, zero, zero, zero);

	Ctridiag V_imag(x_size, zero, zero, zero);
	double x0 = x_list[x_size-1].real() - 30;
	Creal V_imag_value = 1i * 1e-3;
	Potential v_imag_plot(x_size, dx);
	v_imag_plot.init_imag(x_list, x0, V_imag_value, 2.0);
	V_imag = v_imag_plot;

	// propagation time 
	double t_prop_start = -7.5;
	dt = 0.01;
	double t_prop_end = 5;
	const unsigned int t_size = (int)(t_prop_end - t_prop_start) / dt;
	Cvec t_prop(t_prop_start, dt, t_prop_end); // SI
	t_prop.save("t_list.dat");
	t_prop = t_prop / (au_T * 1e15); // au
	dt = t_prop[2].real() - t_prop[1].real();
	

	Cvec E_las(t_prop.get_size());
	
	Creal** psi_time = new Creal * [t_size];
	for (int i = 0; i < t_size; i++) {
		psi_time[i] = new Creal[x_size];
	}

	for (int i = 0; i < x_size; i++) psi_start[i] = sqrt(0.5)*(psi[0][i] + psi[1][i]);
	psi_prop = psi_start;

	Pulse Pulse1(50.0, 1e13, 0.0, 0.0, 2.5);
	Pulse Pulse2(800.0, 1e11, 2.0, 0.0, 2.5);


	for (int ti = 0; ti < t_prop.get_size(); ti++) {
		double t = t_prop[ti].real();

		E_las[ti] = Pulse1.getField(t) + Pulse2.getField(t);

		for (int xi = 0; xi < x_size; xi++) {
			double x = real(x_list[xi]);
			double Vi = au_e * E_las[ti].real() * x;
			V_las.diag[xi] = Vi;
		}

		Ht = H_0 + V_las + V_imag;
		Am = One + Ht * 1i * (dt / 2.0);
		Ap = One - Ht * 1i * (dt / 2.0);

		Ap.matmul(b, psi_prop);	//b = Ap * psi_prop; //
		Am.solve_ref(psi_prop, b, help_am);	// psi_prop = b / Am;

		for (int i = 0; i < x_size; i++) {
			psi_time[ti][i] = psi_prop[i];
		}
	}

	save(psi_time, t_prop.get_size(), x_size, "psi_over_time.dat");




	// Window-Operator
	Cvec E_list(-20.0, 0.05, 40.0);
	N = psi_start.scalar(psi_prop);
	Wavefunction psi_final = psi_prop - psi_start*N;

	int n = 1;
	double gamma = 0.0025;
	double C = pow(2.0, n) * sin(M_PI / pow(2.0, n)) / (2 * M_PI * gamma);






	time(&end);
	double dif = difftime(end, start);
	printf("Elapsed time is %.5lf seconds.", dif);
	cout << endl;





	for (int i = 0; i < max_state; i++) delete[] psi[i];
	delete[] psi;
	for (int i = 0; i < t_prop.get_size(); i++) delete[] psi_time[i];
	delete[] psi_time;
	return 0;
}
