/* SOUniSolve -- Solves the screened Poisson equation by uniscale iteration. */
/* Last edited on 2003-06-24 17:23:08 by ra014852 */

/*
  This program solves a differential equation on the root cell
  
    D(f)(p) == FMap(f(p), p)       (1)
    
  where {D} is a linear differential operator, {FMap} is a given
  function, and {f} is the function to be determined. For example, 
  in the screened Poisson equation on the root cell, the differential
  operator is
  
    D(f)(p) = Lap(f)(p) + c*f(p)       (2)
    
  where {Lap} is the Laplacian operator, and {c} is a given
  constant.
  
  It is assumed that an approximate solution {g} to (2) in some
  finite-dimensional function space {V} can be found by solving the
  non-linear system
  
    H a == b         (3)
    
  where 
  
    {a[i]} is coefficient {i} of the solution 
           {g} (to be determined) in the basis {bas}; 
  
    {b[i] == <F | bas[i]>}, where {F} is the   function
            {p -> FMap(g(p),p)} and {<|>} is the scalar product;
            
    {H[i,j] == <D(bas[i]) | bas[j]>}
           where {D} is the differential operator of eq.(1).
           
  For the screened Poisson equation, we have
  
     {H[i,j] == <Lap(bas[i]) | bas[j]> + c*<bas[i]|bas[j]>}
     
  If the basis functions {bas[i]} are at least C1, the scalar product
  {G[i,j] = <Lap(bas[i]) | bas[j]>} can be reduced to 
  
     {G[i,j] = <Grd(bas[i]) | Grd(bas[j])>}
     
  where {Grd} is the gradient operator. 
        
  Note that, in general, the system (3) is not linear, since the
  right-hand side {b} depends on the function {g}, and therefore on
  the coefficient vector {a}. So (3) must be solved iteratively.
*/

#include <SOGrid.h>
#include <SOMatrix.h>
#include <SOFunction.h>
#include <SOProcFunction.h>
#include <SOLinCombFunction.h>
#include <SOApprox.h>
#include <SOFuncMap.h> 
#include <SOParams.h>
#include <SOBasic.h>
#include <SOPlotParams.h>

#include <rn.h>
#include <js.h>
#include <math.h>
#include <stdio.h>
#include <string.h>
#include <values.h>
#include <limits.h>


typedef struct Options
  { char *rhsName;     /* Name of right-hand-side function. */
    char *basisName;   /* Prefix for basis filenames (minus extension). */
    char *matName;     /* Prefix for matrix filenames (default {basisName}). */
    char *outName;     /* Prefix for output filenames. */
    nat maxIter;       /* Maximum iterations. */
    SOGrid_Dim pDim;   /* Dimension of {f}'s domain. */
    SOGrid_Dim fDim;   /* Dimension of {f}'s values. */
    double relTol;     /* Max change in {a} coeffs between iterations. */
    double absTol;     /* Max error compared to true sol. */
    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. */
    double omega;      /* Relaxation factor for Gauss-Seidel. */
    /* Plotting options: */
    bool plotFinal;    /* TRUE to plot final solution and error. */
    bool plotAll;      /* TRUE to plot every iteration and error. */
    PlotOptions plt;   /* Plot format and style options. */
  } Options;

Options GetOptions(int argn, char **argc);
SOFunction *GetTrueSolution(char *solName, SOGrid_Dim pDim, SOGrid_Dim fDim);
char *SolutionName(char *outName, int iter);

void ComputeRightHandSide
  ( SOGrid_Tree *tree,  /* SOGrid to use for integration. */
    Basis bas,          /* Basis for the approximation space. */
    SOFunction *g,      /* Current solution. */
    FuncMap *FMap,      /* Right-hand side of differential eqn. */
    nat depth,          /* Depth of sampling grid in each triangle */
    double_vec b        /* OUT: Right-hand-side of (3). */ 
  );
  /* Computes the right-hand side of the system (3) for the given
    solution {g} which has coordinates {a} relative to basis {bas}. */

int main(int argn, char **argc)
  { Options o = GetOptions(argn, argc);
    int i;
    int m = o.pDim; 
    int n = o.fDim;

    /* Gets basis */
    char *basName = txtcat(o.basisName, ".bas");
    Basis bas = SOApprox_ReadBasis(basName);
    unsigned int dim = bas.nel;

    /* Gets the basis SOGrid_Tree {tree}. */ 
    SOGrid_Tree *tree;  
    char treeFile[100];
    strcpy(treeFile, o.basisName); i = strlen(treeFile);

    treeFile[i - 5] = treeFile[i - 1]; // max rank number.
    if(treeFile[i - 1] == '0')// || treeFile[i - 1] == '2')
      { treeFile[i - 6] = treeFile[i - 2];  treeFile[i - 5] = treeFile[i - 1]; } // max rank number.

    treeFile[i - 4] = '.'; treeFile[i - 3] = 't'; treeFile[i - 2] = 'r';
    treeFile[i - 1] = 'e'; treeFile[i] = 'e'; treeFile[i + 1] = '\0';
    FILE *rd = open_read(treeFile);
    tree = SOGrid_Tree_read(rd);


    FuncMap NoMap = IdentityMap(n, m);
    SOFuncMap_Data FD = SOFuncMap_FromName(o.rhsName, n, m, n);
    SOFunction *sol = GetTrueSolution(FD.sol, m, n);
    
    double_vec a = double_vec_new(dim); /* Coeffs of current solution. */ // vDim dimensional?????
    double_vec b = double_vec_new(dim); /* Right-hand side vector. */     // vDim dimensional?????

    double_vec r = double_vec_new(dim);
    double_vec s = double_vec_new(dim);
    double_vec t = double_vec_new(dim);

    double_vec aNew = double_vec_new(dim);                                // vDim dimensional?????
    
    SOMatrix H = SOMatrix_Null(dim,dim);  /* The matrix of (3), if needed. */
    SOMatrix HL = SOMatrix_Null(dim,dim); /* Lower Cholesky factor, if needed. */
    
    SOFunction *g;  /* Current solution. */
    
    FILE *errWr = open_write(txtcat(o.outName, ".erp"));
   
    double sPlotMin = 0, sPlotMax = 0;
    double gPlotMin = 0, gPlotMax = 0;    
    double fMax[n], eMax[n];
    double d, norm;
    unsigned int iter = 0;
    char *uName;
    
    /* Obtain the necessary matrices: */
    if (o.cholesky)
      { HL = SOApprox_GetHelmholtzMatrix(o.matName, FD.coeff, TRUE); }
    else
      { H = SOApprox_GetHelmholtzMatrix(o.matName, FD.coeff, FALSE); }

    SOApprox_GuessSol(a); 
    while (TRUE)
      { 
        uName = txtcat(o.outName, txtcat("-", fmtint(iter, 4)));

        g = SOApprox_BuildFunction(bas, a, uName, basName, sol);

        /* Plots the approximation {g}, if required: */
        if((o.plotAll)&&(m == 2)&&(n == 1)) 
          {
            SOApprox_PlotSingleFunction(g, &NoMap, tree, uName,  &(o.plt), &gPlotMin, &gPlotMax);

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

	SOApprox_PrintMaxErrorValues(g, sol, fMax, eMax);

        fprintf(errWr, "%6d %16.12f\n",  iter, eMax[X]);

        if ((iter > 0) && ((d <= o.relTol*norm) && (d <= o.absTol))) { break; }
        if ((eMax[X] <= o.absTol) || (iter >= o.maxIter)) { break; }
 
        SOApprox_ComputeRightHandSide(g, &FD.FMap, bas, NULL, tree, b, TRUE);
        if (o.cholesky)
          { SOApprox_CholeskySolve(HL, 1, b, aNew, r, s, t, TRUE); }
        else
          { assert(FALSE , "no solution method?"); }

        d = SOApprox_UpdateSolution(aNew, a);
        norm = rn_norm(dim, a.el);
        iter++;
      }

    fclose(errWr);

    /* 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 *)sol; error_bas.el[1] = (SOFunction *)g;
    SOFunction *e = (SOFunction *)SOLinCombFunction_Make(m, n,"Errorf.bas", error_bas, errc);
    double ePlotMin = 0, ePlotMax = 0;

    /* 2D Plotting */
    /* Plots the solution {sol} and the final approximation {g}, if required: */
    if((o.plotFinal)&&(m == 2)&&(n == 1)) 
      {

        SOApprox_PlotSingleFunction(sol, &NoMap, tree, txtcat(o.outName, "-sol"),  &(o.plt), &sPlotMin, &sPlotMax);
        printf("observed SOLUTION function extrema:");
        printf(" min = %16.12f max = %16.12f\n", sPlotMin, sPlotMax);
        printf("\n");

        if(!o.plotAll)
          {
            SOApprox_PlotSingleFunction(g, &NoMap, tree, txtcat(o.outName, "-app"),  &(o.plt), &gPlotMin, &gPlotMax);
            printf("observed APPROXIMATION %s function extrema:", uName);
            printf(" min = %16.12f max = %16.12f\n", gPlotMin, gPlotMax);
            printf("\n");
          }

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


    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);
    printf("observed DISTANCE ( |f-g| = sqrt(<f-g, f-g>) ): %.16g \n", distance);

    return 0;
  }
  
SOFunction *GetTrueSolution(char *solName, SOGrid_Dim pDim, SOGrid_Dim fDim)
  { SOFunction *f = (SOFunction *)SOProcFunction_FromName(solName, pDim, fDim);
    if (f == NULL) 
      { fprintf(stderr, "Unknown solution = \"%s\"\n", solName);
        assert(FALSE, "aborted");
      }
    return f;
  }

char *SolutionName(char *outName, int iter)
  {
    if ((iter > 0) && (iter < INT_MAX))
      { return txtcat(outName, txtcat("-", fmtint(iter,4))); }
    else
      { return outName; }
  }

#define PPUSAGE SOParams_SetUsage

Options GetOptions(int argn, char **argc)
  {
    Options o;
    SOParams_T *pp = SOParams_NewT(stderr, argn, argc);

    PPUSAGE(pp, "SOUNiSolve \\\n");
    PPUSAGE(pp, "  -rhsName NAME \\\n");
    PPUSAGE(pp, "  -basisName NAME [ -matName NAME ] \\\n");
    PPUSAGE(pp, "  [ -cholesky | \\\n");
    PPUSAGE(pp, "    -gaussSeidel [-omega NUM] | \\\n");
    PPUSAGE(pp, "    -conjGrad ] \\\n");
    PPUSAGE(pp, "  [ -maxIter NUM ] [ -relTol NUM ] [ -absTol NUM ] \\\n");
    PPUSAGE(pp, "  [ -gaussOrder NUM ] \\\n");
    PPUSAGE(pp, "  -outName NAME \\\n");
    PPUSAGE(pp, "  -pDim pdim \\\n");
    PPUSAGE(pp, "  -fDim fdim \\\n");
    PPUSAGE(pp, "  [ -plotAll ] [ -plotFinal ] \\\n");
    PPUSAGE(pp, SOPlotParams_FunctionHelp " \n");

    SOParams_GetKeyword(pp, "-rhsName");                               
    o.rhsName = SOParams_GetNext(pp);  

    SOParams_GetKeyword(pp, "-basisName");                               
    o.basisName = SOParams_GetNext(pp);  
       
    if (SOParams_KeywordPresent(pp, "-matName"))
      { o.matName = SOParams_GetNext(pp); }
    else
      { o.matName = o.basisName; }

    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"); }

    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, "-relTol"))
      { o.relTol = SOParams_GetNextDouble(pp, 0.0, MAXDOUBLE); }
    else
      { o.relTol = 0.0; }

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

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

    SOParams_GetKeyword(pp, "-pDim");                               
    o.pDim = SOParams_GetNextInt(pp, 1, 4);  

    SOParams_GetKeyword(pp, "-fDim");                               
    o.fDim = SOParams_GetNextInt(pp, 1, 4);

    /* Plotting options: */

    o.plotAll = SOParams_KeywordPresent(pp, "-plotAll");

    o.plotFinal = SOParams_KeywordPresent(pp, "-plotFinal");

    o.plt = SOPlotParams_FunctionParse(pp);
    
    SOParams_Finish(pp);
    return o;
  }