#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <math.h>
#include <time.h>
#include "include.h" 
#include "interact.h"
#include "spectrum.h"
#include "extrapolate.h"

/*
  This file was formerly a part of "extralaso.c"
  The following performs all of the necessary extrapolations
  in K and particle truncation.
*/

/* convert kt into real number. */

#define K_FLOAT(K) (GLUE_PERIODIC?0.5*(element) K:(element) K)

#define PRINTMMA 0 /*Print eigenvalues in extrapolation routines in
                      a Mathematica compatable format */

#define PRINTEXTRAP 0 /* flag for extrapolation routines to print fit*/
                      /* coefficients, associated value for the last sector,
                         and sometimes goodness of fit */

/*
  The following subroutine is for extrapolating in 
  K and p-truncation                              
  
  For the case of GLUE_PERIODIC, we give up and use 1/K since the
  convergence is so poor.
*/

void kpweights(element *w, integer *th, int nsects, specify_calculation *s, 
	       element linc){
  int i;
  for(i=0; i<nsects; i++)
    w[i]=1.0;
  *th=1;
  if(nsects>1){
    for(i=0; i<nsects; i++)
      w[i+nsects]=pow(1.0/K_FLOAT(s[i].kt),
		      (IMPROVED>1 || GLUE_PERIODIC)?1.0:0.5);
    (*th)++;
  }
  if(nsects>2 && (s[0].type==TYPE_GLUE || s[0].type==TYPE_SOURCE)){
    for(i=0; i<nsects; i++)
      w[i+2*nsects]=exp(linc*(element) s[i].nptrunc);
    (*th)++;
  }
#if defined(KSQUARED) /* We can get fancier and fancier, ... */
  if(nsects>3){
    for(i=0; i<nsects; i++)
      w[i+3*nsects]=pow(1.0/K_FLOAT(s[i].kt),
			(IMPROVED>1 || GLUE_PERIODIC)?2.0:1.0);
    (*th)++;
  }
#endif
#if 0 /* We could get fancier and fancier, but we don't*/
  if(nsects>4){
    (*th)++;
    for(i=0; i<nsects; i++)
      w[i+4*nsects]=w[i+nsects]*w[i+2*nsects];
  }
#endif
}

/*  Given a list of truncations and coupling constants,
    find the extrapolated spectrum using laso to diagonalize
    the Hamiltonian.  It also returns the th "subleading terms"
    in the extrapolation.  
    Angle is used to set makephaselist when momentum is
    nonzero or (L,kmax) for heavy sources.  
    This routine can handle the case ns=1, just returning the eigenvalues.
    Although it can do nonzero L, this is best handled by lextrapolate(...).
*/ 

void extrapolate(element **tval, integer ns, specify_calculation  *s, 
		  element **vals, integer *th, int nval, element *pextrap){
  int i,t,j;
  doublereal *w;
#if PRINTMMA
  static int mma=0;
  int k=0;
#endif

  if(*pextrap > -0.25){
     fprintf(stderr,__FILE__":    Bad value of *pextrap in extrapolate\n"
	    "  setting *pextrap=-0.25\n");
    *pextrap=-0.25;
  }

  /* calculate weights used in extrapolation */
  w=malloc(ns*ns*sizeof(doublereal));
  kpweights(w,th,ns,s,*pextrap);
  *tval=realloc(*tval,*th*nval*sizeof(doublereal));
  for(i=0; i< *th*nval; i++)(*tval)[i]=0.0;
  PSEUDOINVERSE(w,th,&ns);

#if 1==0 /* debugging print */
  printf("  extrapolate; beta=%f, c_p=%f weights =\n",
	 beta,*pextrap);
  for(t=0; t<ns; t++){
    for(j=0; j< *th; j++)printf(" %f ",w[t+ns*j]);
    printf("\n");
  }
#endif
#if PRINTMMA /* print eigenvalues in Mathematica format */
  printf("    extra[%i] = {",mma++);
#endif

/*  Main loop -- go through different sectors */
  for(t=0; t<ns; t++){
#if 1==0 /* debugging print */
    printf("  Starting sector %i\n",t);
#endif
#if PRINTMMA /* print eigenvalues in Mathematica format */
    printf("{%i/2,%i,",s[t].kt,s[t].nptrunc);
    for(i=0; i<nval; i++){
      printf("%.10g",vals[t][i]); 
      if(i<nval-1)printf(","); else printf("}");
    }
    if(t<ns-1){
      printf(", "); 
      if(nval>2||(++k)%2==0)printf("\n      ");
    }
    else 
      printf("};\n");
#endif
    for(i=0; i<nval;i++)
      for(j=0; j< *th; j++)
	(*tval)[i+j*nval]+=vals[t][i]*w[t+ns*j];
  }
#if PRINTEXTRAP
  printf("    extrapolate %li sectors using %li parameter fit.\n",ns,*th);
  for(i=0; i<nval;i++){
    printf("      ");
    for(j=0; j<*th; j++)
      printf("%.10g ",(*tval)[j*nval+i]);	   
    printf("\n");
  }     
#endif

  free(w);
  return;
}

/* The following subroutine is for extrapolating in 
   K and p-truncation, and K_max.

   The mathematica results suggest that the Eigenvalues
   tend to converge mainly as (K_max/K)^2.  Thus, we use
   the leading (K_max/K) as well as the nonleading (K_max/K)^2.             
*/

void kmaxweights(element *w, integer *th, int nsects, 
		 specify_calculation *s, element *kmax, element linc){
  int i;
  for(i=0; i<nsects; i++)
    w[i]=1.0;
  *th=1;
  if(nsects>3){
    for(i=0; i<nsects; i++){
      w[i+nsects]=pow(kmax[i]/K_FLOAT(s[i].kt),
		      (IMPROVED>1 || GLUE_PERIODIC)?1:0.5);
      w[i+2*nsects]=exp(linc*(element) s[i].nptrunc);
      /* for the new tau interactions, leading convergence is */
      /* 1/kmax^5.   */
      w[i+3*nsects]=1.0/pow(kmax[i],5);
    }
    *th+=3;
  } 
  if(nsects>4){
    for(i=0; i<nsects; i++)
      w[i+4*nsects]=pow(kmax[i]/K_FLOAT(s[i].kt),
			(IMPROVED>1 || GLUE_PERIODIC)?2:1);
    (*th)++;
  }
}

/*  fitting functions
   for transverse and longitudinal string tension
   Since these can later be used for other values of n or L,
   the form of the fit function must remain fixed.  */

void wweights(element *w, integer *th, int nsects, int (*n)[HINDEX]){
  int i;
  for(i=0; i<nsects; i++){
    *th=3;
    w[i]=(element) dot(n[i],n[i],HINDEX);
    w[i+nsects]=1.0;
    w[i+2*nsects]=1.0/w[i];
  }
}

/*  I decided that the long distance expansion of
    the heavy potential should not contain a coulomb
    term --- since the Coulomb term should be identified
    with a short distance expansion.                     */

void lweights(element *w, integer *th, int nsects, element *ll){
  int i;
  for(i=0; i<nsects; i++){
    *th=3;
    w[i]=ll[i];
    w[i+nsects]=1.0;
    w[i+2*nsects]=1.0/w[i];
  }
}

/*
    Given a list of truncations and coupling constants,
    find the exponential convergence in p-truncation.  
*/

/* We use nonlinear fit for exponential factor only. */

#define MAXN 20  /* maximum number of sectors and fit parameters */
#define MAXPAR 1 /* maximum number of nonlinear fit parameters */
                  /* pass values to fcn2 */
static doublereal pval[MAXN], ptrunc[MAXN], pkt[MAXN];
 
element pextrapolate(integer ns, specify_calculation *s,element **vals){
  int i,t;
#if PRINTMMA
  int k=1;
#endif

  /* vectors must be static, else new memory is set aside by 
     spectrum each time pextrapolate is called. */

  integer npar=MAXPAR,maxfev=300,mode=1,ifail=0,info,nfev,ipvt[MAXPAR];
  doublereal xtol=1.e-10,gtol=0.0,ftol=1.e-10,epsfcn=1.e-10,factor=100.0;
  /* This is the initial guess */
  doublereal x[MAXPAR]={-1.0},fvec[MAXN],diag[MAXPAR],
	  fjac[MAXN*MAXPAR],qtf[MAXPAR],
          wa1[MAXPAR],wa2[MAXPAR],wa3[MAXPAR],wa4[MAXN];

  for(i=0, t=0; i<ns; i++){
    ptrunc[i]=(element) s[i].nptrunc;
    pkt[i]=K_FLOAT(s[i].kt);
  }
#if PRINTMMA /* print eigenvalues in Mathematica format */
  printf("    pextra = {");
#endif
  
  /*  go through different sectors */
  for(i=0; i<ns; i++)
    pval[i]=vals[i][0];
#if PRINTMMA /* print eigenvalues in Mathematica format */
  for(i=0; i<ns; i++){
    printf("{%g,%g,%.10g}",pkt[i],ptrunc[i],pval[i]);
    if(i<ns-1){
      printf(", "); 
      if((++k)%3==0)printf("\n      ");
    }
    else 
      printf("};\n");
  }
#endif
  
/*    Least squares fitting routine    
      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(FCN2_TO_FORTRAN,&ns,&npar,x,fvec,&ftol,&xtol,&gtol,
	&maxfev,&epsfcn,diag,&mode,&factor,&ifail,&info,&nfev,
	fjac,&ns,ipvt,qtf,wa1,wa2,wa3,wa4);
  for(i=0, wa1[0]=0.0; i<ns; i++)wa1[0] += fvec[i]*fvec[i];
  wa1[0]=sqrt(wa1[0]/(element) ns);
  if(info<1||info>4||wa1[0]>0.1){
     fprintf(stderr,__FILE__":    Error in pextrapolate lmdif info=%li\n",info);
     fprintf(stderr,__FILE__":    c=%e\n",x[0]);
    if(wa1[0]>0.1){
       fprintf(stderr,__FILE__":      Bad fit in pextrapolate, RMS=%e,"
	      " setting c=-1\n",*wa1);
      return -1;
    }
  }

#if PRINTEXTRAP
  printf("    pextrapolate %li sectors using %li parameter fit.\n    x=(",
	 ns,npar);
    for(i=0; i<npar; i++)printf(" %.10g",x[i]);	   
    printf(")\n    RMS deviation %g\n",*wa1);
#endif

  if(x[0]<-0.25)
    return x[0];
  else
    return -0.25;
}

/*  Fitting function used to determine exponential
    damping in p-truncation.  The function becomes
    fancier as more data points are given.  We use 
    linear least squares to get other fit parameters --
    otherwise, the fit tends to be unstable and it
    is hard to guess initial values.   */

#define MAXX 4  /* Maximum number of terms in linear fit. */

extern fsub fcn2(integer *m, integer *npar, doublereal *x, 
	       doublereal *fvecc, integer *iflag){
  int i,j;
  integer th;
  element w[MAXN*MAXX],ww[MAXN*MAXX],v[MAXX];
  if(*m>5)th=4; 
  else if(*m>2)th=2; 
  else{
     fprintf(stderr,__FILE__":  error in fcn2\n");
    *iflag=-1;
    return SUBRETURN;
  }
  for(i=0; i<*m; i++){
    w[i]=ww[i]=1.0;
    w[*m+i]=ww[*m+i]=exp(x[0]*ptrunc[i]);
    if(th==4){
      w[*m*2+i]=ww[*m*2+i]=1.0/pkt[i];
      w[*m*3+i]=ww[*m*3+i]=w[*m+i]*w[*m*2+i];
    }
  }
  PSEUDOINVERSE(w,&th,m);
  for(j=0; j<th; j++){
    v[j]=0.0;
    for(i=0; i<*m; i++)v[j]+=w[*m*j+i]*pval[i];
  }
  for(i=0; i<*m; i++){
    fvecc[i]=pval[i];
    for(j=0; j<th; j++)
      fvecc[i]-=v[j]*ww[*m*j+i];
  }
#if 0 /* debugging print */
  for(v[0]=0.0, i=0; i<*m; i++)v[0]+=fvecc[i]*fvecc[i];
  printf("    fcn2:  RMS=%e, c=%f\n",sqrt(v[0]/(element) *m),x[0]);
#endif
  return SUBRETURN;
}



/*  Given a list of truncations and coupling constants,
    find the extrapolated spectrum using laso to diagonalize
    the Hamiltonian.  This is only for the case of nonzero L.
    It also returns the th "subleading terms"
    in the extrapolation.  ll is the longitudinal separation L 
    in units of \bar g.  The vector kmaxlist contains the value
    of kmax in each sector and has length ns.
    This routine can handle the case ns=1, just returning the eigenvalues.
*/

void lextrapolate(element **tval, integer ns, specify_calculation *s, 
		    element **vals, integer *th, integer nval, 
		      element *pextrap, element *kmaxlist){
  int i,t,j;
  /* th is the number of fit terms in the extrapolation */
  doublereal *w;
#if PRINTMMA
  static int mma=0;
  int k=0;
#endif

  if(*pextrap > -0.25){
     fprintf(stderr,__FILE__":    Bad value of *pextrap in lextrapolate\n"
	    "  setting *pextrap=-0.25\n");
    *pextrap=-0.25;
  }

  w=malloc(ns*ns*sizeof(doublereal));
  kmaxweights(w,th,ns,s,kmaxlist,*pextrap);
#if 0 /* debugging print */
  printf("input matrix:  ");
  for(t=0; t<*th * ns; t++){
    if(t%(*th)==0)printf("\n    ");
    printf(" %g",w[t]);
 } 
#endif
  PSEUDOINVERSE(w,th,&ns);
  *tval=realloc(*tval,*th*nval*sizeof(doublereal));
  for(i=0; i< *th*nval;i++)(*tval)[i]=0.0;

#if 0 /*debugging print */
  printf("  lextrapolate exponential factor = %f weights =\n",
	 *pextrap);
  for(t=0; t<ns; t++){
    for(j=0; j< *th; j++)printf(" %f ",w[t+ns*j]);
    printf("\n");
  }
#endif
#if PRINTMMA /* print eigenvalues in Mathematica format */
  printf("    lextra[%i] = {",mma++);
#endif

/*  Main loop -- go through different sectors */
  for(t=0; t<ns; t++){
#if 1==0 /* debugging print */
    printf("  Starting sector %i\n",t);
#endif
#if PRINTMMA /* print eigenvalues in Mathematica format */
    printf("{%i/2,%i,%g,",s[t].kt,s[t].nptrunc,kmaxlist[t]);
    for(i=0; i<nval;i++){
      printf("%.10g",vals[t][i]); 
      if(i<nval-1)printf(","); else printf("}");
    }
    if(t<ns-1){
      printf(", "); 
      if(nval>2||(++k)%2==0)printf("\n      ");
    }
    else 
      printf("};\n");
#endif
    for(j=0; j< *th; j++)
      for(i=0; i<nval;i++)
	(*tval)[i+j*nval]+=vals[t][i]*w[t+ns*j];
  }
#if PRINTEXTRAP
  printf("    lextrapolate %li sectors using %li parameter fit.\n",
	 ns,*th);
  for(i=0; i<nval;i++){
    printf("      ");
    for(j=0; j< *th; j++)
      printf("%.10g ",(*tval)[i+j*nval]);	   
    printf("\n");
  }     
#endif

  free(w);

  return;
}