////////////////////////////////////////////////////////////////////
// シュレディンガー方程式の数値計算（クーロンポテンシャル）
////////////////////////////////////////////////////////////////////
#define _USE_MATH_DEFINES
#include <math.h>
#include <stdlib.h>
#include <time.h>
#include <iostream>
#include <fstream>
#include <sstream>
#include <string>
#include <cstdio>
#include <iomanip>
#include <stdio.h>
#include <complex>
#include <omp.h>
#include <direct.h>

/////////////////////////////////////////////////////////////////
//物理定数
/////////////////////////////////////////////////////////////////
//光速
const double c = 2.99792458E+8;
//真空の透磁率
const double mu0 = 4.0*M_PI*1.0E-7;
//真空の誘電率
const double epsilon0 = 1.0 / (4.0*M_PI*c*c)*1.0E+7;
//プランク定数
const double h = 6.6260896 * 1.0E-34;
double hbar = h / (2.0*M_PI);
//電子の質量
const double me = 9.10938215 * 1.0E-31;
//電子ボルト
const double eV = 1.60217733 * 1.0E-19;
//複素数
const std::complex<double> I = std::complex<double>(0.0, 1.0);
/////////////////////////////////////////////////////////////////
//物理系の設定
/////////////////////////////////////////////////////////////////
//空間の大きさ
const int Lz = 5000, Ld = 10;
//空間分割サイズ
double dz = 1.0E-11;
//重ね合わせの数
const int Npulse = 500;
//パルスの幅
double delta_z = 0.5E-8;
double sigma = 2.0*sqrt(2.0*log(2.0)) / (delta_z);
//波数の間隔
double dk = 30.0 / (delta_z * double(Npulse + 1));
//z方向の原点の場所
const int iz0 = Lz / 2 - Lz / 10;
//電子波のエネルギー
const double E0 = 0.0 * eV;
//波数の中心
double k0 = sqrt(2.0 * me * E0 / pow(hbar, 2));
//角振動数の中心
double omega0 = hbar / (2.0*me) * pow(k0, 2);

//計算ステップ数
int Nt = 40000;
//実空間分布保存
int skip = 100;
//時間間隔
double dt = 1.0 * 1.0E-18;

//ルンゲクッタ波数の個数（E0=10.eVの場合、Npulseの10倍程度が必要）
const int Ncutoff = 4096;
//系のサイズ（周期系）
double L = 2.0 * M_PI / dk;

//波束の初期位置
double x0 = 0;// L / 2.0;//2000*dz;


//中心エネルギー_波束の重ね合わせ数_時間スケール_FFT分割数
std::string fn = "_Coulomb_-1.0_500_4096_2";
//ファイル名
std::string folder = "output";//作成するフォルダ名
std::string filepath1 = folder + "/an" + fn + ".txt";
std::string filepath2 = folder + "/pulse" + fn + ".txt";
std::string filepath3 = folder + "/wavenumber" + fn + ".txt";


const int num_threads = 10;

//クーロン斥力
double V(double x) {

	double r = abs(x - 0.0 * dz);
	if (r < 10 * dz) r = 10 * dz;
	double bar_r = 1000.0 * dz;
	return -1.0*eV / ( r / bar_r);

}
/////////////////////////////////////////////////////////////////
//高速フーリエ変換
/////////////////////////////////////////////////////////////////
void FFT(double* an, double* bn, int N, bool Inverse) {
	/////////////////////////////////////////
	// 入力
	// N  ： 項数（2のべき乗）
	// an : 実数配列（順変換：実数空間データを項数Nで指定、逆変換：展開係数a(n)）
	// bn : 実数配列（順変換：虚数空間データを項数Nで指定、逆変換：展開係数b(n)）
	// Inverse : 逆変換の場合に true
	/////////////////////////////////////////
	// 出力
	// an : 実数配列（順変換：展開係数a(n)、逆変換：実数空間データ）
	// bn : 実数配列（順変換：展開係数b(n)、逆変換：虚数空間データ）
	/////////////////////////////////////////

	int ff = Inverse ? 1 : -1;

	int* rot = new int[N];

	for (int i = 0; i < N; i++) rot[i] = 0;

	int nhalf = N / 2, num = N / 2;
	double sc = 2.0 * M_PI / N;

	while (num >= 1) {
		for (int j = 0; j < N; j += 2 * num) {
			int phi = rot[j] / 2;
			int phi0 = phi + nhalf;
			double c = cos(sc * phi);
			double s = sin(sc * phi * ff);
			for (int k = j; k < j + num; k++) {
				int k1 = k + num;
				double a0 = an[k1] * c - bn[k1] * s;
				double b0 = an[k1] * s + bn[k1] * c;
				an[k1] = an[k] - a0;
				bn[k1] = bn[k] - b0;
				an[k] = an[k] + a0;
				bn[k] = bn[k] + b0;
				rot[k] = phi;
				rot[k1] = phi0;
			}
		}
		num = num / 2;
	}

	for (int i = 0; i < N; i++) {
		int j = rot[i];
		if (j > i) {
			double tmp = an[i];
			an[i] = an[j];
			an[j] = tmp;
			tmp = bn[i];
			bn[i] = bn[j];
			bn[j] = tmp;
		}
	}
	if (!Inverse) {
		for (int i = 0; i < N; i++) {
			an[i] = an[i] / N;
			bn[i] = bn[i] / N;
		}
	}
	delete[] rot;
}


/////////////////////////////////////////////////////////////////
//ルンゲクッタクラス
/////////////////////////////////////////////////////////////////
class RK4 {
private:
public:
	//プロパティ
	double dt = 0.01; //時間間隔
	int N;
	double L;
	std::complex<double>* an, *a1n, *a2n, *a3n, *a4n, *_a1n, *_a2n, *_a3n, *_a4n;
	std::complex<double>* dan;
	std::complex<double>* vn;
	//コンストラクタ
	RK4() {}
	RK4(int _N, double _dt, double _L) {
		N = _N;
		L = _L;
		dt = _dt;
		an = new std::complex<double>[N];
		a1n = new std::complex<double>[N];
		a2n = new std::complex<double>[N];
		a3n = new std::complex<double>[N];
		a4n = new std::complex<double>[N];
		_a1n = new std::complex<double>[N];
		_a2n = new std::complex<double>[N];
		_a3n = new std::complex<double>[N];
		_a4n = new std::complex<double>[N];
		dan = new std::complex<double>[N];
		vn = new std::complex<double>[N];
	}
	//ディストラクタ
	~RK4() {
		delete[] an, a1n, a2n, a3n, a4n, dan;
		delete[] vn;
	}
	//速度ベクトルを与えるメソッド
	void V(double t, std::complex<double> *an, std::complex<double> *out_vs) {

#pragma omp parallel for num_threads(num_threads)
		for (int n = 0; n < N; n++) {

			std::complex<double> kn;

			if (n < N / 2) {
				kn = 2.0 * M_PI / L * n;
			}
			else {
				kn = -2.0 * M_PI / L * (N - n);
			}

			out_vs[n] = hbar * kn * kn / (2.0 * me * I)* an[n];

			for (int l = 0; l < N; l++) {
				if (n >= l) {
					int m = n - l;
					out_vs[n] += vn[l] / (I * hbar) * an[m];
				}
				else {
					int m = N - (l - n);
					out_vs[n] += vn[l] / (I * hbar) * an[m];
				}
			}
		}
	}
	//時間発展を計算するメソッド
	void timeEvolution(double t) {
		//１段目
		V(t, an, a1n);

		for (int i = 0; i < N; i++) {
			_a1n[i] = an[i] + a1n[i] * dt / 2.0;
		}
		//２段目
		V(t + dt / 2.0, _a1n, a2n);
		for (int i = 0; i < N; i++) {
			_a2n[i] = an[i] + a2n[i] * dt / 2.0;
		}
		//３段目
		V(t + dt / 2.0, _a2n, a3n);
		for (int i = 0; i < N; i++) {
			_a3n[i] = an[i] + a3n[i] * dt;
		}
		//４段目
		V(t + dt, _a3n, a4n);
		for (int i = 0; i < N; i++) {
			dan[i] = dt / 6.0 * (a1n[i] + 2.0 * a2n[i] + 2.0 * a3n[i] + a4n[i]);
		}
	}
};


/////////////////////////////////////////////////////////////////
const int precision_N = 5;

int main() {
	{
		double kn = 2.0 * M_PI / L * 50;
		double E = hbar * kn * hbar * kn / (2.0 * me) / eV;
		std::cout << E << " " << exp(-1.0 / 2.0 * pow((kn - k0) / sigma, 2));
	}

	//フォルダ生成
	_mkdir(folder.c_str());

	//順変換用
	double vn_real[Ncutoff], vn_imag[Ncutoff];

	//ポテンシャル障壁の
	for (int n = 0; n < Ncutoff; n++) {
		double x = L / Ncutoff * n;
		if (n < Ncutoff / 2) {
			x = L / Ncutoff * n;
		}
		else {
			x = -L / Ncutoff * (Ncutoff - n);
		}

		vn_real[n] = V(x);
		vn_imag[n] = 0;
	}
	//順変換の実行
	FFT(vn_real, vn_imag, Ncutoff, false); //vn_real, vn_imag にフーリエ変換後の係数が格納される

										   //出力ストリームによるファイルオープン
	std::ofstream fout3;
	fout3.open(filepath3.c_str());
	fout3 << "#x:波数" << std::endl;
	fout3 << "#y:振幅" << std::endl;
	fout3 << "#legend: 波動関数の実部 波動関数の虚部 ポテンシャルの実部 ポテンシャルの虚部" << std::endl;
	fout3 << "#showLines: true true true true false false" << std::endl;
	fout3 << "#showMarkers: false false false false true true" << std::endl;
	fout3 << "#fills: false false false false" << std::endl;
	fout3 << "#times:1.0 1.0 0.2 0.2" << std::endl;

	//fout3 << "#xrange:" << -Lz / 2 << " " << Lz / 2 << " " << Lz / 10 << std::endl;
	//fout3 << "#yrange:" << -0.4 << " " << 1 << " " << 0.1 << std::endl;

	//ルンゲ・クッタ法用オブジェクト
	RK4 rk4(Ncutoff, dt, L);

	for (int n = 0; n < Ncutoff; n++) {

		double kn;
		if (n < Ncutoff / 2) {
			kn = 2.0 * M_PI / L * n;
		}
		else {
			kn = -2.0 * M_PI / L * (Ncutoff - n);
		}
		rk4.an[n] = exp(-I * kn * x0) * exp(-1.0 / 2.0 * pow((kn - k0) / sigma, 2));
		rk4.vn[n] = std::complex<double>(vn_real[n], vn_imag[n]);

		if (n < Ncutoff / 2) {
			fout3 << n;
		}
		else {
			fout3 << -(Ncutoff - n);
		}

		fout3 << " " << rk4.an[n].real() << " " << rk4.an[n].imag() << " " << rk4.vn[n].real() / eV << " " << rk4.vn[n].imag() / eV << std::endl;

	}
	fout3.close();


	//逆変換用
	double _vn_real[Ncutoff], _vn_imag[Ncutoff];
	double _an_real[Ncutoff], _an_imag[Ncutoff];


	//出力ストリームによるファイルオープン
	std::ofstream fout2;
	fout2.open(filepath2.c_str());
	fout2 << "#x:位置" << std::endl;
	fout2 << "#y:確率振幅" << std::endl;
	fout2 << "#legend: ポテンシャル障壁 数値解" << std::endl;
	fout2 << "#showLines: true true true true false false" << std::endl;
	fout2 << "#showMarkers: false false false false true true" << std::endl;
	fout2 << "#fills: false false false false" << std::endl;
	fout2 << "#xrange:" << -2000 << " " << 5000 << " " << 1000 << std::endl;
	fout2 << "#yrange:" << -0.01 << " " << 0.05 << " " << 0.005 << std::endl;
	fout2 << "#all:" << std::endl;
	//fout2 << "-2000 0.03"<< std::endl;
	//fout2 << "-1000 0.03" << std::endl;
	//fout2 << "-1000 0" << std::endl;
	//fout2 << "1000 0" << std::endl;
	//fout2 << "1000 0.03" << std::endl;
	//fout2 << "2000 0.03" << std::endl;


	for (int n = 0; n < Ncutoff; n++) {
		_vn_real[n] = vn_real[n];
		_vn_imag[n] = vn_imag[n];
	}

	//逆変換の実行
	FFT(_vn_real, _vn_imag, Ncutoff, true);

	for (int n = 0; n < Ncutoff; n++) {

		double x = L / Ncutoff * n;
		if (n < Ncutoff / 2) {
			x = L / Ncutoff * n;
		}
		else {
			x = -L / Ncutoff * (Ncutoff - n);
		}

		fout2 << std::setprecision(precision_N);
		fout2 << x / dz << " " << _vn_real[n] / eV / 400.0 << std::endl;
	}

	fout2 << "#skip:4 " << std::endl;

	for (int tn = 0; tn <= Nt; tn++) {
		double t_real = dt * tn;
		std::cout << tn << std::endl;

		//出力
		if (tn % skip == 0) {
			fout2 << "#coma:" << tn << std::endl;

			for (int n = 0; n < Ncutoff; n++) {
				_an_real[n] = rk4.an[n].real();
				_an_imag[n] = rk4.an[n].imag();
			}

			//逆変換の実行
			FFT(_an_real, _an_imag, Ncutoff, true);

			double error = 0;
			for (int n = 0; n < Ncutoff; n++) {

				double x = L / Ncutoff * n;
				if (n < Ncutoff / 2) {
					x = L / Ncutoff * n;
				}
				else {
					x = -L / Ncutoff * (Ncutoff - n);
				}

				fout2 << std::setprecision(precision_N);
				//fout2 << x / dz << " " << sqrt(_an_real[n] * _an_real[n] + _an_imag[n] * _an_imag[n]) / Ncutoff << " " << _vn_real[n] / eV/400.0 <<  std::endl;
				fout2 << x / dz << " " << sqrt(_an_real[n] * _an_real[n] + _an_imag[n] * _an_imag[n]) / Ncutoff << std::endl;

			}
			fout2 << std::endl;

		}

		//ルンゲ・クッタ法による時間発展
		rk4.timeEvolution(t_real);

		for (int n = 0; n < Ncutoff; n++) {
			//位置の更新
			rk4.an[n] = rk4.an[n] + rk4.dan[n];
		}

	}

	fout2.close();

	//出力ストリームによるファイルオープン
	std::ofstream fout1;
	fout1.open(filepath1.c_str());

	for (int n = 0; n < Ncutoff; n++) {
		fout1 << std::setprecision(15);
		fout1 << n << " " << rk4.an[n].real() << " " << rk4.an[n].imag() << std::endl;
	}
	fout1.close();

}

/*
gccコンパイル
g++ main.cpp
*/