# include "psphot.h"

// write out the terms of the given 1D polynomial
void psPolynomial1DDump (psPolynomial1D *poly) {

    for (int i = 0; i < poly->n + 1; i++) {
	fprintf (stderr, "x^%d : %g +/- %g\n", i, poly->coeff[i], poly->coeffErr[i]);
    }
}    

psF32 Polynomial1DEval_EAM(psF32 x, const psPolynomial1D* myPoly)
{
    psS32 loop_x = 0;
    psF32 polySum = 0.0;
    psF32 xSum = 1.0;

    for (loop_x = 0; loop_x < myPoly->n + 1; loop_x++) {
        if (myPoly->mask[loop_x] == 0) {
            polySum += xSum * myPoly->coeff[loop_x];
        }
        xSum *= x;
    }

    return(polySum);
}

psVector *Polynomial1DEvalVector_EAM(const psPolynomial1D *myPoly,
				       const psVector *x)
{
    PS_POLY_CHECK_NULL(myPoly, NULL);
    PS_VECTOR_CHECK_NULL(x, NULL);
    PS_VECTOR_CHECK_TYPE(x, PS_TYPE_F64, NULL);

    psVector *tmp;

    tmp = psVectorAlloc(x->n, PS_TYPE_F64);
    for (psS32 i=0;i<x->n;i++) {
        tmp->data.F64[i] = Polynomial1DEval_EAM(x->data.F64[i], myPoly);
    }

    return(tmp);
}

// XXX EAM : use Nterm = Norder + 1 definition  
// XXX EAM : should we provide both order and nterms in struct?
psPolynomial1D* Polynomial1DAlloc(psS32 nOrder,
				  psPolynomialType type)
{
    PS_INT_CHECK_NON_NEGATIVE(nOrder, NULL);

    psS32 i = 0;
    psS32 nTerm = nOrder + 1;
    psPolynomial1D* newPoly = NULL;

    newPoly = (psPolynomial1D* ) psAlloc(sizeof(psPolynomial1D));
    // p_psMemSetDeallocator(newPoly, (psFreeFcn) polynomial1DFree);
    // XXX EAM : me, being lazy

    newPoly->type = type;
    newPoly->n = nOrder;
    newPoly->coeff = (psF32 *)psAlloc(nTerm * sizeof(psF32));
    newPoly->coeffErr = (psF32 *)psAlloc(nTerm * sizeof(psF32));
    newPoly->mask = (psU8 *)psAlloc(nTerm * sizeof(psU8));
    for (i = 0; i < nTerm; i++) {
        newPoly->coeff[i] = 0.0;
        newPoly->coeffErr[i] = 0.0;
        newPoly->mask[i] = 0;
    }
    return(newPoly);
}

// XXX EAM : my alternate BuildSums1D
static psVector *BuildSums1D(psVector* sums, 
			     psF64 x,
			     psS32 nTerm)
{
    psS32 nSum = 0;
    psF64 xSum = 0.0;

    nSum = 2*nTerm;
    if (sums == NULL) {
        sums = psVectorAlloc(nSum, PS_TYPE_F64);
    }
    if (nSum > sums->n) {
        sums = psVectorRealloc(sums, nSum);
    }

    xSum = 1.0;
    for (int i = 0; i < nSum; i++) {
        sums->data.F64[i] = xSum;
        xSum *= x;
    }
    return (sums);
}

// XXX EAM : test version of 1d fitting
psPolynomial1D* VectorFitPolynomial1DOrd_EAM(psPolynomial1D* myPoly,
					     psVector *mask,
					     const psVector *x,
					     const psVector *y,
					     const psVector *yErr)
{
    // I think this is 1 dimension down
    psImage*  	  A = NULL;
    psVector* 	  B = NULL;
    psVector* xSums = NULL;
    psS32 nTerm;
    psF64 wt;

    psTrace(".psLib.dataManip.VectorFitPolynomial1DOrd", 4,
            "---- VectorFitPolynomial1DOrd() begin ----\n");

    // dump minutiae
# ifndef PS_NO_TRACE
    if (psTraceGetLevel (".psLib.dataManip.VectorFitPolynomial1DOrd") >= 5) {
      FILE *f = psTraceGetDestination ();
      fprintf (f, "VectorFitPolynomial1D()\n");
      for (int i = 0; i < x->n; i++) {
	fprintf (f, "(x, y, yErr) is (%f, %f, %f)\n", x->data.F64[i], y->data.F64[i], yErr->data.F64[i]);
      }
    }
# endif
      
    nTerm = myPoly->n + 1;
    A     = psImageAlloc(nTerm, nTerm, PS_TYPE_F64);
    B     = psVectorAlloc(nTerm, PS_TYPE_F64);

    // Initialize data structures (why is this not a function!)
    for (int i = 0; i < nTerm; i++) {
        B->data.F64[i] = 0.0;
        for (int j = 0; j < nTerm; j++) {
            A->data.F64[i][j] = 0.0;
        }
    }

    // xSums look like: 1, x, x^2, ... x^(2n+1)

    // Build the B and A data structs.
    for (int k = 0; k < x->n; k++) {
        if ((mask != NULL) && mask->data.U8[k]) continue;
	xSums = BuildSums1D(xSums, x->data.F64[k], nTerm);
      
	if (yErr == NULL) {
	    wt = 1.0;
	} else {
	    // this should probably by yErr^2 !!
	    // this should filter yErr == 0 values
	    wt = 1.0 / PS_SQR(yErr->data.F64[k]);
	}
	for (int i = 0; i < nTerm; i++) {
	    B->data.F64[i] += y->data.F64[k] * xSums->data.F64[i] * wt;
	}

	// we could skip half of the array and assign at the end
	// we must handle masked orders
	for (int i = 0; i < nTerm; i++) {
	    for (int j = 0; j < nTerm; j++) {
		A->data.F64[i][j] += xSums->data.F64[i + j] * wt;
	    }
	}
    }
 
    // GaussJordan version
    if (0) {
	// does the solution in place
	psGaussJordan (A, B);
    
	// the first nTerm entries in B correspond directly to the desired 
	// polynomial coefficients.  this is only true for the 1D case
	for (int k = 0; k < nTerm; k++) {
	    myPoly->coeff[k] = B->data.F64[k];
	}
    } 
    else 
    // LUD version of the fit
    {
	psImage *ALUD = NULL;
	psVector* outPerm = NULL;
	psVector* coeffs = NULL;
    
	ALUD = psImageAlloc(nTerm, nTerm, PS_TYPE_F64);
	ALUD = psMatrixLUD(ALUD, &outPerm, A);
	coeffs = psMatrixLUSolve(coeffs, ALUD, B, outPerm);
	for (int k = 0; k < nTerm; k++) {
	    myPoly->coeff[k] = coeffs->data.F64[k];
	}
    }

    psFree(A);
    psFree(B);
    psFree(xSums);

    psTrace(".psLib.dataManip.VectorFitPolynomial1DOrd", 4,
            "---- VectorFitPolynomial1DOrd() begin ----\n");
    return (myPoly);
}

// ********************** 2D polynomial functions ******************

// XXX EAM : this version uses myPoly->nX as Norder, not Nterms
psVector *Polynomial2DEvalVector(const psPolynomial2D *myPoly,
				 const psVector *x,
				 const psVector *y)

{
    PS_POLY_CHECK_NULL(myPoly, NULL);
    PS_VECTOR_CHECK_NULL(x, NULL);
    PS_VECTOR_CHECK_TYPE(x, PS_TYPE_F32, NULL);
    PS_VECTOR_CHECK_NULL(y, NULL);
    PS_VECTOR_CHECK_TYPE(y, PS_TYPE_F32, NULL);

    psVector *tmp;
    psS32 vecLen=x->n;

    // Determine the length of the output vector to by the minimum of the x,y vectors
    if (y->n < vecLen) {
        vecLen = y->n;
    }

    // Create output vector to return
    tmp = psVectorAlloc(vecLen, PS_TYPE_F32);

    // Evaluate the polynomial at the specified points
    for (psS32 i=0; i<vecLen; i++) {
        tmp->data.F32[i] = Polynomial2DEval(myPoly,x->data.F32[i],y->data.F32[i]);
    }

    // Return output vector
    return(tmp);
}

// XXX EAM : this version uses the F64 vectors
psVector *Polynomial2DEvalVectorD(const psPolynomial2D *myPoly,
				  const psVector *x,
				  const psVector *y)

{
    PS_POLY_CHECK_NULL(myPoly, NULL);
    PS_VECTOR_CHECK_NULL(x, NULL);
    PS_VECTOR_CHECK_TYPE(x, PS_TYPE_F64, NULL);
    PS_VECTOR_CHECK_NULL(y, NULL);
    PS_VECTOR_CHECK_TYPE(y, PS_TYPE_F64, NULL);

    psVector *tmp;
    psS32 vecLen=x->n;

    // Determine the length of the output vector to by the minimum of the x,y vectors
    if (y->n < vecLen) {
        vecLen = y->n;
    }

    // Create output vector to return
    tmp = psVectorAlloc(vecLen, PS_TYPE_F64);

    // Evaluate the polynomial at the specified points
    for (psS32 i=0; i<vecLen; i++) {
        tmp->data.F64[i] = Polynomial2DEval(myPoly,x->data.F64[i],y->data.F64[i]);
    }

    // Return output vector
    return(tmp);
}

// XXX EAM : use Nterm = Norder + 1 definition  
// the user requests a polynomial of order Norder
psPolynomial2D* Polynomial2DAlloc(psS32 nXorder, psS32 nYorder,
                                    psPolynomialType type)
{
    PS_INT_CHECK_NON_NEGATIVE(nXorder, NULL);
    PS_INT_CHECK_NON_NEGATIVE(nYorder, NULL);

    psS32 x = 0;
    psS32 y = 0;
    psS32 nXterm = nXorder + 1;
    psS32 nYterm = nYorder + 1;
    psPolynomial2D* newPoly = NULL;

    newPoly = (psPolynomial2D* ) psAlloc(sizeof(psPolynomial2D));
    // p_psMemSetDeallocator(newPoly, (psFreeFcn) polynomial2DFree);
    // XXX EAM : me, being lazy

    newPoly->type = type;
    newPoly->nX = nXorder;
    newPoly->nY = nYorder;

    newPoly->coeff = (psF32 **)psAlloc(nXterm * sizeof(psF32 *));
    newPoly->coeffErr = (psF32 **)psAlloc(nXterm * sizeof(psF32 *));
    newPoly->mask = (psU8 **)psAlloc(nXterm * sizeof(psU8 *));
    for (x = 0; x < nXterm; x++) {
        newPoly->coeff[x] = (psF32 *)psAlloc(nYterm * sizeof(psF32));
        newPoly->coeffErr[x] = (psF32 *)psAlloc(nYterm * sizeof(psF32));
        newPoly->mask[x] = (psU8 *)psAlloc(nYterm * sizeof(psU8));
    }
    for (x = 0; x < nXterm; x++) {
        for (y = 0; y < nYterm; y++) {
            newPoly->coeff[x][y] = 0.0;
            newPoly->coeffErr[x][y] = 0.0;
            newPoly->mask[x][y] = 0;
        }
    }
    return(newPoly);
}

// XXX EAM : BuildSums2D in analogy with BuildSums1D
static psImage *BuildSums2D(psImage* sums,
			    psF64 x,      psF64 y,
			    psS32 nXterm, psS32 nYterm)
{
    psS32 nXsum = 0;
    psS32 nYsum = 0;
    psF64 xSum = 1.0;
    psF64 ySum = 1.0;

    nXsum = 2*nXterm;
    nYsum = 2*nYterm;
    if (sums == NULL) {
        sums = psImageAlloc(nXsum, nYsum, PS_TYPE_F64);
    }
    if ((nXsum != sums->numCols) || (nYsum != sums->numRows)) {
	psFree (sums);
        sums = psImageAlloc(nXsum, nYsum, PS_TYPE_F64);
    }

    ySum = 1.0;
    for (int j = 0; j < nYsum; j++) {
	xSum = ySum;
	for (int i = 0; i < nXsum; i++) {
	    sums->data.F64[i][j] = xSum;
	    xSum *= x;
	}
	ySum *= y;
    }
    return (sums);
}

// XXX EAM : test version of 2d fitting
psPolynomial2D* VectorFitPolynomial2DOrd_EAM(psPolynomial2D* myPoly,
					     psVector* mask,
					     const psVector* x,
					     const psVector* y,
					     const psVector* z,
					     const psVector* zErr)
{
    // I think this is 1 dimension down
    psImage*  	  A = NULL;
    psVector* 	  B = NULL;
    psImage*   Sums = NULL;
    psF64 wt;
    psS32 nTerm, nXterm, nYterm;

    nXterm = myPoly->nX + 1;
    nYterm = myPoly->nY + 1;
    nTerm = nXterm * nYterm;

    A = psImageAlloc(nTerm, nTerm, PS_TYPE_F64);
    B = psVectorAlloc(nTerm, PS_TYPE_F64);

    // Initialize data structures (why is this not a function!)
    for (int i = 0; i < nTerm; i++) {
        B->data.F64[i] = 0.0;
        for (int j = 0; j < nTerm; j++) {
            A->data.F64[i][j] = 0.0;
        }
    }

    // Sums look like: 1, x, x^2, ... x^(2n+1), y, xy, x^2y, ... x^(2n+1)

    // Build the B and A data structs.
    for (int k  = 0; k < x->n; k++) {
      if ((mask != NULL) && mask->data.U8[k]) continue;
	Sums = BuildSums2D(Sums, x->data.F64[k], y->data.F64[k], nXterm, nYterm);
      
	if (zErr == NULL) {
	    wt = 1.0;
	} else {
	    // this should probably by zErr^2 !!
	    // this should filter zErr == 0 values
	    wt = 1.0 / zErr->data.F64[k];
	}

	// we could skip half of the array and assign at the end
	// we must handle masked orders
	for (int n = 0; n < nXterm; n++) {
	    for (int m = 0; m < nYterm; m++) {
		B->data.F64[n+m*nXterm] += z->data.F64[k] * Sums->data.F64[n][m] * wt;
	    }
	}

	for (int i = 0; i < nXterm; i++) {
	    for (int j = 0; j < nYterm; j++) {
		for (int n = 0; n < nXterm; n++) {
		    for (int m = 0; m < nYterm; m++) {
			A->data.F64[i+j*nXterm][n+m*nXterm] += Sums->data.F64[i+n][j+m] * wt;
		    }
		}
	    }
	}
    }
 
    // does the solution in place
    psGaussJordan (A, B);
    
    // XXX: How do we know if these routines were successful?
    // ALUD = psMatrixLUD(ALUD, &outPerm, A);
    // coeffs = psMatrixLUSolve(coeffs, ALUD, B, outPerm);

    for (int n = 0; n < nXterm; n++) {
	for (int m = 0; m < nYterm; m++) {
	    myPoly->coeff[n][m] = B->data.F64[n+m*nXterm];
	}
    }

    psFree(A);
    psFree(B);
    psFree(Sums);

    psTrace(".psLib.dataManip.VectorFitPolynomial2DOrd", 4,
            "---- VectorFitPolynomial2DOrd() begin ----\n");
    return (myPoly);
}

// write out the terms of the given 2D polynomial
void psPolynomial2DDump (psPolynomial2D *poly) {

    for (int i = 0; i < poly->nX + 1; i++) {
	for (int j = 0; j < poly->nY + 1; j++) {
	    fprintf (stderr, "x^%d y^%d : %g +/- %g\n", i, j, poly->coeff[i][j], poly->coeffErr[i][j]);
	}
    }
}    

psPolynomial2D* RobustFit2D_nomask(psPolynomial2D* poly,
			    const psVector* x,
			    const psVector* y,
			    const psVector* z,
			    const psVector* dz)
{
    psVector *X;
    psVector *Y;
    psVector *Z;
    psVector *dZ;

    psVector *zFit   = NULL;
    psVector *zResid = NULL;
    psStats  *stats  = NULL;

    X  = psVectorCopy (NULL, x, PS_TYPE_F64);
    Y  = psVectorCopy (NULL, y, PS_TYPE_F64);
    Z  = psVectorCopy (NULL, z, PS_TYPE_F64);
    dZ = psVectorCopy (NULL, dz, PS_TYPE_F64);

    for (int N = 0; N < 3; N++) {
	// XXX EAM : this would be better defined with an element mask
	poly   = VectorFitPolynomial2DOrd_EAM (poly, NULL, X, Y, Z, dZ);
	zFit   = Polynomial2DEvalVectorD (poly, x, y);
	zResid = (psVector *) psBinaryOp (NULL, (void *) z, "-", (void *) zFit);

	stats = psStatsAlloc (PS_STAT_CLIPPED_MEAN | PS_STAT_CLIPPED_STDEV);
	stats  = psVectorStats (stats, zResid, NULL, NULL, 0);
	psTrace (".psphot.RobustFit", 4, "residual stats for robust fit:  %g +/- %g (%d pts)\n", stats->clippedMean, stats->clippedStdev, stats->clippedNvalues);

	// re-create X, Y, Z, dZ if pts are valid
	int n = 0;
	for (int i = 0; i < zResid->n; i++) {
	    if (fabs(zResid->data.F64[i] - stats->clippedMean) > 3*stats->clippedStdev) {
	      continue;
	    } 
	    X->data.F64[n]  =  x->data.F64[i];
	    Y->data.F64[n]  =  y->data.F64[i];
	    Z->data.F64[n]  =  z->data.F64[i];
	    dZ->data.F64[n] = dz->data.F64[i];
	    n++;
	}
	X->n = n;
	Y->n = n;
	Z->n = n;
	dZ->n = n;
    }
    return (poly);
}

// XXX EAM : be careful here with F32 vs F64 vectors 
psPolynomial2D* RobustFit2D(psPolynomial2D* poly,
			    psVector* mask,
			    const psVector* x,
			    const psVector* y,
			    const psVector* z,
			    const psVector* dz)
{
    PS_VECTOR_CHECK_NULL(mask, NULL);
    PS_VECTOR_CHECK_NULL(x, NULL);
    PS_VECTOR_CHECK_NULL(y, NULL);
    PS_VECTOR_CHECK_NULL(z, NULL);
    PS_VECTOR_CHECK_NULL(dz, NULL);

    psVector *zFit   = NULL;
    psVector *zResid = psVectorAlloc (x->n, PS_TYPE_F64);
    psStats  *stats  = psStatsAlloc (PS_STAT_SAMPLE_MEAN | PS_STAT_SAMPLE_STDEV);

    for (int N = 0; N < 3; N++) {
	poly   = VectorFitPolynomial2DOrd_EAM (poly, mask, x, y, z, dz);
	zFit   = Polynomial2DEvalVectorD (poly, x, y);
	zResid = (psVector *) psBinaryOp (zResid, (void *) z, "-", (void *) zFit);

	stats  = psVectorStats (stats, zResid, NULL, mask, 1);
	psTrace (".psphot.RobustFit", 4, "residual stats for robust fit:  %g +/- %g\n", 
		 stats->sampleMean, stats->sampleStdev);

	// set mask if pts are not valid
	// we are masking out any point which is out of range 
	// recovery is not allowed with this scheme
	for (int i = 0; i < zResid->n; i++) {
	  if (mask->data.U8[i]) continue;
	  if (fabs(zResid->data.F64[i] - stats->sampleMean) > 3*stats->sampleStdev) {
	    mask->data.U8[i] = 1;
	    continue;
	  }       
	}
	psFree (zFit);
    }
    psFree (zResid);
    psFree (stats);
    return (poly);
}

// XXX EAM : VectorFitPolynomial2DOrd and Polynomial2DEvalVector require different types (F32 vs F64)

# if (0) // moved to psLib
// XXX EAM : this version uses myPoly->nX as Norder, not Nterms
psF32 Polynomial2DEval(const psPolynomial2D* myPoly,
		       psF32 x,
		       psF32 y)
{
    PS_POLY_CHECK_NULL(myPoly, NAN);

    psS32 loop_x = 0;
    psS32 loop_y = 0;
    psF32 polySum = 0.0;
    psF32 xSum = 1.0;
    psF32 ySum = 1.0;

    // XXX EAM : nX is order, not nTerms
    for (loop_x = 0; loop_x < myPoly->nX + 1; loop_x++) {
        ySum = xSum;
	// XXX EAM : nX is order, not nTerms
        for (loop_y = 0; loop_y < myPoly->nY + 1; loop_y++) {
            if (myPoly->mask[loop_x][loop_y] == 0) {
                polySum += ySum * myPoly->coeff[loop_x][loop_y];
            }
            ySum *= y;
        }
        xSum *= x;
    }

    return(polySum);
}
# endif