/* INTERP_EXTRAP_WTS.C 
          interpolates/extrapolates a function whose values are
          known on a hypergrid. Interpolation is either
	  multidimensional linear interpolation or the the simplicial 
	  interpolation algorithm given in Algorithm 6.5 (p. 243) of Judd, 1998, 
	  Numerical methods in Economics (MIT Press).

          Procedure also computes the gradient vector and extrapolates 
	  extrapolates when the point is outside of the hypergrid in one 
  	  or more dimensions. Extrapolation is a fraction of a linear extrapolation 
	  based on the gradient of the function near the boundary of the 
	  hypergrid hypercube.

	  This function differs from interp_extrap.c by returning a vector p that
	  represents the interpolation/extrapolation, see below

	  John Rust, University of Maryland, June 2005 

    Inputs:

    s:       point (dx1 vector) the function is to be interpolated or extrapolated at
    vf:      function to be interpolated, an (nx1) vector storing the function at n
             grid points on a rectangular hyergrid, where n=n1*n2*...*nd, where nj is
	     the number of grids in dimension j, j=1,...,d
    grid:    gridpoints at which vf is given 
             grid=(x(1,1),...,x(1,n1),....,x(d,1),....,x(d,nd))
    gridnum: dx1 vector with number of grid points used in each dimension
             gridnum=(n1,n2,...,nd)

    intmeth: interpolation method, 1 for simplicial interpolation
				   0 for multilinear interpolation
    Output:

    V(s):      interpolated/extrapolated value of function at point s
    grad:      gradient of the interpolated function at point s.
    p:         a vector of dimension n x 1, same size as vf, that "represents" the
               interpolation operation, i.e. v(s)=p'vf. Note that when extrapolation
	       occurs (i.e. one or more components of s are outside the min or max
	       grid values), then v( s) is not necessarily equal to p'vf.
    
    We assume the elements of grid and vf are stored in column order, i.e.

        k=1;
        do i1=1,...,n1
	do i2=1,...,n2
	   ...
        do id=1,...,nd
	    k=k+1;
	    vf[k]=v(x(i1,1),x(i2,2),...,x(id,d)),

        where the grid points in each dimension satisfy:

	 x(j,1) < x(j,2) < .... < x(j,nj)

    Parameter: extfrac (extfrac is a parameter that controls the extrapolation
                        of the function when the point s is outside the grid.
			If extfrac=0 we are using a flat extrapolation, i.e.
			assuming the function value equals the interpolated i
			value of the function on the nearest boundary face of 
			the grid hypercube that is closest to s.

			If extrac=1, then a linear extrapolation is used, i.e.
			starting at the interpolated value of the function at
			the nearest grid hypercube cast to s, we project linearly
			using the gradient of the interpolated function at the
			nearest point to s in the hypercube.

			If extfrac is between 0 and 1 then a convex combination
			of the above two extrapolation strategies is used.

    Required procedures:

    binsucc.c    computes binary successor operation
    power.c      computes binary exponential function, 2^d
    vertwt.c     computes the probability weights (convex combination)
	         for the vertices of the grid hypercube so that the weighted
		 average of the vertices equals the point s to be interpolated
    kseq.c       locates the sequence number in a multidimensional index

    Note: this routine is for 2 and higher dimensional problems only.
          Use vint1d.c for 1-dimensional problems. In 1-dimensional 
	  problems there is no difference between linear and simplicial
          interpolation.
*/

double interp_extrap_wts(double *s,double *grad,double *p,double *vf,double *grid,
		       int *gridnum,int d,int intmeth)
{
        int i,j,ci,li,dd,n,k,lb,ub,lbs,twod,*indx,*gindx,*arg,*argl,*lbound,*ubound,*extrap;
	double adj,ivf,cvf,lvf,extfrac,*vsav,*wts,*swts,*ts,*runv,*grsav;

	extfrac=1.0;
	adj=0.0;

	ts=(double *) malloc(d*sizeof(double));
	runv=(double *) malloc(d*sizeof(double));
	wts=(double *) malloc(d*sizeof(double));
	lbound=(int *) malloc(d*sizeof(int));
	ubound=(int *) malloc(d*sizeof(int));
	arg=(int *) malloc(d*sizeof(int));
	extrap=(int *) malloc(d*sizeof(int));

	n=1;
	for (i=0; i<d; i++)
        {
          n=n*gridnum[i];
	  extrap[i]=0;
        }

	for (i=0; i<n; i++)
	{
           p[i]=0.0;
        }

	/* loop through each coordinate s and determine the upper and 
	   lower bounds of the grid slice containing it, or if is is
	   outside the grid, record whether s is the the left (extrap < 0)
	   of the minimum grid point, or to the right (extrap > 0) of the 
	   maximum grid point for the coordinate in question. */

        for (i=0; i<d; i++)
        {
	    ub=0;
	    for (j=0;j<=i;j++)
            {
	      ub=ub+gridnum[j];
	    } 

	    if (i > 0)
            {
              lb=1;
              for (j=0; j<=i-1; j++)
              {
	        lb=lb+gridnum[j];
              }
	    }
	    else
            {
	       lb=1;
	    }

	    lb--;
	    ub--;
	    lbs=lb;

	    if (s[i] > grid[ub])
            {
	       extrap[i]=i+1;
	       ts[i]=grid[ub];
            }
	    else if (s[i] < grid[lb])
            {
	       extrap[i]=-i-1;
	       ts[i]=grid[lb];
            }
	    else
            {
               ts[i]=s[i];
            }

	    for (j=lb; j<=ub; j++)
            {
               if (grid[j] >= s[i])
               {
		   ub=j;
               }
	    }

	    if (ub == lb)
            {
               ub=lb+1;
            }
	    else
            {
	       lb=ub-1;
            }

	    lbound[i]=lb-lbs;
	    ubound[i]=ub-lbs;

	    wts[i]=(ts[i]-grid[lb])/(grid[ub]-grid[lb]);
	    runv[i]=grid[ub]-grid[lb];
	}

/* The following branch uses either simplicial interpolation or
   multidimensional linear interpolation depending on the value of the
   intmeth argument */

   if (intmeth == 1)
   {
        /* Under simplicial interpolation, we start at the "top" of the
	   hypercube (i.e. the largest gridpoint in the vector sense)
	   and then move down the vertices of the hypercube for d more
	   times according to the sorted weight vector, wts, above, and
	   recursively taking convex combinations so that the final
	   interpolated value is a convex combination of d+1 of the 
	   hypercube vectices. The gradient, however, will be discontinuous
	   and constant within sub-simplices of the grid hypercube. */

        swts=(double *) malloc(d*sizeof(double));
	indx=(int *) malloc(d*sizeof(int));
	argl=(int *) malloc(d*sizeof(int));

	for (k=0; k<d; k++)
        {
	  arg[k]=ubound[k];
	  argl[k]=arg[k];
        }

	k=kseq(arg,gridnum,d);

	ivf=vf[k];
	p[k]=1.0;
	li=k;

	lvf=ivf;

	k=mysort(wts,swts,indx,d);

	for (k=0; k<d; k++)
        {
	    arg[indx[k]]=argl[indx[k]]-ubound[indx[k]]+lbound[indx[k]];
	    for (j=0; j<d; j++)
            {
	      argl[j]=arg[j];
            }
	    ci=kseq(arg,gridnum,d);
	    cvf=vf[ci];
	    grad[indx[k]]=(lvf-cvf)/runv[indx[k]];
	    ivf=ivf+(1-swts[k])*(cvf-lvf);
	    p[ci]=1-swts[k];
	    p[li]=p[li]-(1-swts[k]);
	    li=ci;
	    lvf=cvf;
	}

	lvf=0.0;

	free(argl);
	free(swts);
	free(indx);
   }
   else
   {
       /* Under multidimensional linear interpolation we start by
          retrieving the value of the function on each of the
          corners of the grid hypercube that contains the point
          that we want to interpolate at, and then recursively compute
          convex combinations to arrive at the interpolated value.
	  Multidimensional interpolation is thus subject to a curse of
	  dimensionality since it is a convex combination of all 2^d 
	  values of the hypercube vertices. However the gradient is
	  a continuous function of s. It is constant (i.e. independent
	  of s(1)) for the first coordinate, but the gradient with
	  respect to s(j) depends on the values of s(i) for i < j. */

        twod=power(d);
	i=twod-1;

        vsav=(double *) malloc(twod*sizeof(double));
        grsav=(double *) malloc(i*sizeof(double));
	indx=(int *) malloc(d*sizeof(int));
	gindx=(int *) malloc(d*sizeof(int));

	gindx[0]=0;
	for (i=1; i<=d-1; i++)
        {
           gindx[i]=gindx[i-1]+power(d-i);
        }

	for (k=0; k<d; k++)
        {
           indx[k]=0;
        }

        for (dd=d; dd>0; dd--)
        {
          if (dd == d)
          {
             for(i=0; i<twod; i++)
             {
               for (k=0; k<d; k++)
               {
	         arg[k]=lbound[k]*(1-indx[k])+indx[k]*ubound[k];
               }

	       ci=kseq(arg,gridnum,d);
  	       vsav[i]=vf[ci];
	       p[ci]=vertwt(wts,indx,d);

	       if (i % 2 == 1)
               {
                 k=(i+1)/2;
		 grsav[k-1]=(vsav[i]-vsav[i-1])/runv[d-1];
	         vsav[k-1]=(1-wts[d-1])*vsav[i-1]+wts[d-1]*vsav[i];
	       } 

	       if (i < twod-1)
               {
                 binsucc(indx,d);
               }
	   
             }
          }
          else
          {
	     twod=power(dd);

             for(i=0; i<twod; i++)
             {
                if (i % 2 == 1)
                {
                   k=(i+1)/2;

		   grsav[gindx[d-dd]+k-1]=(vsav[i]-vsav[i-1])/runv[dd-1];
	           vsav[k-1]=(1-wts[dd-1])*vsav[i-1]+wts[dd-1]*vsav[i];

		   for (j=dd+1; j<=d; j++)
                   {
                      grsav[gindx[d-j]+k-1]=(1-wts[dd-1])*grsav[gindx[d-j]+i-1]+
			                    wts[dd-1]*grsav[gindx[d-j]+i];
                   }
		}
	     }
	  }
	}
	for (j=d-1; j>=0; j--)
        {
          grad[j]=grsav[gindx[d-1-j]];
        }

	ivf=vsav[0];

        free(vsav);
        free(grsav);
	free(indx);
	free(gindx);
   }

/* The final step is to determine if any extrapolation is necessary, and if
   so, extrapolate and then return the interpolated/extrapolated function 
   value and the gradient vector */

        k=0;
	for (j=0; j<d; j++)
        {
           if (extrap[j] != 0)
           {
              k=1;
           }
        }

        if (k > 0)
	{ 
	   adj=0;

	   for (i=0; i<d; i++)
           {
	     if (extrap[i] < 0)
             {
   	       adj=adj+grad[-extrap[i]-1]*(s[-extrap[i]-1]-ts[-extrap[i]-1]);
	     } 
             if (extrap[i] > 0)
	     { 
   	       adj=adj+grad[extrap[i]-1]*(s[extrap[i]-1]-ts[extrap[i]-1]);
             }
	   } 
	   ivf=ivf+extfrac*adj;
	}  

	free(ts);
	free(runv);
	free(wts);
	free(lbound);
	free(ubound);
	free(arg);
	free(extrap);

	return ivf;
}