//this is a C implementation to solve Ex 5.14 in DeVries
//the ode solver used will be RK-Fehlberg with adaptative step size
//the output are in four files (x and phi)
//	even1_phi, odd1_phi, even2_phi, potential
//	the wavefunctions will be normalized
#include <stdio.h>
#include <math.h>
#include "numcip.h"

//global variables for quantum well
double hbar;
double *phi;//phi[0]-x value, phi[1]-phi(x), phi[2]-phi'(x)
double xmin;//starting value for x
double phiprime;//starting value for phi'

//global variables for RKF
double *y4,*y5,*ystep,*yerr;
double **f;
double **a,*b,*c,*d;
double hstart;//initial estimate of step size
double tol;//tolerance level

double qwell(double estart,int flag);
double logdevfn(int flag);
void propagator(double eng);
void outputphi(double eng,int flag,FILE *fpointer);
double potential(double xx);
void qdev(double *dev,double xx,double yy1,double yy2,double eng);
double norm(double *yy);
double rkf(double eng, double h);
void rkpara(void);

main()
{
	double evenE1,evenE2,oddE;
	double x;
	double h;
	int N,i;

	FILE *outfile;

	//userio
	xmin=-5.0;
	phiprime=1e-5;
	hstart=1e-5;
	tol=1e-9;


	//memory allocations
	phi=dvector(0,2);
	y4=dvector(1,2);
	y5=dvector(1,2);
	ystep=dvector(1,2);
	yerr=dvector(1,2);
	f=dmatrix(0,5,1,2);
	a=dmatrix(1,5,1,6);
	b=dvector(1,4);
	c=dvector(1,5);
	d=dvector(1,5);
	rkpara();//initializing rk parameters

	//1st even energy
	evenE1=qwell(0.1,0);//start searching the 1st even energy at ~0 and 0-even
	outfile=fopen("qwellen1","w");
	fprintf(outfile,"xx\teven1\n");
	outputphi(evenE1,0,outfile);
	fclose(outfile);

	//1st odd energy
	oddE=qwell(evenE1+0.1,1);//start searching the 1st odd energy at evenE1+0.1 and 1-odd
	outfile=fopen("qwellod1","w");
	fprintf(outfile,"xx\todd1\n");
	outputphi(oddE,1,outfile);
	fclose(outfile);

	//2nd even energy
	evenE2=qwell(oddE+0.1,0);//start searching the 2nd even energy at oddE+0.1 and 0-even
	outfile=fopen("qwellen2","w");
	fprintf(outfile,"xx\teven2\n");
	outputphi(evenE2,0,outfile);
	fclose(outfile);

	//plot out potential
	outfile=fopen("qwellpol","w");
	fprintf(outfile,"xx\tpot\n");
	x=xmin;//starting point for plot
	N=500;//number of points to plot
	h=2.0*fabs(xmin)/N;
	for(i=1;i<=N+1;i++)
		{
		fprintf(outfile,"%20.15lg\t%20.15lg\n",x,potential(x));
		x+=h;
		}
	fclose(outfile);

	//other output
	outfile=fopen("qwellpar","w");
	fprintf(outfile,"Aharmonic Oscillator Eigenenergies\n");
	fprintf(outfile,"\n[1] 1st Even:%20.15lg\n",evenE1);
	fprintf(outfile,"\n[2] 1st Odd:%20.15lg\n",oddE);
	fprintf(outfile,"\n[3] 2nd Even:%20.15lg\n",evenE2);
	fclose(outfile);


	return 0;
}

double qwell(double estart,int flag)
//flag = 0 -> Even, flag = 1 -> Odd
//estart starting energy for search
//to search for the logarithmic derivative will be done by bisection
{
	double logdev,logdevL,logdevR;
	double eng,engL,engR;

	//first we need to bound the zero
	//to start with estart and move 0.5 eV up until sign changes
	engL=estart;
	propagator(engL),
	logdevL=logdevfn(flag);
	engR=estart;
	do
		{
		engR+=0.5;
		propagator(engR),
		logdevR=logdevfn(flag);
		}
	while(logdevL*logdevR>0.0);

	//now we refine the zero by bisection
	engL=engR-0.5;
	propagator(engL);
	logdevL=logdevfn(flag);
	do
		{
		eng=(engR+engL)/2.0;
		propagator(eng);
		logdev=logdevfn(flag);
		if(logdev*logdevL<0.0)
			{
			engR=eng;
			logdevR=logdev;
			}
		else
			{
			engL=eng;
			logdevL=logdev;
			}
		}
	while(fabs(logdev)>tol);//this check for absolute error

	return eng;
}

double logdevfn(int flag)
{
	if(flag)
		return phi[1]/phi[2];
	else
		return phi[2]/phi[1];
}

void propagator(double eng)
{
	double h;

	phi[0]=xmin;
	phi[1]=0.0;
	phi[2]=phiprime;
	h=hstart;

	do//inner loop to propogate from xmin<0 to x=0
		{
		//check if next step will be at center of well
		if(phi[0]+h>0.0)
			h=-phi[0];//adjust so that final step is exactly at x=0

		h=rkf(eng,h);//rkf will move trajectory (x,phi(x)) one step forward and return new step h
		}
	while(phi[0]<0.0);
}

void outputphi(double eng,int flag,FILE *fpointer)
{
	double h;
	double **outwave;
	double nfactor;
	int i,n;
	int sign=1;//sign for even or odd symmetry

	//set up temp storage
	outwave=dmatrix(1,2,1,500);

	//generate wave and find nfactor
	phi[0]=xmin;//starting from negative x
	phi[1]=0;//phi(xmin)
	phi[2]=phiprime;//phi'(xmin)
	h=hstart;//initial guess of step size
	n=0;
	do
		{
		//check if next step will be at center of well
		if(phi[0]+h>0.0)
			h=-phi[0];//adjust so that final step is exactly at x=0

		h=rkf(eng,h);//rkf will move trajectory (x,phi(x)) one step forward and return new step h

		//put wave in storage
		n++;
		outwave[1][n]=phi[0];
		outwave[2][n]=phi[1];
		}
	while(phi[0]<0.0);

	//find normalization factor by integrating phi*phi with trapezoid rule
	nfactor=0.0;
	for(i=2;i<=n;i++)
		{
		//area of trapezoid
		nfactor+=0.5*(outwave[1][i]-outwave[1][i-1])*
			(outwave[2][i]*outwave[2][i]+outwave[2][i-1]*outwave[2][i-1]);
		}
	nfactor*=2.0;//half of the wave
	nfactor=sqrt(nfactor);//phi*phi

	//output wave
	//first half
	for(i=1;i<=n;i++)
		fprintf(fpointer,"%20.15lg\t%20.15lg\n",outwave[1][i],outwave[2][i]/=nfactor);
	//second half
	if(flag)
		sign=-1;//odd symmetry for odd wavefunction
	for(i=n-1;i>=1;i--)
		fprintf(fpointer,"%20.15lg\t%20.15lg\n",-outwave[1][i],sign*outwave[2][i]);
	
	//free up allocation
	free_dmatrix(outwave,1,2,1,500);
}

double potential(double xx)
{
	double temp;

	temp=xx*xx;
	return 0.5*temp+0.25*temp*temp;
}

void qdev(double *dev,double xx,double yy1,double yy2,double eng)
{
	dev[1]=yy2;
	dev[2]=2.0/hbar*(potential(xx)-eng)*yy1;
}

double norm(double *yy)
{
	return sqrt(yy[1]*yy[1]+yy[2]*yy[2]);
}

double rkf(double eng, double h)
//the complication is that our ode is 2 dimensional
//	so that y and the derivative of y are 2D vectors
{
	double xstep;
	//ystep[2],y4[2],y5[2] are global variables
	//f[][] are global variables
	double thiserr,maxerr;
	double hnew;
	double temp;
	int n,i,j;

	//this is an implementation of eqs. 5.41-5.46
	qdev(f[0],phi[0],phi[1],phi[2],eng);

	do
	{
		for(n=1;n<=5;n++)
			{
			xstep=a[n][1]*h;
			for(i=1;i<=2;i++)
				{
				ystep[i]=a[n][2]*f[0][i];
				for(j=3;j<=n+1;j++)
					ystep[i]+=a[n][j]*f[j-2][i];
				ystep[i]*=h;
				}
			qdev(f[n],phi[0]+xstep,phi[1]+ystep[1],phi[2]+ystep[2],eng);
			}

		//calculate RK4 RK5 and error
		//RK4
		for(i=1;i<=2;i++)
			{
			y4[i]=phi[i];
			temp=b[1]*f[0][i];
			for(j=2;j<=4;j++)
				temp+=b[j]*f[j][i];
			y4[i]+=temp*h;
			}

		//RK5
		for(i=1;i<=2;i++)
			{
			y5[i]=phi[i];
			temp=c[1]*f[0][i];
			for(j=2;j<=5;j++)
				temp+=c[j]*f[j][i];
			y5[i]+=temp*h;
			}

		//error
		for(i=1;i<=2;i++)
			{
			yerr[i]=d[1]*f[0][i];
			for(j=2;j<=5;j++)
				yerr[i]+=d[j]*f[j][i];
			yerr[i]*=h;
			}
		
		//calculate new step
		maxerr=h*tol;
		thiserr=norm(yerr);
		hnew=0.9*h*sqrt(sqrt(maxerr/thiserr));

		//if thiserr is larger maxerr, reduce h to hnew and redo
		if(thiserr>maxerr)
			h=hnew;
		}
	while(thiserr>maxerr);

	//if thiserr is less than maxerr, accept estimates
	phi[0]+=h;
	phi[1]=y5[1];
	phi[2]=y5[2];

	return hnew;
}

void rkpara(void)
//inputing the coefficients for the RKF in Eqs. 5.41-5.49
//	treating the coefs in Eq. as an array as printed on pg. 220
{
	hbar=7.6199682; //hbar^2 in units of eV and A

	//a[][] is for f1,f2,...,f5
	a[1][1]=a[1][2]=0.25;

	a[2][1]=3.0/8.0;
	a[2][2]=3.0/32.0;
	a[2][3]=9.0/32.0;

	a[3][1]=12.0/13.0;
	a[3][2]=1932.0/2197.0;
	a[3][3]=-7200.0/2197.0;
	a[3][4]=7296.0/2197.0;

	a[4][1]=1.0;
	a[4][2]=439.0/216.0;
	a[4][3]=-8.0;
	a[4][4]=3680.0/513.0;
	a[4][5]=-845.0/4104.0;

	a[5][1]=0.5;
	a[5][2]=-8.0/27.0;
	a[5][3]=2.0;
	a[5][4]=-3544.0/2565.0;
	a[5][5]=1859.0/4104.0;
	a[5][6]=-11.0/40.0;

	//b[] is for the 4th order RK
	b[1]=25.0/216.0;
	b[2]=1408.0/2565.0;
	b[3]=2197.0/4104.0;
	b[4]=-0.2;

	//c[] is for the 5th order RK
	c[1]=16.0/135.0;
	c[2]=6656.0/12825.0;
	c[3]=28561.0/56430.0;
	c[4]=-9.0/50.0;
	c[5]=2.0/55.0;

	//d[] is for the error
	d[1]=1.0/360.0;
	d[2]=-128.0/4275.0;
	d[3]=-2197.0/75240.0;
	d[4]=1.0/50.0;
	d[5]=2.0/55.0;
}