#include <complex> 
using namespace std;

#ifndef __CINT__
#include <fftw3.h>
#endif 

#define nomomentumtrans
//#define momentumtransbackandforth

#ifdef nomomentumtrans
#undef momentumtransbackandforth
#endif

typedef std::complex<double> cdouble;

//Solves for a vector u[0..n-1] the tridiagonal linear set given by equation (2.4.1 numerical recipies in c)
//a[0..n-1],b[0..n-1],c[0..n-1] and r[0..n-1] are input vectors and are not modified
//the tridiagonal matrix looks as following (all indices must be decreased by one when accessed by e.g. a[i]) :
// [ b1 c1 0   ...                ] [u1  ] [r1  ]
// | a2 b2 c2  ...                | |u2  | |r2  |
// |           ...                |*|... |=|... |
// |           ... aN-1 bN-1 cN-1 | |uN-1| |rN-1|
// [           ... 0    aN   bN   ] [uN  ] [rN  ]

extern void tridag(cdouble a[],cdouble b[],cdouble c[],cdouble r[],cdouble u[],cdouble gam[], unsigned long n);

extern void  numerov( double *g,long glen,double h,double *Y,bool secure=true);
typedef double (* functionPtr) (double x, void * params);

//extern double findminimum (functionPtr fptr,double start,double stop,int max_iter=100 );

extern char *itoa2(int val,char* str, int base);
typedef double *pdouble;

class TimePropagator{
public:
	TimePropagator(
		int xanz_,int xanzzoomed_,double Dx_,double Dt_,double hbar_,double mass_,int dim=1);
	~TimePropagator();
	virtual void forward(cdouble *phi,cdouble *summedVs,double *momentum=0,cdouble *summedV=0)=0;
	virtual void backward(cdouble *phi,cdouble *summedVs,double *momentum=0,cdouble *summedV=0)=0;
	
public:
	double *zero;
	int dim;
	int xanz;
	int xanzzoomed;
	double Dx;
	double Dt;
	double hbar;
	double mass;
};

#ifndef __CINT__
typedef fftw_complex *pfftw_complex;
#endif

class SplitOpChebyshev:public TimePropagator{//as fftw is used use fftw_malloc for array so that proper alignment is guaranteed
public:
	SplitOpChebyshev(int xanz_,int xanzzoomed_,double Dx_,double Dt_,double hbar_,double mass_,int dim=1,double Vmin_=-10,double Vmax_=10);
	~SplitOpChebyshev();
	virtual void propagate(cdouble *phi,cdouble *summedVs,double dir=1.,double *momentum=0,cdouble *summedV=0);//summedVs is at time t, summedV at time t+dt, we will use summedVs
	virtual void forward(cdouble *phi,cdouble *summedVs,double *momentum=0,cdouble *summedV=0);
	virtual void backward(cdouble *phi,cdouble *summedVs,double *momentum=0,cdouble *summedV=0);
protected:
#ifndef __CINT__
	fftw_plan fftw_forward_plan;
	fftw_plan fftw_backward_plan;
	int MaxCoeff;
	cdouble *CoeffTbl;
	cdouble *phitemp;
	cdouble *phiadd;
	fftw_complex *tmpdata;
	pfftw_complex psi[3];
	double *kinop;
	double kstep;
	double Vmin;
	double Vmax;
	double Tmin;
	double Tmax;
	double Emax;
	double Emin;
	double DT;
	double DV;
	double DE;
	double Pi;
#endif
};

class SplitOp:public TimePropagator{//as fftw is used use fftw_malloc for array so that proper alignment is guaranteed
public:
	SplitOp(int xanz_,int xanzzoomed_,double Dx_,double Dt_,double hbar_,double mass_,int dim=1);
	~SplitOp();
	virtual void propagate(cdouble *phi,cdouble *summedVs,double dir=1.,double *momentum=0,cdouble *summedV=0);//summedVs is at time t, summedV at time t+dt, we will use summedVs
	virtual void forward(cdouble *phi,cdouble *summedVs,double *momentum=0,cdouble *summedV=0);
	virtual void backward(cdouble *phi,cdouble *summedVs,double *momentum=0,cdouble *summedV=0);
protected:
#ifndef __CINT__
	fftw_plan fftw_forward_plan;
	fftw_plan fftw_backward_plan;
	fftw_complex *tmpdata;
	fftw_complex *tmpdata2;

	double Pi;
	cdouble *kinop;
	cdouble *kinopback;

	double kstep;
#endif
};

//for the timedependent schroedinger equation:
// I hbar d phi/dt = (-hbar^2/(2 m) d^2/dx^2+Sum_i( Vi(x)*SCALEi(t) ) phi
//written in crank nicholson scheme: with wavefunction u_position^time
// (I hbar/Dt -hbar^2/(Dx^2 2 m)-1/2 Sum_i(Vi_j SCALEi^{n+1}) u_j^{n+1})+hbar^2/(Dx^2 4 m) u_{j+1}^{n+1}+hbar^2/(Dx^2 4 m) u_{j-1}^{n+1}
//=(I hbar/Dt +hbar^2/(Dx^2 2 m)+1/2 Sum_i(Vi_j SCALEi^{n}) u_j^{n}  )-hbar^2/(Dx^2 4 m) u_{j+1}^{n}  -hbar^2/(Dx^2 4 m) u_{j-1}^{n}
//
//this can be written in a tridiagonal form (see function tridiag) with the following coefficients:
//
// a_j=hbar^2/(Dx^2 4 m)       b_j=(I hbar/ Dt - hbar^2/(Dx^2 2 m) - 1/2 Sum_i(Vi_j SCALEi^{n})    c_j=a_j
// as_j=-hbar^2/(Dx^2 4 m)     bs_j=(I hbar/ Dt + hbar^2/(Dx^2 2 m) + 1/2 Sum_i(Vi_j SCALEi^{n+1})   cs_j=-as_j
//
// TridiagonalMatrixOfa_b_c_ * (u_1^n+1,u_2^n+1,...)=TridiagonalMatrixOf_as_bs_cs * (u_1^n,u_2^n,...)
//
//where u_j^n are known and u_j^{n+1} are unknown for time forward propagation
class CrankNicholson:public TimePropagator{
public:
	CrankNicholson(int xanz_,int xanzzoomed_,double Dx_,double Dt_,double hbar_,double mass_,int dim=1);
	~CrankNicholson();
	virtual void forward(cdouble *phi,cdouble *summedV,double *momentum=0,cdouble *summedVs=0);
	virtual void backward(cdouble *phi,cdouble *summedV,double *momentum=0,cdouble *summedVs=0);
protected:
	double a;
	cdouble b1;
	cdouble cc;
	double as;
	cdouble bs1;
	double cs;

	cdouble *rr;
	cdouble *array_a;
	cdouble *array_b;
	cdouble *array_c;
	cdouble *phiadd;
	cdouble *gam;
	cdouble *phitemp;
};

extern void shiftwfn(cdouble *phi,int shift,int xanzzoomed);

class QMOperators{
public:
	QMOperators(TimePropagator *tp);
	~QMOperators();
	double MomentumExpectationValue(cdouble *phi,int dimnum=0);
	double PositionExpectationValue(cdouble *phi,int dimnum=0);
	void MomentumDisplace(cdouble *phi,double displacement,int offset,int dimnum=0);
	void FFTPowerSpectrum(cdouble *phi,double *spec);
protected:
#ifndef __CINT__
	int dim;
	int xanz;
	double kstep;
	double Dx;
	double Dt;
	double hbar;
	double mass;
	int xanzzoomed;
	fftw_plan fftw_forward_plan;
	fftw_plan fftw_backward_plan;
	fftw_complex *psi;
	fftw_complex *tmpdata;
#endif
};

enum timepropagationmethod { cranknicholson = 0 , splitopchebyshev = 1 , splitop = 2 };

extern void makevideo(
//input from matlab
double *phi_ini,double *phi_final,double **V,double Dx,double **SCALE,double Dt,int nanz,double hbar,double mass,double **oneOverLambda,
//info about data
int xanz,int tanz,int amountOfElectrodes,
//output to matlab
double *scale_optimized_vid,double *phi_r_vid,double *phi_i_vid,double *overlap,
int *phi_reached_offset,int xanzzoomed=-1,timepropagationmethod method=splitopchebyshev,int optTransOrSplit=1);

extern void optimalcontrol(
//input from matlab
double *phi_ini,double *phi_final,double **V,double Dx,double **SCALE,double Dt,int nanz,double hbar,double mass,double **oneOverLambda,
//info about data
int xanz,int tanz,int amountOfElectrodes,
//output to matlab
double **SCALEnew,cdouble *phi_reached,double *overlap,
int &phi_reached_offset,int xanzzoomed=-1,timepropagationmethod method=splitopchebyshev);

