/* SOComputeApprox - Spline approximation of an arbitrary function over a dyadic grid. */
/* Last edited on 2003-07-08 12:49:26 by stolfi */

/*
  This program computes a least squares approximation {g}, for a 
    function {f}, of the form
  
    g = SUM { u[i]*bas[i] : i = 0..dim-1 }       (1)
    
  where {bas[0..dim-1]} is an arbitrary (but add-compatible) basis.
  The coefficients {u[i]} are found by solving the linear system
    
    G u == b    (2)
    
  where {G} is the rigidity matrix, {G[i,j] == <bas[i] | bas[j]>},
  and the vector {b} is given by {b[i] == <f | bas[i]>}.
*/
    
#include <SOGrid.h>
#include <SOMatrix.h>
#include <SOFunction.h>
#include <SOLinCombFunction.h>
#include <SOApprox.h>
#include <SOParams.h>
#include <SOPlotParams.h>
#include <SO1DPlot.h>
#include <SOBasic.h>
#include <SOIntegral.h>

#include <js.h>
#include <vec.h>

#include <stdio.h>
#include <math.h>
#include <stdlib.h>
#include <string.h>
#include <values.h>
#include <limits.h>

#define MAXSTR 50
#define MAX_PLOT_STEPS 10

typedef struct Options
  { char *funcName;       /* Filename of function to approximate (minus ".fun"). */
    char *matName;        /* Prefix of matrix file names (default {basisName}). */
    char *outName;        /* Prefix for output file names. */
    int minRank;          /* Minimum rank for grid generation. */  
    int maxRank;          /* Maximum rank for grid generation. */  
    double diff_limit;    /* Parameter used for grid refinement. */
    bool maximal;         /* TRUE creates a maximal basis. Defaults is FALSE (minimal basis). */       
    nat gaussOrder;       /* Gaussian quadrature order for dot products. */
    /* Method used to solve the linear system: */
    bool cholesky;        /* TRUE uses direct method based on Cholesky factzn. */
    bool gaussSeidel;     /* TRUE uses Gauss-Seidel iteration. */
    bool conjGrad;        /* TRUE uses the conjugate gradient method. */
    /* Parameters for Gauss-Seidel iteration: */
    double omega;         /* Relaxation factor. */
    nat maxIter;          /* Maximum iterations. */
    double relTol;        /* Stopping criterion: relative change. */
    double absTol;        /* Stopping criterion: absolute change. */
    /* Plotting options: */
    bool plot;            /* TRUE to generate function and error plots. */
    PlotOptions plt;      /* Plot format and style options. */
  } Options;

Options *GetOptions(int argn, char **argc);

void ShatterNode
  ( SOGrid_Dim pDim,
    SOGrid_Node *p,
    SOGrid_Rank r,
    double *lo, 
    double *hi,
    double *maxhi,
    int minRank,
    int maxRank,
    double diff_limit,
    SOFunction *f,
    int *leaf_cells
  );

void SetGenDotMatrices
  ( MatEntry_vec *evaldot, 
    SOMatrix Meval, 
    unsigned int offset,
    unsigned int maxindices
  );

void SetDotMatrices
  ( unsigned int dim, 
    unsigned int offset,
    Basis bas, 
    MatEntry_vec *evaldot, 
    MatEntry_vec *Hev_ents
  );


int main(int argn, char **argc)
  { SOFunction *w = (SOFunction *)NULL; 
    Options *o = GetOptions(argn, argc);
    int i, leaf_cells=601;

    SOFunction *f = SOApprox_ReadFunction(txtcat(o->funcName, ".fun")); 
/*     SOFunction *refine_f = SOApprox_ReadFunction("fun/temp.fun"); */
    int pDim = f->d->pDim;
    int fDim = f->d->fDim;


    /* Builds the {SOGrid} tree. */
    SOGrid_Tree *tree = SOGrid_Tree_new(pDim);

    double lo[pDim], hi[pDim], maxhi[pDim];
    for(i = 0; i < pDim; i++){ lo[i] = 0.0; hi[i] = 1.0/pow(2.0, (double)i/pDim); maxhi[i] =  hi[i]; }

    /* Finds the lower rank grid to build the initial functions. */
    while(leaf_cells > 600)
      {
        o->diff_limit = o->diff_limit + 0.000001;
        if( tree->root->c[LO] != NULL && tree->root->c[HI] != NULL )
          { SOGrid_subtree_free(tree->root->c[LO]); tree->root->c[LO] = NULL;
            SOGrid_subtree_free(tree->root->c[HI]); tree->root->c[HI] = NULL;
          }
        leaf_cells=1;
/*         ShatterNode (pDim, tree->root, 0, lo, hi, maxhi, o->minRank, o->maxRank, o->diff_limit, refine_f, &leaf_cells); */
        ShatterNode (pDim, tree->root, 0, lo, hi, maxhi, o->minRank, o->maxRank, o->diff_limit, f, &leaf_cells);
        printf("## DEBUG> o->diff_limit: %f, leaf_cells: %d \n\n", o->diff_limit, leaf_cells);
        if(o->minRank == o->maxRank) break;
      }

    /* Saves the used tree. */
    FILE *wr = open_write(txtcat(o->outName, ".tree"));
    SOGrid_Tree_write(wr, tree);
    fclose(wr);

    /* Gets basis. */
    SOTent_vec tents;
    Basis bas;
    unsigned int dim, maxdim = (1 << o->maxRank);

    if (o->maximal) tents = SOTent_maximal_basis(tree);
    else tents = SOTent_minimal_basis(tree);

    bas = SOFunctionRef_vec_new(tents.nel);
    dim = bas.nel;
    char *basName = txtcat(o->outName, ".bas");
    for (i = 0; i < tents.nel; i++)
      { SOTent t = tents.el[i];
        bas.el[i] = (SOFunction *)SOTentFunction_Make(pDim, t.cx, t.r);
      }

    /* Gets the pre-calculated dot product matrices. */
    SOMatrix Hev, Meval;

    /* <Meval> gets all dot product combinations between minRank to maxRank basis. */
    Meval = SOApprox_ReadMatrix(txtcat(o->matName, "-ev.mat"));
    unsigned int offset = (1 << o->minRank), maxindices;

    /* <Hev> may have only up to the total maxRank basis (grid all refined). */
    Hev.ents = MatEntry_vec_new(maxdim * maxdim);
    Hev.rows = dim; Hev.cols = dim;
    //    Hev.rows = maxdim; Hev.cols = maxdim;

    /* <evaldot> keeps <Meval> in a direct access organization. */
    maxindices = 0;
    for(i = o->minRank; i <= o->maxRank; i++) maxindices += (1 << i);

    if (o->maximal)
      { offset = (1 << pDim);
        maxindices = 0;
        for(i = pDim; i <= o->maxRank; i++) maxindices += (1 << i);
      }

    MatEntry_vec evaldot[maxindices];
    SetGenDotMatrices(evaldot, Meval, offset, maxindices);

    /* Gets the dot product values for the main matrice <Hev>. */
    SetDotMatrices( dim, offset, bas, evaldot, &(Hev.ents) );


    /* Allocates vectors. */ 
    double_vec u = double_vec_new(fDim * dim);
    double_vec b = double_vec_new(fDim * dim); // to receive <basis | f>, f's results on R^fDim
    double_vec uj = double_vec_new(dim); // Column {j} of {u}.
    double_vec bj = double_vec_new(dim); // Column {j} of {b}.
    double_vec y = double_vec_new(dim);
    
    FuncMap NoMap = IdentityMap(fDim, pDim);
    
    SOIntegral_GaussOrder = o->gaussOrder; 

    { for (i = 0; i < b.nel; i++) { b.el[i] = 0.0; } } 

    SOApprox_ComputeRightHandSide(f, &NoMap, bas, w, tree, b, TRUE); 
 
    if (o->cholesky)
      { 
        SOMatrix Mevch = SOMatrix_Cholesky(Hev, 0.0);        
        SOApprox_CholeskySolve(Mevch, fDim, b, u, bj, uj, y, TRUE);
      }
    /* 
    else if (o->conjGrad)
      { SOMatrix G = SOApprox_GetBasisMatrix(o->matName, FALSE);
        SOApprox_GuessSol(v);
        SOApprox_ConjugateGradientSolve(G, b, v, u, bj, uj, TRUE);
      }
    else if (o->gaussSeidel)
      { SOMatrix G = SOApprox_GetBasisMatrix(o->matName, FALSE);
        SOApprox_GuessSol(v);
        SOApprox_GaussSeidelSolve
          ( G, b, 
            o->omega, o->maxIter, o->absTol, o->relTol,
            v, u, bj, uj, TRUE
          );
      }
    */
    else
      { assert(FALSE , "unspecified/invalid solution method"); }
    
    /* Output final solution: */
    SOFunction *g = SOApprox_BuildFunction(bas, u, o->outName, basName, f);
    double fMax[fDim], eMax[fDim];
    SOApprox_PrintMaxErrorValues(f, g, fMax, eMax);

    /* 2D Plotting */
    if((o->plot)&&(pDim == 2)&&(fDim == 1))
      {
        /* Builds error function: */
        Basis error_bas = SOFunctionRef_vec_new(2);
        double_vec errc = double_vec_new(2); errc.el[0] = 1; errc.el[1] = -1;
        error_bas.el[0] = (SOFunction *)f; error_bas.el[1] = (SOFunction *)g;
        SOFunction *e = 
          (SOFunction *)SOLinCombFunction_Make(pDim, fDim,"Errorf.bas", error_bas, errc);
      
        double fPlotMin = 0, fPlotMax = 0;
        double gPlotMin = 0, gPlotMax = 0;
        double ePlotMin = 0, ePlotMax = 0;

        /* Plots the true solution {f}, the approximation {g}, and the error {e}: */
        SOApprox_PlotSingleFunction(f, &NoMap, tree, txtcat(o->outName, "-sol"),  &(o->plt), &fPlotMin, &fPlotMax);
        SOApprox_PlotSingleFunction(g, &NoMap, tree, txtcat(o->outName, "-apr"),  &(o->plt), &gPlotMin, &gPlotMax);

	/* Cortes 1D. */
	double initP[2] = {0.0, 0.5}, finP[2] = {1.0, 0.5}, D1min, D1max;
        SO1DPlot_SingleFunction
        (TRUE, NULL, 300, 200, g, &NoMap, pDim, initP, finP, -0.5, 12.0, 8, 0.10, "corteX", &D1min, &D1max);
	initP[0] = 0.25; initP[1] = 0.0; finP[0] = 0.25; finP[1] = SQRTHALF;
        SO1DPlot_SingleFunction
        (TRUE, NULL, 300, 200, g, &NoMap, pDim, initP, finP, -0.5, 12.0, 8, 0.10, "corteY", &D1min, &D1max);

	o->plt.fRange = o->plt.fRange / 10.0;
	o->plt.fStep = o->plt.fStep / 10.0;
        SOApprox_PlotSingleFunction(e, &NoMap, tree, txtcat(o->outName, "-err"),  &(o->plt), &ePlotMin, &ePlotMax);

        fprintf(stderr, "observed SOLUTION function extrema:");
        fprintf(stderr, " min = %16.12f max = %16.12f\n", fPlotMin, fPlotMax);
        fprintf(stderr, "\n");

        fprintf(stderr, "observed APPROXIMATION function extrema:");
        fprintf(stderr, " min = %16.12f max = %16.12f\n", gPlotMin, gPlotMax);
        fprintf(stderr, "\n");

        fprintf(stderr, "observed ERROR function extrema:");
        fprintf(stderr, " min = %16.12f max = %16.12f\n", ePlotMin, ePlotMax);
        fprintf(stderr, "\n");

        fprintf(stderr, "Basis with %d tents; Grid with %d leaf cells.\n", dim, leaf_cells);

        double distance = 0.0;
        FILE *rd2 = open_read("../SOFindTentBasis/regular10.tree");
        SOGrid_Tree *tree10 = SOGrid_Tree_read(rd2);
        SOFunction_Dot(e, &NoMap, e, &NoMap, (SOFunction *)NULL, tree10, &distance, TRUE);
        distance = sqrt(distance);
        fprintf(stderr, "observed DISTANCE ( |f-g| = sqrt(<f-g, f-g>) ): %.16g \n", distance);

      } 

    return 0;
  } 

#define PPUSAGE SOParams_SetUsage

Options *GetOptions(int argn, char **argc)
  {
    Options *o = (Options *)notnull(malloc(sizeof(Options)), "no mem");
    SOParams_T *pp = SOParams_NewT(stderr, argn, argc);

    PPUSAGE(pp, "SOComputeApprox \\\n");
    PPUSAGE(pp, "  -funcName NAME \\\n");
    PPUSAGE(pp, "  -matName NAME \\\n");
    PPUSAGE(pp, "  -minRank NUM \\\n");
    PPUSAGE(pp, "  -maxRank NUM \\\n");
    PPUSAGE(pp, "  -diff_limit NUM \\\n");
    PPUSAGE(pp, "  [ -maximal ] \\\n");
    PPUSAGE(pp, "  [ -cholesky | \\\n");
    PPUSAGE(pp, "    -gaussSeidel [-omega NUM] [-tol NUM] [-maxIter NUM] | \\\n");
    PPUSAGE(pp, "    -conjGrad ] \\\n");
    PPUSAGE(pp, "  [ -gaussOrder NUM ] \\\n");
    PPUSAGE(pp, "  -outName NAME \\\n");

    /* Plotting options instructions. */
    PPUSAGE(pp, "  [ -plot ] \\\n");
    PPUSAGE(pp, SOPlotParams_FunctionHelp " \n");

    SOParams_GetKeyword(pp, "-funcName");                               
    o->funcName = SOParams_GetNext(pp);  
       
    SOParams_KeywordPresent(pp, "-matName");
    o->matName = SOParams_GetNext(pp);
       
    SOParams_GetKeyword(pp, "-outName");                               
    o->outName = SOParams_GetNext(pp);  
                                                 
    o->cholesky = SOParams_KeywordPresent(pp, "-cholesky");
    o->gaussSeidel = SOParams_KeywordPresent(pp, "-gaussSeidel");
    o->conjGrad = SOParams_KeywordPresent(pp, "-conjGrad");
    if (! (o->cholesky || o->gaussSeidel || o->conjGrad))
      { SOParams_Error(pp, "must specify a linear solution method"); }
    else if ((o->cholesky + o->gaussSeidel + o->conjGrad) > 1)
      { SOParams_Error(pp, "please specify only one linear solution method"); }

    SOParams_KeywordPresent(pp, "-minRank");
    o->minRank = SOParams_GetNextInt(pp, 0, INT_MAX);

    SOParams_KeywordPresent(pp, "-maxRank");
    o->maxRank = SOParams_GetNextInt(pp, 0, INT_MAX);

    SOParams_KeywordPresent(pp, "-diff_limit");
    o->diff_limit = SOParams_GetNextDouble(pp, 0, MAXDOUBLE);

    o->maximal = SOParams_KeywordPresent(pp, "-maximal");

    if (o->gaussSeidel)
      { 
        if (SOParams_KeywordPresent(pp, "-omega"))
          { o->omega = SOParams_GetNextDouble(pp, 0.0, MAXDOUBLE); }
        else
          { o->omega = 1.0; }

        if (SOParams_KeywordPresent(pp, "-maxIter"))
          { o->maxIter = SOParams_GetNextInt(pp, 0, INT_MAX); }
        else
          { o->maxIter = 20; }

        if (SOParams_KeywordPresent(pp, "-absTol"))
          { o->absTol = SOParams_GetNextDouble(pp, 0.0, MAXDOUBLE); }
        else
          { o->absTol = 0.0; }

        if (SOParams_KeywordPresent(pp, "-relTol"))
          { o->relTol = SOParams_GetNextDouble(pp, 0.0, MAXDOUBLE); }
        else
          { o->relTol = 0.0; }
      }

    SOParams_GetKeyword(pp, "-gaussOrder");
    o->gaussOrder = SOParams_GetNextInt(pp, 1, INT_MAX);

    /* Plotting options: */

    o->plot = SOParams_KeywordPresent(pp, "-plot");
    o->plt = SOPlotParams_FunctionParse(pp);

    SOParams_Finish(pp);
    return o;
  }


void ShatterNode
  ( SOGrid_Dim pDim,
    SOGrid_Node *p,
    SOGrid_Rank r,
    double *lo, 
    double *hi,
    double *maxhi,
    int minRank,
    int maxRank,
    double diff_limit,
    SOFunction *f,
    int *leaf_cells
  )
  {
    bool split;//, border = FALSE;
    int i, j;
    double mid_val, maxdiff=0, diff, lo_child[pDim], hi_child[pDim]; //, abs_mid_val;
    double pnt_mid[pDim], pnt_lo[pDim], pnt_hi[pDim], pnt_lo_val, pnt_hi_val;

    for(i = 0; i < pDim; i++) pnt_mid[i] = (double) (hi[i] + lo[i]) / 2.0;
    f->m->eval(f, pnt_mid, &mid_val); 

    for(j = 0; j < pDim; j++)
      {
        for(i = 0; i < pDim; i++)
          if(i != j){ pnt_lo[i] =  pnt_mid[i];  pnt_hi[i] =  pnt_mid[i]; }
 
        pnt_lo[j] =  lo[j];  f->m->eval(f, pnt_lo, &pnt_lo_val);
        pnt_hi[j] =  hi[j];  f->m->eval(f, pnt_hi, &pnt_hi_val);

	//        printf(" ###### DEBUG-> pnt_lo[0]: %f, pnt_lo[1]: %f \n\n", pnt_lo[0], pnt_lo[1]);
	//        printf(" ###### DEBUG-> pnt_hi[0]: %f, pnt_hi[1]: %f \n\n", pnt_hi[0], pnt_hi[1]);

        diff = fabs(mid_val - (pnt_lo_val + pnt_hi_val)/2.0);
        if(diff > maxdiff) maxdiff = diff;

	//        if(lo[j] == 0.0 || hi[j] == maxhi[j]){ border = TRUE; abs_mid_val = fabs(mid_val); }
      }

    //    printf(" ###### DEBUG maxdiff: %.16g  diff_limit: %.16g \n\n", maxdiff, diff_limit);

    split = (maxdiff > diff_limit);// || (border && abs_mid_val > 2.0 * diff_limit));

    if (r < minRank){split = TRUE;}
    if (r >= maxRank){split = FALSE;}

    if (split)
      { 
        *leaf_cells = *leaf_cells + 1;
        if( p->c[LO] == NULL && p->c[HI] == NULL )SOGrid_Node_split(p);

        /* Compute low and high coordinates of low child. */
        for (i = 0; i < pDim; i++)
          { lo_child[i] = lo[i]; 
            hi_child[i] = (i == (r%pDim) ? (lo[i]+hi[i])/2 : hi[i]);
          }  
        ShatterNode(pDim, p->c[LO], r+1, lo_child, hi_child, maxhi, minRank, maxRank, diff_limit, f, leaf_cells);
        
        /* Compute low and high coordinates of high child. */
        for (i = 0; i < pDim; i++)
          { lo_child[i] = (i == (r%pDim) ? (lo[i]+hi[i])/2 : lo[i]); 
            hi_child[i] = hi[i];
          }
        ShatterNode(pDim, p->c[HI], r+1, lo_child, hi_child, maxhi, minRank, maxRank, diff_limit, f, leaf_cells);
      }
 
  }


void SetGenDotMatrices
  ( MatEntry_vec *evaldot, 
    SOMatrix Meval, 
    unsigned int offset,
    unsigned int maxindices
  )
  {
    unsigned int i, index, indexi;

    for(i = 0; i < maxindices; i++)
      {
        evaldot[i] = MatEntry_vec_new(Meval.cols);
        evaldot[i].nel = 0;
      }

    i = 0;
    while(i < Meval.ents.nel)
      {
        index = Meval.ents.el[i].row;
        while(Meval.ents.el[i].row == index)
	  {
            indexi = Meval.ents.el[i].row - offset;
 
            evaldot[indexi].el[evaldot[indexi].nel].row = Meval.ents.el[i].row;
            evaldot[indexi].el[evaldot[indexi].nel].col = Meval.ents.el[i].col;
            evaldot[indexi].el[evaldot[indexi].nel].va = Meval.ents.el[i].va;
            evaldot[indexi].nel++;
            i++;
	  }
      }
  }


void SetDotMatrices
  ( unsigned int dim, 
    unsigned int offset,
    Basis bas, 
    MatEntry_vec *evaldot, 
    MatEntry_vec *Hev_ents
  )
  {
    int i, j, k, hi_ind, low_ind;
    SOTentFunction *itf, *jtf;

    Hev_ents->nel = 0;

    for(i = 0; i < dim; i++)
      for(j = 0; j < dim; j++)
        {
          itf = (SOTentFunction *)bas.el[i];
          jtf = (SOTentFunction *)bas.el[j];
          
          if(itf->d->index > jtf->d->index){ hi_ind = itf->d->index; low_ind = jtf->d->index; }
          else { hi_ind = jtf->d->index; low_ind = itf->d->index; }
   
          k = 0;
          while(k < evaldot[hi_ind - offset].nel)
            {
              if(evaldot[hi_ind - offset].el[k].col == low_ind) 
                { 
                  Hev_ents->el[Hev_ents->nel].row = i;
                  Hev_ents->el[Hev_ents->nel].col = j;
                  Hev_ents->el[Hev_ents->nel].va = evaldot[hi_ind - offset].el[k].va;
                  Hev_ents->nel++;
		}
              k++;
	    }

	}

  }