#include <stdio.h>
#include <string.h>
#include <time.h>
#include <math.h>
#include "include.h"
#include "run.h"
doublereal dot2(const doublereal *,const doublereal *,const doublereal *,int);

/*  Print timing and convergence information  */
#define PRINTIT 1==1

/** This file runs the lanczos algorithm for finding the lowest
    Eigenvalues of the Hamiltonian with extrapolation for nonzero P_1.
    The results are then analyzed for consistant dispersion relations,
    and a best fit is found based on this criterion.

    This is a minor modification of disperse.c where 
    we look at two nonzero values of momentum.
    
**/

#define NJPC 2 /* Number of JPC sectors */
#define TOTALSTATES 8 /*Total number of states in all NJPC sectors */
#define NP1 3  /* Number of different values of transverse momentum */
#define MM TOTALSTATES*(NP1-1) /* Number of terms in least squares fit */
#if 1==0  /* Truncations used in paper */
#define NS 7
  const static int kt[NS]={20,20,20,22,24,26,28}, np[NS]={4,6,8,6,6,6,6};
#elif 1==1
#define NS 5
  const static int kt[NS]={14,10,10,12,14}, np[NS]={4,6,8,6,6};
#else
#define NS 5
  const static int kt[NS]={20,20,22,22,32}, np[NS]={6,8,6,8,6};
#endif
const int nvals[NJPC]={4,4};
static struct sectors s[NJPC][NS],t[NJPC][NS];
#if 1==0
  element par[NPARAMS]={0.1,1.0,-0.2,-0.23,10.0}, rescale=1.0;
#elif 1==1 /* Best fit couplings from paper  */
  element par[NPARAMS]={0.0649839, 1.,-0.201606821,-0.550384068,116.657},
	  rescale=3.46911443596316;
#elif 1==0 /*  This is the new best fit using shuffle */
  element par[NPARAMS]={0.064984,1.000000,-0.147802,-0.611392,116.657000},
	  rescale=3.410446;
#else /*  This is the new best fit using larger basis. */
  element par[NPARAMS]={0.064984, 1.000000, -0.152583, -0.582787, 116.657000},
	     rescale=3.552178;
#endif
static int wpar[NPARAMS];
static doublereal longvals[NP1][TOTALSTATES],
  p1a[NP1][HINDEX]={{0.0},{0.5},{1.0}};

static const doublereal err2[NP1-1][TOTALSTATES]=
             {{1.0,10000.0,10000.0,2.0, 1.0,1.0,10000.0,10000.0},
              {1.0,10000.0,10000.0,2.0, 1.0,1.0,10000.0,10000.0}};
static const doublereal teper[TOTALSTATES]={
       16.63019968814394, 39.77813397880572, 45.7690369882667,
          45.7690369882667,
       35.73301963006418, 55.454956994115, 61.83048213326445, 
          61.83048213326445},
	 specterr2[TOTALSTATES]={
        1.020792692835208, 7.16636029411156, 7.995737368805813,
	     7.995737368805813, 
       16.58584836076752, 65.06810610315358, 67.44195640455192,
	    67.44195640455192};
   

int MAIN(){
  char out[]="baba.out";
  integer mm=MM,npar=1,maxfev,mode=1,
    ifail=0,info,nfev,ldfjac=MM;
  integer ipvt[NPARAMS];
  /*  fcntol should be a little less than the Lanczos accuracy
      tol is the desired precision of the coupling constants  */
  doublereal tol,fcntol=1.0e-7,factor=100.0;
  doublereal vpar[NPARAMS],fvec[MM],diag[NPARAMS],
    fjac[MM*NPARAMS],qtf[NPARAMS],
    wa1[NPARAMS],wa2[NPARAMS],wa3[NPARAMS],wa4[MM];
  
  clock_t t1;
  int i,j,o[NJPC]={1,1},o2[NJPC]={-1,-1},ht[HINDEX]={0},
	     multi[NJPC]={1,2},multi2[NJPC]={2,1};
  int new=(1==1);
  element chi2;
  FILE *fp;

#ifdef BABBAGE  /* Routine to initialize Fortran on Babbage */
  if(hf_fint((char*)NULL)) exit(1);
#endif

  loadelements();
  if((fp = fopen(out,"w")) == NULL){
     fprintf(stderr,__FILE__":  Can't open file %s in monster, exiting\n",out);
    exit(55);
  }
  for(i=0, j=0; i<NJPC; i++)j+=nvals[i];
  if(j!=TOTALSTATES){
     fprintf(stderr,__FILE__":  Number of states don't match in Jessica\n");
    exit(5);
  }
  fprintf(fp,"*** Eigenvalues for the 2+1 transverse lattice  ***\n");
  fprintf(fp,"*** Best fit to Teper's lattice spectrum        ***\n");
  fprintf(fp,"%i starting couplings: (m^2, G^2 N, lambda_i/g^2 N,...):\n  ",
	  NIN);
  for(i=0;i<NIN;i++)fprintf(fp," %f",par[i]);
  fprintf(fp,"\nExtrapolation using (K,p)=");
  for(i=0;i<NS;i++)fprintf(fp,"(%i/2,%i) ",kt[i],np[i]);
  fprintf(fp,"\n");
  fprintf(fp,"Multiplets (o,multiplet, # states):  ");
  for(i=0;i<NJPC;i++)fprintf(fp,"(%i,%i,%i) ",o[i],multi[i],nvals[i]);
  fprintf(fp,"\nThe following is chi^2, # steps, rescaling factor\n");
  fprintf(fp,"The %i couplings (without rescaling) and which were fit\n\n",
	  NIN);

  printf("Constructing %i bases",2*NJPC*NS);
  t1=clock();
  for(j=0; j<NJPC; j++)
    for(i=0; i<NS; i++){
      initbasis(s[j]+i,o[j],ht,kt[i],np[i],multi[j],new);
      initbasis(t[j]+i,o2[j],ht,kt[i],np[i],multi2[j],new);
    }
  new=(1==0);
  printf(" in %f CPU seconds\n",
       (float)(clock()-t1)/(float) CLOCKS_PER_SEC);

  for(;;){
    printf("How many to fit (0 to %i, -1=stop)\n",NIN-2);
    scanf("%li",&npar);
    if(npar<0)
      break;
    else if(npar==0)
      fcn(&mm,&npar,vpar,fvec,&info);
    else{
      printf("Input which %li couplings (0=m^2, ...)\n",npar);
      for(i=0;i<npar;i++)scanf("%i",wpar+i);
      for(i=0;i<npar;i++)vpar[i]=par[wpar[i]];
      maxfev=200*(npar+1);

/*      subroutine lmdif(fcn,m,n,x,fvec,ftol,xtol,gtol,maxfev,epsfcn,
  *                 diag,mode,factor,nprint,info,nfev,fjac,ldfjac,
  *                 ipvt,qtf,wa1,wa2,wa3,wa4)  */
 
    
      LMDIF(FCN_TO_FORTRAN,&mm,&npar,vpar,fvec,&tol,&tol,
	    &fcntol,&maxfev,&fcntol,
	    diag,&mode,&factor,&ifail,&info,&nfev,fjac,&ldfjac,
	    ipvt,qtf,wa1,wa2,wa3,wa4);
      if(info<1||info>4) fprintf(stderr,__FILE__":  Error in lmdif info=%li\n",info);
      for(i=0; i<npar; i++)par[wpar[i]]=vpar[i];
    }
    chi2=0.0;
    for(i=0; i<mm; i++)chi2+=fvec[i]*fvec[i];
    printf("chi2=%e in %li steps, rescale=%e\n",chi2,npar?nfev:1,rescale);
    fprintf(fp,"%e %li %e\n",chi2,nfev,rescale);
    printf("Final couplings: ");
    for(i=0; i<NIN; i++){
      printf(" %f",par[i]); fprintf(fp,"%f ",par[i]);
    }
    if(npar>0)printf("\nwhile fitting");
    for(i=0; i<npar; i++){
      printf(" %i",wpar[i]); fprintf(fp," %i",wpar[i]);
    }
    printf("\nMomenta: "); fprintf(fp,"\n");
    for(j=0; j<NP1; j++){
      printf(" %f",p1a[j][0]); fprintf(fp," %f",p1a[j][0]);
    }
    printf("\n"); fprintf(fp,"\n");
    for(i=0; i<TOTALSTATES; i++){
      for(j=0; j<NP1; j++){
	printf(" %f",longvals[j][i]); fprintf(fp," %f",longvals[j][i]);
      }
      printf("\n"); fprintf(fp,"\n");
    }
    fprintf(fp,"\n"); printf("\n");
  }
  fclose(fp);
  return 0;
}

extern int fcn(integer *m, integer *npar, doublereal *xp, 
	       doublereal *fvecc, integer *iflag){
  int i,k,j,l;
  int ipt[HINDEX]={0};
  clock_t tf;
  element *vals,chi2=0.0,ideal;
  tf=clock();
  for(i=0;i<*npar;i++)par[wpar[i]]=xp[i];
#if PRINTIT
  printf("  Parameters: ");
  for(i=0; i<NIN; i++)printf(" %f",par[i]);
#endif
  for(l=0; l<NP1; l++)
    for(i=0, j=0; i<NJPC; i++){
      vals=extrapolate(NS,s[i],t[i],nvals[i],NIN,par,ipt,p1a[l]);
      for(k=0; k<nvals[i]; k++, j++)longvals[l][j]=vals[k];  
    }

  /*  Calculate the desired dispersion with overall scale
      determined by a fit to Teper's spectrum.      */
  rescale=dot2(teper,longvals[0],specterr2,TOTALSTATES)/
    dot2(longvals[0],longvals[0],specterr2,TOTALSTATES);
  rescale=fabs(rescale);
  for(l=0; l<NP1; l++)
    for(i=0; i<TOTALSTATES; i++)longvals[l][i]*=rescale;
  for(l=1, j=0; j<*m; l++){
    ideal=pow(0.1974*rescale*p1a[l][0],2)-
          pow(0.1974*rescale*p1a[0][0],2);
    for(i=0; i<TOTALSTATES; i++, j++){
      fvecc[j] = longvals[l][i]-longvals[0][i]-ideal; 
      fvecc[j] /= sqrt(err2[l-1][i]);
      chi2 += fvecc[j]*fvecc[j];
    }
   }
#if PRINTIT
  for(l=1, j=0; l<NP1; l++){
    printf("\n ");
    for(i=0; i<TOTALSTATES; i++, j++)printf(" %f",fvecc[j]);
  }
  printf("\n  Chi^2=%f rescale=%f;",chi2,rescale);
  printf(" %i extrapolations in %f CPU seconds\n\n",
	 NJPC*NP1,(float)(clock()-tf)/(float) CLOCKS_PER_SEC);
#endif
  return 0;
}


doublereal dot2(const doublereal *x,const doublereal *y,
		const doublereal *z,int n){
  int i;
  doublereal t=0;
  for(i=0; i<n; i++)t+=x[i]*y[i]/z[i];
  return t;
}