#include <stdio.h>
#include <math.h>
#include <time.h>

#include <omp.h>
#define CHUNKSIZE 5

#include <gsl/gsl_math.h>
#include <gsl/gsl_matrix.h>
#include <gsl/gsl_odeiv.h>
#include <gsl/gsl_errno.h>

#include <essl.h>
#define CSET(x,a,b) (RE(x)=a, IM(x)=b)

/*************************************System parameters**********************************************/
#define L 3.5 /*Length of the box for particle-in-a-box states*/
#define U0 -10.0 /* Pseudopotential interaction strength*/
#define V0 5.0 /*Double Well Depth*/

#define EPSILON0 4.0 /*STIRAP Pulse Amplitude, analytically kept small wrt V0 for truncability*/

/****************************************************************************************************/

#define STIRAPWIDTH 2700.00 /*STIRAP pulse width T*/
#define TSTOKE 9000.00   /*centroid of the stoke pulse */
#define TPUMP 18000.00  /*centroid of the pump pulse */

/*Adiabatic Time*/
#define TFIXMIN 5600.00 
#define TFIXMAX 8317.9752  /*Up to ttot of 28009.465*/
#define TFIXINC 67.94938 /*Exact total of 20 oscillations b4 next increment (same as 1 osc) ideal  */ 

/******Duffing Drive Parameters*********************/
#define OMEGASTOKE 3.6987449
#define OMEGAPUMP 9.2468622 

/*omegastoke/omegapump is n1/n2 frequencies MUST BE COMMENSURATE CHECK THIS THOROUGHLY*/
#define N1 2
#define N2 5
/*******************************/

/***********************************Numerical Truncation Parameters********************************/ 
#define ALPHA 11 /*ALPHA+ ALPHA*(ALPHA-1)/2 is DIM. ALPHA is the # of 2-free-particles-in-a-box-states 8 works*/
#define DIM 66   /*Truncated Dimension of the Floquet Matrix generated from 2-free-particles-in-a-box states 36 works*/

/***********************************************ERROR TOLERANCES***********************************/ 
/*
The step-size adjustment procedure for this method begins by computing the desired error level D_i for each component,

D_i = ABSERROR + RELERROR * (YERROR |y_i| + YPRIMEERROR*DT |y'_i|)

and comparing it with the observed error E_i = |yerr_i|. If the observed error E exceeds the desired error level D by more than 10% for any component then the 
method reduces the step-size by an appropriate factor
*/
#define DT 1e-6
#define ABSERROR 1e-6
#define RELERROR 1e-8
#define YERROR 0.5
#define YPRIMEERROR 0.2

/********The Parameter epsilon & the tower of 2-particle states are stored in this datatype********/
typedef struct
        {
         double v0;
         double u0;
         double epsilonstoke;
         double epsilonpump;
         double omegastoke;
         double omegapump;
         int tower[DIM][2];
         }drive_and_tower;

/*****************************Defines Kronecker Delta**********************************************/
double kdel(int i,int j)
      {
        int result;
        if(i==j)
           {
             result=1.0;  
	   }
        else result=0.0;
        return(result);
      }	       

/*****************Calculates the Adiabatic Amplitude of the STIRAP Pulse**************************/

double amp(double tfix, double tmean)
          {
           double ampl;
           double beta=1.0/(2.0*STIRAPWIDTH*STIRAPWIDTH);
           ampl=EPSILON0*exp(-beta*(tfix-tmean)*(tfix-tmean));
           return(ampl);
          }

/*****************This is the lth single particle-in-a-box energy level*****************************/
double En(int l)
         {
          double energy;
          l=l+1;
          energy=M_PI/2.0;
          energy=energy*energy;
          energy=energy*l*l;
          energy=energy/(L*L);
          return(energy);
         }

/**************************** This is a matrix element of x*****************************************/
double f1(int m, int n)
        {
         double value;
         m=m+1;/*Arrays are counted from n=0, but levels from n=1*/
         n=n+1;
         if((m+n)%2==1)
           {
            value=16.0*m*n;
            value=value/(M_PI*M_PI);
            value=value/(m*m-n*n);
            value=value/(m*m-n*n);
            value=value*L;          
           }
         else
           {
            value=0.0;
           }
         return(value);
        }

/**************************** This is a matrix element of x^2**************************************/
double f2(int m, int n)
         {
           double value,temp;
           m=m+1; /*Arrays are counted from n=0, but levels from n=1*/
           n=n+1;
           if((m+n)%2==0)
             {
              if(m==n)
                {
                 value=((1/3.0)-(2.0/(n*n*M_PI*M_PI)))*L*L;
		}
              else
                {   
                 value=32.0*m*n*L*L/(M_PI*M_PI*(m*m-n*n)*(m*m-n*n));
		}
             }
           else
             {
              value=0.0;
             }
	   return(value);
         }

/**************************** This is a matrix element of x^4***************************************/
double f4(int m, int n)
         {
           double value;
           m=m+1; /*Arrays are counted from n=0, but levels from n=1*/
           n=n+1;
           if((m+n)%2==0)
             {
              if(m==n)
                {
		 value=L*L*L*L*((1/5.0)-(4.0/(n*n*M_PI*M_PI))+(24.0/(n*n*n*n*M_PI*M_PI*M_PI*M_PI)));
                }
              else
                {
                 value=(64.0*m*n/(M_PI*M_PI*M_PI*M_PI))*(pow(M_PI/(m*m-n*n),2)-(48.0*(m*m+n*n)/pow(m*m-n*n,4)))*L*L*L*L;
                } 
            }
           else
             {
              value=0.0;
             }
           return(value);
         }

/*Matrix elements of the Double Well Potential*/
double pot(int n, int m)
          {
           double f2(int,int),f4(int,int);
           double potelm;
           potelm=(-2.0*f2(n,m)+f4(n,m));
           return(potelm);
          }

/*Matrix Element of the Interaction*/
double in(int n1, int n2, int n3, int n4)
         {
          double interelm=0.0;
          n1++;
          n2++;
          n3++;
          n4++;
          if((n1==n2)&&(n2==n3)&&(n3==n4))
            {
             interelm=3.0/(4.0*L);
            }
          if((n1==n2)&&(n3==n4))
            {
             interelm=1.0/(2.0*L);
            }
         if((n1==n3)&&(n2==n4))
           {
             interelm=1.0/(2.0*L);
           }
         if((n1==n4)&&(n2==n3))    
           {
             interelm=1.0/(2.0*L);
           }
        return(interelm);
        }

/*This uses the above subroutines to evaluate the hamiltonian matrix elements*/
double hamilt(void *param, int nt, int mt, double t)
             {
              double kdel(int,int),En(int),in(int,int,int,int),pot(int,int);
              double matrixelm;
              int n1,n2,m1,m2,i,j;
              drive_and_tower *p=(drive_and_tower *)param;
              /*Get epsilon & towerofstates from structure param*/
              double v0=p->v0;
              double u0=p->u0;
              double epsilonstoke=p->epsilonstoke;
              double epsilonpump=p->epsilonpump;
              double omegastoke=p->omegastoke;
              double omegapump=p->omegapump;
              double towerofstates[DIM][2];
              for(i=0;i<DIM;i++)
                 {
                  towerofstates[i][0]=p->tower[i][0];
                  towerofstates[i][1]=p->tower[i][1];
                 }

              n1=towerofstates[nt][0];
              n2=towerofstates[nt][1];
              m1=towerofstates[mt][0];
              m2=towerofstates[mt][1];

              matrixelm=En(n1)+En(n2);
              matrixelm=matrixelm*(kdel(m1,n1)*kdel(m2,n2)+kdel(m1,n2)*kdel(m2,n1));

              matrixelm=matrixelm+v0*pot(n1,m1)*kdel(n2,m2);
              matrixelm=matrixelm+v0*pot(n1,m2)*kdel(n2,m1);
              matrixelm=matrixelm+v0*pot(n2,m1)*kdel(n1,m2);
              matrixelm=matrixelm+v0*pot(n2,m2)*kdel(n1,m1);
     
              matrixelm=matrixelm+2.0*u0*in(n1,n2,m1,m2);

              matrixelm=matrixelm+(epsilonstoke*sin(omegastoke*t)+epsilonpump*sin(omegapump*t))*f1(n1,m1)*kdel(n2,m2);
              matrixelm=matrixelm+(epsilonstoke*sin(omegastoke*t)+epsilonpump*sin(omegapump*t))*f1(n1,m2)*kdel(n2,m1);
              matrixelm=matrixelm+(epsilonstoke*sin(omegastoke*t)+epsilonpump*sin(omegapump*t))*f1(n2,m1)*kdel(n1,m2);
              matrixelm=matrixelm+(epsilonstoke*sin(omegastoke*t)+epsilonpump*sin(omegapump*t))*f1(n2,m2)*kdel(n1,m1);

              return(matrixelm);
             }

/*********************This evaluates the Differential Equation System*****************************/ 
int func(double t, const double y[], double dydt[], void *param)
        { 
          double y_re[DIM],y_im[DIM],dydt_re[DIM],dydt_im[DIM];
          int nt,mt;
	  /*Write out the (complicated) expression for RHS & assign to dydt[i]*/	
          for(nt=0;nt<DIM;nt++)
             {
              y_re[nt]=y[nt];
              y_im[nt]=y[nt+DIM];
             }
           for(nt=0;nt<DIM;nt++)
             {
               dydt_re[nt]=0.0;
               dydt_im[nt]=0.0;
               for(mt=0;mt<DIM;mt++)
                  {
                   dydt_re[nt]=dydt_re[nt]+hamilt(param,nt,mt,t)*y_im[mt];
                   dydt_im[nt]=dydt_im[nt]-hamilt(param,nt,mt,t)*y_re[mt];
                } 
              }
          for(nt=0;nt<DIM;nt++)
             {
              dydt[nt]=dydt_re[nt];
              dydt[DIM+nt]=dydt_im[nt];
             }
          return GSL_SUCCESS;
	 }		

/*****This is a DUMMY FUNCTION. IT DOES NOTHING. IT'S ONLY NEEDED FOR BSIMP & I'M USING RK8PD******/
int jac(double t, const double y[], double *dfdy, double dfdt[], void *param)
       {
         gsl_matrix_view dfdy_mat = gsl_matrix_view_array(dfdy,2*DIM,2*DIM);
         gsl_matrix *jacobian=&dfdy_mat.matrix;
	 return GSL_SUCCESS;	
	}

/*****This function actually runs the full integration of a particulat set of IC's from 0-T**********/ 
 void integrate(double *input, double initial, double final,void *param)
      {
       int i,j,status;
       double dt,t;
       /*This sets up the gsl ODE system structure*/
        const gsl_odeiv_step_type *T=gsl_odeiv_step_rk8pd;
	gsl_odeiv_step *s=gsl_odeiv_step_alloc(T,2*DIM);
        gsl_odeiv_control *c=gsl_odeiv_control_standard_new(ABSERROR,RELERROR,YERROR,YPRIMEERROR);
        gsl_odeiv_evolve *e=gsl_odeiv_evolve_alloc(2*DIM);
        gsl_odeiv_system sys={func, jac, 2*DIM, param};
	dt=DT;
        t=initial;
        while(t<final)
              {
                status=gsl_odeiv_evolve_apply(e,c,s,&sys,&t,final,&dt,input);
                if(status!=GSL_SUCCESS)
                  {		   
	           printf("\nGSL execution of differential equation solver failed."); 
                   printf("\nCheck program for bugs, bailing... \n");
                   break;
	          }
              }
        gsl_odeiv_evolve_free(e);
        gsl_odeiv_control_free(c);
        gsl_odeiv_step_free(s);
       }

/*************************************************************************************************/
int main (void)
      {
        time_t begin, end, rawtime;
        struct tm * timeinfo;
        drive_and_tower param;
        int i,j,select[DIM],naux=3*(2*DIM+1),towerofstates[DIM][2],buffer[2];
        int true,chunk=CHUNKSIZE,nthreads,tid,count;
        long errorcount=0;
	double comtest1,comtest2;
	double y[2*DIM],y_out[2*DIM],aux[3*(2*DIM+1)];
        double quasi[DIM],shuldbeone,temp;
        double tfix,period;
        double v0,u0,epsilonstoke,epsilonpump,omegastoke,omegapump;
        double pumptime=TPUMP, stoketime=TSTOKE;
        double eni,enj;	
        dcmplx floquet4essl[DIM][DIM],floqueteigenvalues[DIM],dummy[3];	

        FILE *floqueteigenval;

	gsl_matrix *floquetmatrix=gsl_matrix_alloc(2*DIM,DIM);

	floqueteigenval=fopen("result/stirap_fine_01.dat","w");
        begin=time(NULL);
        printf("\n=========================================================================");

        /*Compute the adiabatic period of the STIRAP pulse */
	comtest1=N1/N2;
	comtest2=(OMEGASTOKE/OMEGAPUMP);
        if((comtest1-comtest2)>=ABSERROR)
          {
	   printf("\n ERROR!!!\n");
           printf("\nFrequencies are not commensurate, or N1, N2 entries are wrong, bailing...");
	   printf("\n=========================================================================");
           exit(1);
          }

         /*Build the 2-particle tower of states*/
         count=0;
         for(i=0;i<ALPHA;i++)
            {
             for(j=0;j<=i;j++)
                {
                  towerofstates[count][0]=i;
                  towerofstates[count][1]=j;
                  count++;
                }
            }
         /*Now sort them in order of increasing free 2-particle energy n1^2+n2^2*/
         for(i=0;i<DIM;i++)
            {
             for(j=0;j<=i;j++)
                {
                 eni=(towerofstates[i][0]+1)*(towerofstates[i][0]+1)+(towerofstates[i][1]+1)*(towerofstates[i][1]+1);
                 enj=(towerofstates[j][0]+1)*(towerofstates[j][0]+1)+(towerofstates[j][1]+1)*(towerofstates[j][1]+1);
                 if(eni<enj)
                   {
                    buffer[0]=towerofstates[i][0];
                    buffer[1]=towerofstates[i][1];

                    towerofstates[i][0]=towerofstates[j][0];
                    towerofstates[i][1]=towerofstates[j][1];

                    towerofstates[j][0]=buffer[0];
                    towerofstates[j][1]=buffer[1];
                   } 
                }
            }
       u0=U0;
       v0=V0;
       omegastoke=OMEGASTOKE;
       omegapump=OMEGAPUMP;
       period=M_PI*((N1/omegastoke)+(N2/omegapump));
       printf("\nCOMMENSURATE TIME PERIOD OF BOTH PULSES = %lf",period);
       /*tfix  loop begins here*/
       for(tfix=TFIXMIN;tfix<=TFIXMAX;tfix=tfix+TFIXINC)
          {
           /*Now, put towerofstates, epsilon, vo & u0 in structure param*/
           param.v0=v0;
           param.u0=u0;
           param.omegastoke=omegastoke;
           param.omegapump=omegapump;
   
           epsilonstoke=amp(tfix,stoketime);
           param.epsilonstoke=epsilonstoke;
           epsilonpump=amp(tfix,pumptime);
           param.epsilonstoke=epsilonstoke;
      
           for(i=0;i<DIM;i++)
              {
               param.tower[i][0]=towerofstates[i][0];
               param.tower[i][1]=towerofstates[i][1];
              }

           printf("\nELAPSED TIME = %d sec",time(NULL)-begin);
           printf("\n\nEVOLVING FLOQUET MATRIX WITH TFIX =  %lf\n",tfix);

           /****************HERE IS WHERE THE PARALLEL FUN BEGINS********************************/
           #pragma omp parallel shared(period) private(i,j,y)
            {
             /*Build the Floquet Matrices (Real & Imaginary separate) using RK method in gsl*/
             #pragma omp for schedule(dynamic,chunk)
             for(j=0;j<DIM;j++)
                {
                 tid=omp_get_thread_num();
                 /*This is the jth column of the matrix*/
                 /* Initial conditions, */      
                 for(i=0;i<2*DIM;i++)
                    {
                     if(i<DIM)
	               {    
	                y[i]=kdel(i,j);
	               }
	             else
	               {
		        y[i]=0.0;
		       }	 
	             }		    
                 /*DONE. Now evolve the jth column of the floquet matrix (stored in array y[j]) across time period */ 

                 integrate(y,0.0,period,&param);

	         /*DONE, Now print out the final values of floquet matrix for that array into the gsl matrix*/
                 for(i=0;i<2*DIM;i++)
	            {
	             gsl_matrix_set(floquetmatrix,i,j,y[i]);
                    }	
	          /*DONE*/
                  if(tid==0)
                    {
                     nthreads=omp_get_num_threads();
                    }
	       }  
               /*End*/
         }
           /****************HERE IS WHERE THE PARALLEL FUN ENDS********************************/

           /*Formulate the floquet matrix from gsl form to essl form*/
           for(i=0;i<DIM;i++)
              {
               for(j=0;j<DIM;j++)
                  {
                   CSET(floquet4essl[i][j],gsl_matrix_get(floquetmatrix,i,j),gsl_matrix_get(floquetmatrix,i+DIM,j));
                  }
              }
           /* End the formulation of the essl floquet matrix in floquet4essl array*/

           /*We want to diagonalize the floquet matrix now*/
           zgeev(0,floquet4essl,DIM,floqueteigenvalues,dummy,2,select,DIM,aux,naux);
           /*End Diagonalization*/

	   /*Dump out the floquet quasienergies */
           printf("\n=========================================================================");
	   printf("\nChecking for Unitarity of Floquet Matrix...\n");
           for(i=0;i<DIM;i++)
              {
               quasi[i]=-atan2(IM(floqueteigenvalues[i]),RE(floqueteigenvalues[i]))/period;     
               shuldbeone=RE(floqueteigenvalues[i])*RE(floqueteigenvalues[i])+IM(floqueteigenvalues[i])*IM(floqueteigenvalues[i]); 
               if(fabs(shuldbeone-1.0)>=ABSERROR) 
                 {
                  errorcount=errorcount+1;
                  printf("Eigenvalue %d =  %lf +I %lf , quasienergy = %lf has non-unit modulus = %lf\n",i,RE(floqueteigenvalues[i]),IM(floqueteigenvalues[i]),quasi[i],shuldbeone);
                     }
               }
        printf("\n=========================================================================");
           fprintf(floqueteigenval,"\n%lf",tfix);
           for(i=0;i<DIM;i++)
              {
               fprintf(floqueteigenval,"\t%10.6lf",quasi[i]);
              }
       } 
     /*tfix loop ends here*/
     end=time(NULL);

     printf("\n=========================================================================");
     printf("\nTIME PERIOD OF DRIVE = %lf",period);
     printf("\n=========================================================================");
     printf("\n                 PROGRAM RAN SUCCESSFULLY (I THINK)                      ");
     printf("\n NUMBER OF TIMES THE IMAGINARY PART OF THE QUASIENERGY GOT TOO BIG = %ld ",errorcount);
     printf("\n                    TIME ELAPSED = %ld sec                               ",end-begin);
     printf("\n                   NUMBER OF THREADS = %ld                               ",nthreads);
     printf("\n=========================================================================\n");

     time ( &rawtime );
     timeinfo = localtime ( &rawtime );
     fclose(floqueteigenval);	
     return 0;
    }