#include <bitmapfonts.c>
//#include <stdio.h>
#include <squarematrix.c>
#include <rnd/sq.c>

typedef struct {REAL xc,xcp,*b,*bp;} OrbitFields;

OrbitFields OrbitFields_new(const int N)
{
	OrbitFields ret; ret.b=array_1d_zeroed(N);
	ret.bp=array_1d_zeroed(N);
	ret.xc=ret.xcp=0;
	return ret;
}

OrbitFields OrbitFields_copy(const OrbitFields x)
{
	OrbitFields ret; ret.xc=x.xc; ret.xcp=x.xcp;
	ret.b=array_1d_copy(x.b);
	ret.bp=array_1d_copy(x.bp);
	return ret;
}

void OrbitFields_addmul(OrbitFields *px,const OrbitFields y,const double z)
{
	px->xc+=y.xc*z; px->xcp+=y.xcp*z;
	for (int n=array_1d_m(y.b)-1;n>=0;n--)
		{px->b[n]+=y.b[n]*z; px->bp[n]+=y.bp[n]*z;}
}

void OrbitFields_print(const OrbitFields x)
{
	int n,N=array_1d_m(x.b);
	printf("xc=%lg x'c=%lg\n",x.xc,x.xcp);
	for (n=0;n<N;n++) printf("b%d=%lg%c",n,x.b[n],(n==N-1?'\n':' '));
	for (n=0;n<N;n++) printf("b'%d=%lg%c",n,x.bp[n],(n==N-1?'\n':' '));
}

void OrbitFields_free(OrbitFields x)
{
	array_1d_free(x.b);
	array_1d_free(x.bp);
}

void OrbitFields_initialise(OrbitFields *px,const REAL *d,const REAL p0)
{ // Produces consistent initial conditions; user enters b[0..N-2], bp[0..N-1], xc; this calculates b[N-1] and xcp
	int n,N=array_1d_m(px->b); REAL t=0;
	for (n=N-2;n>=0;n--) t+=px->b[n];
	px->b[N-1]=-t;
	REAL dr=d[N-1]; t=0;
	for (n=N-2;n>=0;n--)
	{
		t+=px->b[n]*dr;
		dr+=d[n];
	}
	px->xcp=-t/(p0*dr);
}

OrbitFields OrbitFields_deriv(const OrbitFields x,const REAL *d,const REAL p)
{
	int n,i,j,N=array_1d_m(x.b); OrbitFields ret=OrbitFields_new(N);
	REAL *dx[3],*dxp[3];
	for (n=2;n>=0;n--)
	{
		dx[n]=array_1d(N+1); dxp[n]=array_1d(N+1);
		dx[n][0]=(n==0); dxp[n][0]=(n==1);
	}
	for (n=1;n<=N;n++)
	{
		for (i=2;i>=0;i--) dx[i][n]=dx[i][n-1]+d[n-1]*dxp[i][n-1];
		for (i=1;i>=0;i--) dxp[i][n]=dxp[i][n-1]+x.bp[n-1]*dx[i][n]/p;
		dxp[2][n]=dxp[2][n-1]-x.b[n-1]/(p*p)+x.bp[n-1]*dx[2][n]/p;
	}
	REAL D=2.0-dx[0][N]-dxp[1][N];
	ret.xc=((1.0-dxp[1][N])*dx[2][N]+dx[1][N]*dxp[2][N])/D; // dxc/dp
	ret.xcp=(dxp[0][N]*dx[2][N]+(1.0-dx[0][N])*dxp[2][N])/D; // dxcp/dp
	REAL *dpx=array_1d(N+1);
	for (n=N;n>=0;n--) dpx[n]=dx[0][n]*ret.xc+dx[1][n]*ret.xcp+dx[2][n];
	for (n=N-1;n>=0;n--) ret.b[n]=x.bp[n]*dpx[n+1];
	REAL **ddx[5],**ddxp[5]; // 0=x0x0 1=x0x'0 2=x'0x'0 3=x0p 4=x'0p [n][b''_j coefficient (0=none)]
	for (n=4;n>=0;n--)
	{
		ddx[n]=array_2d(N+1,N+1); ddxp[n]=array_2d(N+1,N+1);
		for (j=N;j>=0;j--) {ddx[n][0][j]=0; ddxp[n][0][j]=0;}
	}
	for (n=1;n<=N;n++)
	{
		for (i=4;i>=0;i--) for (j=N;j>=0;j--)
		{
			ddx[i][n][j]=ddx[i][n-1][j]+d[n-1]*ddxp[i][n-1][j];
			ddxp[i][n][j]=ddxp[i][n-1][j]+x.bp[n-1]*ddx[i][n][j]/p;
		}
		ddxp[0][n][n]+=sq(dx[0][n])/p;
		ddxp[1][n][n]+=dx[0][n]*dx[1][n]/p;
		ddxp[2][n][n]+=sq(dx[1][n])/p;
		ddxp[3][n][0]-=x.bp[n-1]*dx[0][n]/(p*p);
		ddxp[4][n][0]-=x.bp[n-1]*dx[1][n]/(p*p);
		ddxp[3][n][n]+=dx[0][n]*dx[2][n]/p;
		ddxp[4][n][n]+=dx[1][n]*dx[2][n]/p;
	}
	for (n=2;n>=0;n--) {array_1d_free(dx[n]); array_1d_free(dxp[n]);} // Freeing here, but perhaps could export auxiliary info about what phi_x, phi_y are?
	REAL cx,cy,*ax=array_1d(N),*ay=array_1d(N); // cx,cy are minus because they are moved to opposite side of formula
	cx=-((ddx[0][N][0]+ddxp[1][N][0])*ret.xc+(ddx[1][N][0]+ddxp[2][N][0])*ret.xcp+ddx[3][N][0]+ddxp[4][N][0]);
	for (n=N;n>0;n--)
		ax[n-1]=(ddx[0][N][n]+ddxp[1][N][n])*ret.xc+(ddx[1][N][n]+ddxp[2][N][n])*ret.xcp+ddx[3][N][n]+ddxp[4][N][n];
	for (n=4;n>=0;n--) {array_2d_free(ddx[n]); array_2d_free(ddxp[n]);}
	REAL *dy[2],*dyp[2];
	for (n=1;n>=0;n--)
	{
		dy[n]=array_1d(N+1); dyp[n]=array_1d(N+1);
		dy[n][0]=(n==0); dyp[n][0]=(n==1);
	}
	for (n=1;n<=N;n++) for (i=1;i>=0;i--)
	{
		dy[i][n]=dy[i][n-1]+d[n-1]*dyp[i][n-1];
		dyp[i][n]=dyp[i][n-1]-x.bp[n-1]*dy[i][n]/p;
	}
	REAL **ddy[2],**ddyp[2]; // 0=y0p 1=y'0p [n][b''_j coefficient (0=none)]
	for (n=1;n>=0;n--)
	{
		ddy[n]=array_2d(N+1,N+1); ddyp[n]=array_2d(N+1,N+1);
		for (j=N;j>=0;j--) {ddy[n][0][j]=0; ddyp[n][0][j]=0;}
	}
	for (n=1;n<=N;n++)
	{
		for (i=1;i>=0;i--) for (j=N;j>=0;j--)
		{
			ddy[i][n][j]=ddy[i][n-1][j]+d[n-1]*ddyp[i][n-1][j];
			ddyp[i][n][j]=ddyp[i][n-1][j]-x.bp[n-1]*ddy[i][n][j]/p;
		}
		for (i=1;i>=0;i--)
		{
			ddyp[i][n][0]+=x.bp[n-1]*dy[i][n]/(p*p);
			ddyp[i][n][n]-=dpx[n]*dy[i][n]/p;
		}
	}
	cy=-(ddy[0][N][0]+ddyp[1][N][0]);
	for (n=N;n>0;n--) ay[n-1]=ddy[0][N][n]+ddyp[1][N][n];
	for (n=1;n>=0;n--)
	{
		array_1d_free(dy[n]); array_1d_free(dyp[n]);
		array_2d_free(ddy[n]); array_2d_free(ddyp[n]);
	}
	REAL **aa=array_2d(2,2),t; // Matrix of dot products
	for (i=1;i>=0;i--) for (j=1;j>=0;j--)
	{
		t=0;
		for (n=N-1;n>=0;n--) t+=(i?ay:ax)[n]*(j?ay:ax)[n];
		aa[i][j]=t;
	}
	invertmatrix(aa);
	REAL cax=aa[0][0]*cx+aa[0][1]*cy,cay=aa[1][0]*cx+aa[1][1]*cy;
	array_2d_free(aa); // No N-2 additional tuning functions, so only good for N=2 so far
	for (n=N-1;n>=0;n--) ret.bp[n]=(cax*ax[n]+cay*ay[n]/*+gam[n]*/)*dpx[n+1];
	array_1d_free(ax); array_1d_free(ay);
	array_1d_free(dpx);
	return ret;
}

void OrbitFields_rk4step(OrbitFields *f,const REAL *d,const REAL p,const REAL dp)
{
	OrbitFields d1=OrbitFields_deriv(*f,d,p),fa=OrbitFields_copy(*f);
	OrbitFields_addmul(&fa,d1,dp*0.5);
	OrbitFields d2=OrbitFields_deriv(fa,d,p+dp*0.5); OrbitFields_free(fa);
	fa=OrbitFields_copy(*f); OrbitFields_addmul(&fa,d2,dp*0.5);
	OrbitFields d3=OrbitFields_deriv(fa,d,p+dp*0.5); OrbitFields_free(fa);
	fa=OrbitFields_copy(*f); OrbitFields_addmul(&fa,d3,dp);
	OrbitFields d4=OrbitFields_deriv(fa,d,p+dp); OrbitFields_free(fa);
	OrbitFields_addmul(f,d1,dp/6.0); OrbitFields_free(d1);
	OrbitFields_addmul(f,d2,dp/3.0); OrbitFields_free(d2);
	OrbitFields_addmul(f,d3,dp/3.0); OrbitFields_free(d3);
	OrbitFields_addmul(f,d4,dp/6.0); OrbitFields_free(d4);
}

REAL **OrbitFields_orbit(const OrbitFields x,const REAL *d,const REAL p)
{ // Returns a [2][N+1] array_2d of positions (first and last should be equal)
	int n,N=array_1d_m(x.b);
	REAL **ret=array_2d(2,N+1);
	ret[0][0]=x.xc; ret[1][0]=x.xcp;
	for (n=1;n<=N;n++)
	{
		ret[0][n]=ret[0][n-1]+d[n-1]*ret[1][n-1];
		ret[1][n]=ret[1][n-1]+x.b[n-1]/p;
	}
	return ret;
}

#include <muon1/PV2.c>

PV2 OrbitFields_track(const OrbitFields x,const REAL *d,const REAL p,PV2 s,const int n0,const int n1)
{
	REAL **o=OrbitFields_orbit(x,d,p);
	for (int n=n0;n<n1;n++)
	{
		s.x+=d[n]*s.xp;
		s.xp+=(x.b[n]+x.bp[n]*(s.x-o[0][n+1]))/p; // Linear approximation to the field around the closed orbit
		s.y+=d[n]*s.yp;
		s.yp-=x.bp[n]*s.y/p;
	}
	array_2d_free(o);
	return s;
}

PV2 OrbitFields_trackall(const OrbitFields x,const REAL *d,const REAL p,const PV2 s)
{
	return OrbitFields_track(x,d,p,s,0,array_1d_m(x.b));
}

#include <muon1/celloptics.c>

CellOptics OrbitFields_CellOptics(const OrbitFields x,const REAL *d,const REAL p)
{
	int i,j; const REAL eps=1e-6;
	CellOptics ret=CellOptics_new();
	ret.xo.x=x.xc; ret.xo.xp=x.xcp; ret.xo.y=ret.xo.yp=0;
	PV2 s,xoo=OrbitFields_trackall(x,d,p,ret.xo); // Strictly xoo ought to equal ret.xo but there may be p-integration errors
	for (i=0;i<4;i++)
	{
		s=ret.xo; s.e[i]+=eps;
		s=OrbitFields_trackall(x,d,p,s);
		for (j=0;j<4;j++) ret.m[j][i]=(s.e[j]-xoo.e[j])/eps;
	}
	CellOptics_calctwiss(&ret);
	return ret;
}

char *momentummarker(REAL p1,REAL p2)
{
	const char *mks_multipliers="a..f..p..n..u..mcd Dhk..M..G..T..P..E";
	p1*=299.792458e6;
	p2*=299.792458e6;
	int e=floor(log10(max(p1,p2))),
		i1=floor(min(p1,p2)/pow(10,e)),
		i2=floor(max(p1,p2)/pow(10,e));
	if (i1!=i2) if (i2==1 || i2==2 || i2==5)
	{
		char str[100],*s;
		sprintf(str,"%d %ceV/c",(int)(i2*pow(10,e%3)),mks_multipliers[18+e/3*3]);
		streqi(&s,str); return s;
	}
	return NULL;
}

OrbitFields fodo(const REAL *d,const REAL p)
{
	OrbitFields ret=OrbitFields_new(2);
	ret.b[0]=1;
	if (KEY(VK_SPACE))
	{
		ret.bp[0]=1; ret.bp[1]=-1;
	}
	else
	{
		ret.bp[0]=4.0*my/yres-2.0; ret.bp[1]=-(4.0*mx/xres-2.0);
	}
	ret.xc=0;
	OrbitFields_initialise(&ret,d,p);
	return ret;
}

//void main(void)
WINMAIN
{
	/*OrbitFields f=OrbitFields_new(2); REAL d[2]={1,1},p0=1; // p has units of 299.792458 MeV/c if b=Bl
	f.b[0]=1;
	f.bp[0]=1; f.bp[1]=-1;
	f.xc=0;
	OrbitFields_initialise(&f,d,p0);
	OrbitFields_print(f);
	OrbitFields fp=OrbitFields_deriv(f,d,p0);
	OrbitFields_print(fp);
	REAL **o=OrbitFields_orbit(f,d,p0);
	for (int n=0;n<=2;n++) printf("(x,x')_%d = (%lg,%lg)\n",n,o[0][n],o[1][n]);
	array_2d_free(o);
	getchar();*/
	int n,i,px,py;
	GRAPHICS(0,0,"Integration of Tune Fixing Equations [Stephen Brooks 2010]");
	loadfont("font.dat","normal");
	unsigned *col=palto24bit();
	REAL zoom=xres/4,bzoom=yres/6;
	LOOP
	{
		line(0,yres/2,xres,yres/2,0x808080); //line(xres/2,0,xres/2,yres,0x808080);
		REAL d[2]={1,1},p0=1;
		for (i=-1;i<=1;i+=2)
		{
			REAL p=p0,dp,dpp=0.01/**exp(5.0*my/yres)*/; // p has units of 299.792458 MeV/c if b=Bl
			OrbitFields f=fodo(d,p);//,fp;
			while (fabs(log(p/p0))<4)
			{
				dp=dpp*i*p;
				REAL **o=OrbitFields_orbit(f,d,p);
				char *s=momentummarker(p,p+dp);
				for (n=0;n<=2;n++)
				{
					putpix(px=xres/2+o[0][n]*zoom,py=yres/2-f.b[n%2]*bzoom,col[9+n]);
					if (s)
					{
						bfwrite(px,py,s,GFX_white);
						CellOptics co=OrbitFields_CellOptics(f,d,p);
						char str[100]; sprintf(str,"Qx = %lg\tQy = %lg\nDelta(x,x') (%lg,%lg)",
							co.phix/M_TWOPI,co.phiy/M_TWOPI,
							o[0][2]-o[0][0],o[1][2]-o[1][0]);
						bfwrite(px,py+BF_height,str,0xFF00FF);
						CellOptics_free(co);
					}
				}
				if (s) free(s);
				array_2d_free(o);
				/*fp=OrbitFields_deriv(f,d,p);
				OrbitFields_addmul(&f,fp,dp);
				OrbitFields_free(fp);*/
				OrbitFields_rk4step(&f,d,p,dp);
				p+=dp;
			}
			OrbitFields_free(f);
		}
		ccl();
	}
}