// simeq_newton5.java  solve nonlinear system of equations
//                     method: newton iteration using Jacobian
//                     use list for higher order terms
//
//  a maximal term can have up to power of 3 divided by power of 2
//  in this implementation (if we had symbolic differentiation, no limit)
//  a term could be  x1 , x1*x3 , x2/x5 , x1*x2*x3/(x4*x5) ,  x2^3/x4^2 
//
//  Solve by initial guess at values of x1, x2, x3 computing products
//    X_next = X_initial - J_initial^-1 * (A *  X_initial - Y)
//    in general X_next = X_prev - (J_prev^-1 * (A * X_prev - Y))*b
//                                 where 0 < b < 1, often 0.5, for stability
//
//  solved when  abs sum each row A * X_next -Y < epsilon
//
//  It may stall, stop if abs(X_next-X_prev)<epsilon
//  It may diverge, stop, indicate no solution
//                        (or try a different initial guess)
//  It may oscillate, stop, indicate no solution
//                        (or try a different initial guess)
//
//  The matrix equation could be:
//
//   x1    x2    x3   x1*x2 x3*x2^2 x2/x3 x1^3/x3^2 ... 
//
// | A1,1  A1,2  A1,3  A1,4   A1,5  A1,6  A1,7 ...| | x1         | |y1 |
// | A2,1  A2,2  A2,3  A2,4   A2,5  A2,6  A2,7 ...| | x2         | |y2 |
// | A3,1  A3,2  A3,3  A3,4   A3,5  A3,6  A3,7 ...| | x3         | |y3 |
// | A4,1  A4,2  A4,3  A4,4   A4,5  A4,6  A4,7 ...| | x1*x2      | |y4 |
// | A5,1  A5,2  A5,3  A5,4   A5,5  A5,6  A5,7 ...|*| x3*x2^2    |=|y5 |
// | A6,1  A6,2  A6,3  A6,4   A6,5  A6,6  A6,7 ...| | x2/x3      | |y6 |
// | A7,1  A7,2  A7,3  A7,4   A7,5  A7,6  A7,7 ...| | x1^3/x3^2  | |y7 |
// | ...                                       ...| |            | |   |
//
//  The Jacobian, J is the numerically computed derivatives based on A, X_prev
//  The derivative coefficients may be computed once based on A and the
//  unknown variables. The numeric value plugs in the X_prev values.
//
//  J1,1 = A1,1 + A1,4*x2 + 3*A1,7*x1^2/x3^2               deriv Row 1 wrt x1
//  J1,2 = A1,2 + A1,4*x1 + 2*A1,5*x3*x2 + A1,6/x3         deriv Row 1 wrt x2
//  J1,3 = A1,3 + A1,5*x2^2 + -A1,6*x2/x3^2 + -2*A1,7*x1^2/x3^3  
//                                                         deriv Row 1 wrt x3
//  J2,1 = A2,1 + A2,4*x2 + 3*A2,7*x1^2/x3^2               deriv Row 2 wrt x1
//  J2,2 = A2,2 + A2,4*x1 + 2*A2,5*x3*x2 + A2,6/x3         deriv Row 2 wrt x2
//  J2,3 = A1,3 + A2,5*x2^2 + -A2,6*x2/x3^2 + -2*A2,7*x1^2/x3^3  
//                                                         deriv Row 2 wrt x3
//  J3,1 = A3,1 + A3,4*x2 + 3*A3,7*x1^2/x3^2               deriv Row 3 wrt x1
//  J3,2 = A3,2 + A3,4*x1 + 2*A3,5*x3*x2 + A3,6/x3         deriv Row 3 wrt x2
//  J3,3 = A3,3 + A3,5*x2^2 + -A3,6*x2/x3^2 + -2*A3,7*x1^2/x3^3  
//                                                         deriv Row 3 wrt x3
//  ...
//  etc.
//
//  Code or a data structure must be available to know the equation
//  of the entries in the Y vector. Symbolic computation of
//  derivatives is assumed to be available, when needed.
//

public class simeq_newton5
{
  // derivatives computed from var1, var2, var3, vari1, vari2 to
  // load Ja, initial guess in X[0..n-1], returned solution in X[0..n+nlin-1]
  // n is number of variables, nlin is number of nonlinear terms
  // var1 etc are [0] to [n+nlin-1]

  double eps = 1.0e-6; // (default)
  double b = 1.0; // stability factor (default)
  int maxiter = 8; // maximum number of iterations (default)
  boolean debug = false;  // should add monitor control
  boolean varerr = false; // input consistency flag

  public simeq_newton5()
  {
    // use default setup
  }

  // optional initialization
  public simeq_newton5(double eps, double b, int maxiter, boolean debug)
  {
    this.eps = eps;
    this.b = b;
    this.maxiter = maxiter;
    this.debug = debug;
  }

  public void simeq(int n, int nlin, double A[][], double Y[],
                    int var1[], int var2[], int var3[],
                    int vari1[], int vari2[], double X[])
  {
    double resid = 1.0e12;  // residual from last iteration
    double presid; // residual from prior to last iteration
    double Ja[][] = new double[n][n]; // Jacobian inverted in place
    double X_resid[] = new double[n];
    double Next_resid[] = new double[n];
    double X_tmp2[] = new double[n+nlin];
    double X_next[] = new double[n+nlin];
    double T; // term

    if(debug) System.out.println("simeq_newton5 running n="+n+", nlin="+nlin);

    // check for consistent input
    for(int i=0; i<n; i++) // single variables
    {
      if(var1[i]!=i)
      {
        System.out.println("var1["+i+"]="+var1[i]+" must be -1");
        varerr=true;
      }
      if(var2[i]!=-1)
      {
        System.out.println("var2["+i+"]="+var2[i]+" must be -1");
        varerr=true;
      }
      if(var3[i]!=-1)
      {
        System.out.println("var3["+i+"]="+var3[i]+" must be -1");
        varerr=true;
      }
      if(vari1[i]!=-1)
      {
        System.out.println("vari1["+i+"]="+vari1[i]+" must be -1");
        varerr=true;
      }
      if(vari2[i]!=-1)
      {
        System.out.println("vari2["+i+"]="+vari2[i]+" must be -1");
        varerr=true;
      }
    } // end single variable check
    for(int i=n; i<n+nlin; i++) // nonlinear terms
    {
      if(var1[i]<-1 || var1[i]>=n)
      {
	System.out.println("var1["+i+"] must be -1 to "+(n-1));
        varerr=true;
      }
      if(var2[i]<-1 || var2[i]>=n)
      {
	System.out.println("var2["+i+"] must be -1 to "+(n-1));
        varerr=true;
      }
      if(var3[i]<-1 || var3[i]>=n)
      {
	System.out.println("var3["+i+"] must be -1 to "+(n-1));
        varerr=true;
      }
      if(vari1[i]<-1 || vari1[i]>=n)
      {
	System.out.println("vari1["+i+"] must be -1 to "+(n-1));
        varerr=true;
      }
      if(vari2[i]<-1 || vari2[i]>=n)
      {
	System.out.println("vari2["+i+"] must be -1 to "+(n-1));
        varerr=true;
      }
    } // end nonlinear term check

    if(varerr)
    {
      System.out.println("var1  var2  var3  vari1 vari2");
      for(int i=0; i<n+nlin; i++)
      {
	System.out.println(var1[i]+"    "+var2[i]+"    "+var3[i]+
                           "    "+vari1[i]+"   "+vari2[i]);
      }
      System.exit(1);
    }

    // setup nonlinear terms, this does not need to be set up by caller
    for(int i=n; i<n+nlin; i++)
    {
      X[i] = 1.0; // there should be no entry for null term
      if(var1[i]>=0) X[i] *= X[var1[i]];
      if(var2[i]>=0) X[i] *= X[var2[i]];
      if(var3[i]>=0) X[i] *= X[var3[i]];
      if(vari1[i]>=0) X[i] /= X[vari1[i]];
      if(vari2[i]>=0) X[i] /= X[vari2[i]];
    }
    // debug print
    if(debug)
    {
      System.out.println("initial guess at solution, can be bad");
      for(int i=0; i<n+nlin; i++)
      {
        System.out.println("X["+i+"]="+X[i]);
      }
      System.out.println(" ");
    }

    // compute residual
    resid = 0.0; 
    for(int i=0; i<n; i++)
    {
      X_resid[i] = 0.0;
      for(int j=0; j<n+nlin; j++)
      {
        X_resid[i] += A[i][j]*X[j];
      }
      X_resid[i] -= Y[i]; // X_resid used later if not solved
      resid += Math.abs(X_resid[i]);
    }
    // debug print
    if(debug)
    {
      for(int i=0; i<n; i++)
      {
        System.out.println("residual X_resid["+i+"]="+X_resid[i]);
      }
      System.out.println(" ");
      System.out.println("simeq_newton5 itr "+0+",         initial residual="
                         +resid);
    }
    presid = resid;

    // iterate to find solution
    for(int itr=1; itr<=maxiter; itr++)
    {
      // check for convergence
      if(resid<eps) break;

      // compute Jacobian
      for(int i=0; i<n; i++) // equation
      {
	for(int j=0; j<n; j++) // variable
        {
	  Ja[i][j] = A[i][j]; // each linear term
	  for(int k=n; k<n+nlin; k++) // nonlinear term
	  {
            /* add derivative of nonlinear term to Ja[i][j] */
            /* first group with no derivative of inverse term */
	    if(vari1[k]!=j && vari2[k]!=j) /* no derivative, may contribute*/
	    { 
	      if(var1[k]==j && var2[k]==j && var3[k]==j)       /* X^3  */
	      {
	        T = A[i][k]*3.0*X[var2[k]]*X[var3[k]];         /* 3X^2 */
	        if(vari1[k]>=0) T /= X[vari1[k]];              /* vari1 */
                if(vari2[k]>=0) T /= X[vari2[k]];              /* vari2 */
	        Ja[i][j] += T;
	      }
	      else if(var1[k]==j && var2[k]==j)                /* X^2 x */
	      {
	        T = A[i][k]*2.0*X[var2[k]];                    /* 2X x  */
	        if(var3[k]>=0)  T *= X[var3[k]];               /* var3 */
	        if(vari1[k]>=0) T /= X[vari1[k]];              /* vari1 */
                if(vari2[k]>=0) T /= X[vari2[k]];              /* vari2 */
	        Ja[i][j] += T;
	      }
	      else if(var2[k]==j && var3[k]==j)                /* x X^2 */
	      {
	        T = A[i][k]*2.0*X[var3[k]];                    /* 2X x  */
	        if(var1[k]>=0)  T *= X[var1[k]];               /* var1 */
	        if(vari1[k]>=0) T /= X[vari1[k]];              /* vari1 */
                if(vari2[k]>=0) T /= X[vari2[k]];              /* vari2 */
	        Ja[i][j] += T;
	      }
              else if(var3[k]==j && var1[k]==j)                /* x X^2 */
	      {
	        T = A[i][k]*2.0*X[var1[k]];                    /* 2X x  */
	        if(var2[k]>=0)  T *= X[var2[k]];               /* var2 */
	        if(vari1[k]>=0) T /= X[vari1[k]];              /* vari1 */
                if(vari2[k]>=0) T /= X[vari2[k]];              /* vari2 */
	        Ja[i][j] += T;
	      }
              else if(var1[k]==j)                              /* only X */
	      {
		T = A[i][k];
	        if(var2[k]>=0) T *= X[var2[k]];                /* var2 */
                if(var3[k]>=0) T *= X[var3[k]];                /* var3 */
	        if(vari1[k]>=0) T /= X[vari1[k]];              /* vari1 */
                if(vari2[k]>=0) T /= X[vari2[k]];              /* vari2 */
		Ja[i][j] += T;
	      }
              else if(var2[k]==j)                              /* only X */
	      {		
		T = A[i][k];
                if(var1[k]>=0) T *= X[var1[k]];                /* var1 */
                if(var3[k]>=0) T *= X[var3[k]];                /* var3 */
	        if(vari1[k]>=0) T /= X[vari1[k]];              /* vari1 */
                if(vari2[k]>=0) T /= X[vari2[k]];              /* vari2 */
		Ja[i][j] += T;
              }
              else if(var3[k]==j)                              /* only X */
	      {
		T = A[i][k];
	        if(var2[k]>=0) T *= X[var2[k]];                /* var2 */
                if(var1[k]>=0) T *= X[var1[k]];                /* var1 */
	        if(vari1[k]>=0) T /= X[vari1[k]];              /* vari1 */
                if(vari2[k]>=0) T /= X[vari2[k]];              /* vari2 */
		Ja[i][j] += T;
	      }
	    } /* end no derivative of inverse term */
            else if(vari1[k]==j || vari2[k]==j)/* do derivative inverse term */
	    {
	      /* three cases
	      /* d/dx 1/x = -1/x^2    d/dx 1/x^2 = -2/x^3 */
	      if(vari1[k]==j && vari2[k]==j)                   /* 1/x^2 */
	      {
		T = -2.0*A[i][k]/(X[vari1[k]]*X[vari1[k]]*X[vari1[k]]);
		                                               /* -2/x^3 */
                if(var1[k]>=0) T *= X[var1[k]];                /* var1 */
                if(var2[k]>=0) T *= X[var2[k]];                /* var2 */
                if(var3[k]>=0) T *= X[var3[k]];                /* var3 */
		Ja[i][j] += T;
	      }
	      else if(vari1[k]==j)                             /* 1/xi1 */
	      {
		T = -A[i][k]/(X[vari1[k]]*X[vari1[k]]);
		                                               /* -1/x^2 */
                if(var1[k]>=0) T *= X[var1[k]];                /* var1 */
                if(var2[k]>=0) T *= X[var2[k]];                /* var2 */
                if(var3[k]>=0) T *= X[var3[k]];                /* var3 */
                if(vari2[k]>=0) T /= X[vari2[k]];              /* vari2 */
		Ja[i][j] += T;
	      }
	      else if(vari2[k]==j)                             /* 1/xi2 */
	      {
		T = -A[i][k]/(X[vari2[k]]*X[vari2[k]]);        /* -1/x^2 */
                if(var1[k]>=0) T *= X[var1[k]];                /* var1 */
                if(var2[k]>=0) T *= X[var2[k]];                /* var2 */
                if(var3[k]>=0) T *= X[var3[k]];                /* var3 */
                if(vari1[k]>=0) T /= X[vari1[k]];              /* vari1 */
		Ja[i][j] += T;
	      } 
	    } // end of inverse derivative 
	  } // end k
	} // end j
      } // end i
      if(debug && n<2)
      {
        System.out.println("Ja computed ");
        for(int i=0; i<n; i++)
        {
          for(int j=0; j<n; j++)
          {
	    System.out.println("Ja["+i+"]["+j+"]="+Ja[i][j]);
	  }
        }
        System.out.println(" ");
      }

      // invert Ja
      new invert(Ja);
      if(debug && n<2)
      {
        System.out.println("Ja inverted ");
        for(int i=0; i<n; i++)
        {
          for(int j=0; j<n; j++)
          {
	    System.out.println("Ja["+i+"]["+j+"]="+Ja[i][j]);
	  }
        }
        System.out.println(" ");
      }

      // X_resid = A X - Y        error vector
      // X_tmp2 = J^-1 * X_resid  raw correction vector
      // X_next = X - (X_tmp2)*b weighted correction vector
      // X      = X_next         next X with computed higher terms

      for(int i=0; i<n; i++)
      {
	X_tmp2[i] = 0.0;
        for(int j=0; j<n; j++)
        {
	  X_tmp2[i] += Ja[i][j]*X_resid[j]; 
        }
      }
      // debug print
      if(debug & n<2)
      {
        for(int i=0; i<n; i++)
        {
	  System.out.println("change X_tmp2["+i+"]="+X_tmp2[i]);
        }
        System.out.println(" ");
      }
      for(int i=0; i<n; i++)
      {
        X_next[i] = X[i] - b*X_tmp2[i];
      }
      for(int i=n; i<n+nlin; i++) // compute non linear terms
      {
	X_next[i] = 1.0;
	if(var1[i]>=0) X_next[i] = X_next[var1[i]];
        if(var2[i]>=0) X_next[i] *= X_next[var2[i]];
        if(var3[i]>=0) X_next[i] *= X_next[var3[i]];
        if(vari1[i]>=0) X_next[i] /= X_next[vari1[i]];
        if(vari2[i]>=0) X_next[i] /= X_next[vari2[i]];
      }

      if(debug && n<6)
      {
        for(int i=0; i<n+nlin; i++)
        {
          System.out.println("X_next["+i+"]="+X_next[i]);
        }
        System.out.println(" ");
      }

      // compute residual
      resid = 0.0; 
      for(int i=0; i<n; i++)
      {
        Next_resid[i] = 0.0;
        for(int j=0; j<n+nlin; j++)
        {
	  Next_resid[i] += A[i][j]*X_next[j];
        }
        Next_resid[i] -= Y[i]; // X_resid used later if not solved
        resid += Math.abs(Next_resid[i]);
      }
      // debug print
      if(debug)
      {
        for(int i=0; i<n; i++)
        {
	  System.out.println("residual X_resid["+i+"]="+X_resid[i]);
        }
        System.out.println(" ");
      }
      System.out.println("simeq_newton5 itr "+itr+", prev="+presid+
                         ", residual="+resid);
      if(resid<presid) // progress, increase b
      {
	if(b<1.0)
	{
          b = b * 1.4143;
          if(b>1.0) b = 1.0;
          System.out.println("b increased to "+b);
	}
        // iterate
        presid = resid; 
        for(int i=0; i<n+nlin; i++) X[i] = X_next[i];
        for(int i=0; i<n; i++) X_resid[i] = Next_resid[i];
      }
      else // worse, reduce b
      {
	b = b*0.5;
        System.out.println("b reduced to "+b);
      }
      if(debug) System.out.println(" ");
    } // end iteration    

    if(debug)
    {
      if(resid<eps) System.out.println("converged last residual = "+resid);
      else System.out.println("not converged last residual = "+resid);
      System.out.println("simeq_newton5.java finished");
    }
  }
} // end class simeq_newton5