/*
 * This file is part of tela the Tensor Language.
 * Copyright (c) 1994-1996 Pekka Janhunen
 */

/*
	numerics.ct
	Numerical analysis functions.
	Preprocess with ctpp.
	C-tela code is C++ equipped with []=f() style function definition.
*/

extern Treal MachineEpsilon;	// from tela.C

#if 0
static int intpol1D(Treal& u, const Treal A[], int N, Treal x)
// 1D linear interpolation from a vector.
// Allowed range: 0<=x<N-1.
// Returns: 0 on success, nonzero on range overflow.
{
	int i = int((1-2*MachineEpsilon)*x);
	if (i<0 || i>=N-1) return 1;
	Treal t = x - i;
	u = (1-t)*A[i] + t*A[i+1];
	return 0;
}

static int intpol2D(Treal& u, const Treal A[], int Nx, int Ny, Treal x, Treal y)
// 2D linear interpolation from a matrix.
// Allowed range: 0<=x<Nx-1, 0<=y<Ny-1.
// Returns: 0 on success, nonzero on range overflow.
{
	int i = int((1-2*MachineEpsilon)*x);
	if (i<0 || i>=Nx-1) return 1;
	int j = int((1-2*MachineEpsilon)*y);
	if (j<0 || j>=Ny-1) return 1;
	Treal t = x - i;
	Treal s = y - j;
	u = (1-t)*(1-s)*A[i*Ny+j] + t*(1-s)*A[(i+1)*Ny+j] + (1-t)*s*A[i*Ny+j+1] + t*s*A[(i+1)*Ny+j+1];
	return 0;
}

static int intpolND(Treal& u, const Treal A[], const TDimPack N, const Treal x[], int r)
{
	Tint is[MAXRANK], mult[MAXRANK], isupper[MAXRANK];
	Treal ts[MAXRANK];
	for (int d=0; d<r; d++) {
		is[d] = int((1-2*MachineEpsilon)*x[d]);
		if (is[d] < 0 || is[d] >= N[d]-1) return 1;
		ts[d] = x[d] - is[d];
	}
	mult[r-1] = 1;
	for (d=r-1; d>=1; d--) mult[d-1] = N[d]*mult[d];
	const int Nterms = (1 << r);
	u = 0;
	int baseindex = 0, index;
	Treal coeff;
	for (d=0; d<r; d++) baseindex+= is[d]*mult[d];
	for (d=0; d<r; d++) isupper[d] = 0;
	for (int a=0; a<Nterms; a++) {
		for (d=0; d<r; d++)
			isupper[d] = (a & (1 << d));
		index = baseindex;
		coeff = 1;
		for (d=0; d<r; d++)
			if (isupper[d]) {
				index+= mult[d];
				coeff*= ts[d];
			} else
				coeff*= 1 - ts[d];
		u+= coeff*A[index];
	}
	return 0;
}
#endif

static int intpolND(Treal u[], int M, const Treal A[], const TDimPack N, const Treal *x[MAXRANK], int r)
// N-dimensional linear interpolation from array A to M-length array u
// Arguments:
//    u     - result array
//    M     - length of result array
//    A     - array from which to interpolate
//    N     - dimensions of A
//    x     - array of pointers to real index arrays
//    r     - rank of A
{
	int a,d;
	const int Nterms = (1 << r);	// 2^r terms affect each result
	const Treal almostunity = 1 - 2*MachineEpsilon;
	Tint ind, mult[MAXRANK], baseindex, isupper[MAXRANK][1 << MAXRANK], inds[1 << MAXRANK];
	Treal interpolant, ts[MAXRANK], coeffs[1 << MAXRANK];
	mult[r-1] = 1;
	for (d=r-1; d>=1; d--) mult[d-1] = N[d]*mult[d];
	// compute isupper table of 1's and 0's: isupper==1 if (index+1), 0 if (index)
	for (d=0; d<r; d++) {
		const int mask = (1 << d);
		for (a=0; a<Nterms; a++)
			isupper[d][a] = a & mask;
	}
	int i;
	for (i=0; i<M; i++) {	// Loop over result values
		baseindex = 0;			// index of lowest corner of interpolation hypercube
		for (d=0; d<r; d++) {
			ind = int(almostunity*x[d][i]);
			if (ind < 0 || ind >= N[d]-1) return 1;
			ts[d] = x[d][i] - ind;
			baseindex+= ind*mult[d];
		}
		// now ts, baseindex are ready
		// a numbers the 2^r terms in the sum that affects the result
		for (a=0; a<Nterms; a++)
			inds[a] = baseindex;
		// loop for d==0:
		for (a=0; a<Nterms; a++) {
			if (isupper[0][a]) {
				coeffs[a] = ts[0];
				inds[a]+= mult[0];
			} else
				coeffs[a] = 1 - ts[0];
		}
		// loop for d>0:
		for (d=1; d<r; d++)	{
			const Treal t = ts[d], oneminust = 1-t; const Tint m = mult[d];
			for (a=0; a<Nterms; a++) {
				if (isupper[d][a]) {
					coeffs[a]*= t;
					inds[a]+= m;
				} else
					coeffs[a]*= oneminust;
			}
		}
		interpolant = 0;		// gather result in this variable
		for (a=0; a<Nterms; a++)
			interpolant+= coeffs[a]*A[inds[a]];
		u[i] = interpolant;		// one result ready
	}
	global::nops+= r*(2*M + 3*M*Nterms/2) + 2*Nterms*M;
	return 0;
}

static int intpol(Tobject& y, const Tobject& A, const TConstObjectPtr indices[], int Nindices)
// For error codes, see following C-tela function.
// In this function, A is always real array.
{
	int errcode;
	Treal result;
	if (A.rank() != Nindices) return -2;
	int ScalarIndices = 1;
	TDimPack IndexDims;
    int p;
	for (p=0; p<Nindices; p++) {
		const Tkind ik = indices[p]->kind();
		if (ik!=Kint && ik!=Kreal && ik!=KIntArray && ik!=KRealArray) return -3;
		if (p == 0) {
			ScalarIndices = (ik==Kint || ik==Kreal);
			if (!ScalarIndices) IndexDims = indices[p]->dims();
		} else {
			if (ScalarIndices) {
				if (ik!=Kint && ik!=Kreal) return -4;
			} else {
				if (indices[p]->dims() != IndexDims) return -4;
			}
		}
	}
	if (ScalarIndices) {
		Treal inds[MAXRANK];
		for (p=0; p<Nindices; p++)
			inds[p] = ((indices[p]->kind()==Kint) ? Treal(indices[p]->IntValue()) : indices[p]->RealValue()) - ArrayBase;
		const Treal *indarray[MAXRANK];
		for (p=0; p<Nindices; p++) indarray[p] = &inds[p];
		errcode = intpolND(&result,1, A.RealPtr(), A.dims(), indarray, Nindices);
#if 0
		errcode = intpolND(result, A.RealPtr(), A.dims(), inds, Nindices);
#endif
		if (errcode) return -6;
		y = result;
#if 0
		switch (A.rank()) {
		case 1:
			errcode = intpol1D(result, A.RealPtr(), A.length(), inds[0]);
			if (errcode) return -6;
			y = result;
			break;
		case 2:
			errcode = intpol2D(result, A.RealPtr(), A.dims()[0], A.dims()[1], inds[0], inds[1]);
			if (errcode) return -6;
			y = result;
			break;
		default:
			return -5;
		}
#endif
	} else {	// Array indices
		Tobject inds[MAXRANK];
		const Tobject offset = Treal(-ArrayBase);
		for (p=0; p<Nindices; p++)
			Add(inds[p],*indices[p],offset);
		y.rreserv(indices[0]->dims());
		const Treal *indarray[MAXRANK];
		for (p=0; p<Nindices; p++) indarray[p] = inds[p].RealPtr();
		errcode = intpolND(y.RealPtr(),y.length(), A.RealPtr(), A.dims(), indarray, Nindices);
		if (errcode) return -6;
#if 0
		Treal indarray[MAXRANK];
		for (int i=0; i<y.length(); i++) {
			for (p=0; p<Nindices; p++) indarray[p] = inds[p].RealPtr()[i] - ArrayBase;
			errcode = intpolND(y.RealPtr()[i], A.RealPtr(), A.dims(), indarray, Nindices);
			if (errcode) return -6;
		}
#endif
#if 0
			switch (A.rank()) {
			case 1:
				errcode = intpol1D(y.RealPtr()[i], A.RealPtr(), A.length(), inds[0].RealPtr()[i]-ArrayBase);
				if (errcode) return -6;
				break;
			case 2:
				errcode = intpol2D(y.RealPtr()[i], A.RealPtr(), A.dims()[0], A.dims()[1],
								   inds[0].RealPtr()[i]-ArrayBase, inds[1].RealPtr()[i]-ArrayBase);
				if (errcode) return -6;
				break;
			default:
				return -5;
			}
#endif
	}
	return 0;
}

// Use C-tela functions 'Re'  and 'Im' from std.ct
extern "C" int Refunction(const TConstObjectPtr[], const int, const TObjectPtr[], const int);
extern "C" int Imfunction(const TConstObjectPtr[], const int, const TObjectPtr[], const int);

[y] = intpol(A...)
/* intpol(A,index1,index2...) is a general interpolation
   function. A must be an array from which values are interpolated.
   The rank of A must equal the number of index arguments.
   Each index argument may be a real scalar or real array.
   All index arguments must mutually agree in type and rank.
   The array A may also be complex. The result y is of same
   rank and size as each of the index arguments.

   intpol(A,i,j,...) is a generalization of mapped indexing
   A<[i,j,...]> for non-integral indices. The function benefits
   from vectorization even more than most other Tela functions.

   Currently intpol uses linear interpolation.
   Error codes:
   -1: First arg not a numerical array
   -2: Rank of first arg does not match number of index args
   -3: Non-real index arg
   -4: Dissimilar index args
   -6: Range overflow
*/
{
	// The actual job is performed by intpol above.
	// This function however treats the case of complex A.
	int errcode;
	if (A.kind()==KComplexArray) {
		Tobject ReA, ImA, Rey, Imy;
		TConstObjectPtr inputs[1]; TObjectPtr outputs[1];
		inputs[0] = &A;
		outputs[0] = &ReA;
		Refunction(inputs,1,outputs,1);
		inputs[0] = &A;
		outputs[0] = &ImA;
		Imfunction(inputs,1,outputs,1);
		errcode = intpol(Rey,ReA,argin+1,Nargin-1);
		if (errcode) return errcode;
		errcode = intpol(Imy,ImA,argin+1,Nargin-1);
		if (errcode) return errcode;
		const Tobject i = Tcomplex(0,1);
		Mul(Imy,i,Imy);
		Add(y,Rey,Imy);
	} else if (A.kind()==KRealArray) {
		errcode = intpol(y,A,argin+1,Nargin-1);
	} else if (A.kind()==KIntArray) {
		Tobject A1;
		const Tobject zero = 0.0;
		Add(A1,A,zero);		// now A1 is real array
		errcode = intpol(y,A1,argin+1,Nargin-1);
	} else
		return -1;
	return errcode;
}

[Lu] = stencil2d_4(u,ap0,am0,a0p,a0m)
/* stencil2d_4(u,ap0,am0,a0p,a0m) computes the two-dimensional
   five-point "molecule" where the coefficient of the central
   term is unity:

   Lu = u[i,j]
      + ap0*u[i+1,j] + am0*u[i-1,j]
	  + a0p*u[i,j+1] + a0m*u[i,j-1];

   where the indices i and j run from 2..nx and 2..ny where
   [nx,ny] = size(u). The size of ap0,am0,a0p,a0m must be two
   less than the size of u in both directions.
   Error codes:
   -1: One of the arguments is not a real matrix
   -2: One of the coefficient args has bad size
*/
{
    int p;
	for (p=0; p<5; p++) {
		if (argin[p]->kind()!=KRealArray) return -1;
		if (argin[p]->rank()!=2) return -1;
	}
	const Tint nx = u.dims()[0];
	const Tint ny = u.dims()[1];
	for (p=1; p<5; p++)
		if (argin[p]->dims()[0]!=nx-2 || argin[p]->dims()[1]!=ny-2) return -2;
	const Treal *U = u.RealPtr();
	const Treal *AP0 = ap0.RealPtr();
	const Treal *AM0 = am0.RealPtr();
	const Treal *A0P = a0p.RealPtr();
	const Treal *A0M = a0m.RealPtr();
	const Tint nx2=nx-2, ny2=ny-2;
	Lu.rreserv(TDimPack(nx2,ny2));
	Treal *LU = Lu.RealPtr();
	Tint i,j;
	VECTORIZED for (i=1; i<nx-1; i++) {
		VECTORIZED for (j=1; j<ny-1; j++) {
			LU[(i-1)*ny2+j-1] =
				U[i*ny+j]
				+ AP0[(i-1)*ny2+j-1]*U[(i+1)*ny+j] + AM0[(i-1)*ny2+j-1]*U[(i-1)*ny+j]
				+ A0P[(i-1)*ny2+j-1]*U[i*ny+j+1]   + A0M[(i-1)*ny2+j-1]*U[i*ny+j-1];
		}
	}
	global::nops+= 8*nx2*ny2;
	return 0;
}