// test_simeq_newton3 .swift  solve  A X = Y  given matrix A and vector Y

#if os(OSX) || os(iOS) // for libraries, portable for OSX and Linux
import Foundation  // or Playground
#elseif os(Linux)
import Glibc
#endif 

// needs simeq_newton3.swift imported, included below
// needs matinv.swift        imported, included below

print("test_simeq_newton3.swift running")

let n = 4    // total number of unknown variables  rows of A matrix
let nlin = 6 // nonlinear unknowns, products var2
let ntot = n + nlin    // columns of A matrix
let var1 = [ 0, 1, 2, 3, 0, 0, 0, 1, 1, 2] // zero based subscript
let var2 = [-1,-1,-1,-1, 1, 2, 3, 2, 3, 3] // e.g. last is x3*x4
let var3 = [-1,-1,-1,-1,-1,-1,-1,-1,-1,-1] // possible cube or 3 term
print("n=\(n)  nlin=\(nlin)")
print("variables = x0, x1, x2, x3")
print("nonlinear = x0*x1, x0*x2, x0*x3, x1*x2, x1*x3, x2*x3")
var A =  [[Double]](repeating:[Double](repeating:0.0,count:n+nlin),count:n)
var X =  [Double](repeating:0.0,count:n+nlin)
var X1 = [Double](repeating:0.0,count:n+nlin)
var Y =  [Double](repeating:0.0,count:n)
var X_soln = [Double](repeating:0.0,count:n+nlin)
let Xname = ["X0", "X1", "X2", "X3"] // all linear terms

// build first test case
for i in 0..<n {
  for j in 0..<n+nlin {
    A[i][j] = rnd()
    print("A[\(i)][\(j)]= \(A[i][j])")
  } // end j
} // end i
print("A generated random")
print("accuracy of results vary with generated random matrix")
// choose x1=1.1 x2=1.2 x3=1.4 x4 =1.5, compute x1*x2, x3*x4
// arbitrary test solution for random A matrix
X_soln[0] = 1.1
X_soln[1] = 1.2
X_soln[2] = 1.4
X_soln[3] = 1.5
print("solve system of equations A X = Y for X = ")
for i in 0..<n {
  print("\(Xname[i]) +")
} // end i
for i in n..<n+nlin {
  print("\(Xname[var1[i]]) ")
  if var2[i]>=0 {
    print("* \(Xname[var2[i]]) +")
  } // end if
  else {
    print(" +")
  } // end else
} // end i

for i in n..<n+nlin {
   X_soln[i] = X_soln[var1[i]] // at least one required
   if var2[i]>=0 { X_soln[i] *= X_soln[var2[i]] }
   if var3[i]>=0 { X_soln[i] *= X_soln[var3[i]] }
} // end i
// debug print
for i in 0..<n {
  print("X_soln[\(i)]=\(X_soln[i])")
} // end i
print(" ")

// set up Y
for i in 0..<n {
  Y[i] = 0.0
  for j in 0..<n+nlin {
    Y[i] += A[i][j]*X_soln[j]
  } // end j
} // end i
// debug print
for i in 0..<n {
  print("Y[\(i)]=\(Y[i])")
} // end i
print(" ")
print("initial guess")
for i in 0..<n+nlin {
  X1[i] = 1.0
  if i>=n {
    X1[i] = 0.0
  } // end if
  print("X1[\(i)]=\(X1[i])")
} // end i
print(" ")
print("var1=\(var1)")
print("var2=\(var2)")
print("var3=\(var3)")
print(" ")
print("calling simeq_newton3")
print("(n, nlin, A, Y, X1, var1, var2, var3, 10, 1.0, 1.0e-6, 0)")
X = simeq_newton3(n, nlin, A, Y, X1, var1, var2, var3, 10, 1.0, 1.0e-6, 0 )

print(" ")
print("Returned solution")

for i in 0..<n {
  print("X[\(i)]=\(X[i])   err=\(X[i]-X_soln[i])")
} // end i
print(" ")
print("first test finished \n")

print("test_simeq_newton3.swift finished")

func rnd() -> Double {  // random numbers 0.001 to 1.0
  var val = 0.0
  for _ in 0..<100 {
    val = Double(random())/1.0e10
    if val<1.0 && val>0.001 {
      return val
    } // end if
  } // end i
  return 0.01
} // end rnd
	
func simeq_newton3(_ n: Int, _ nlin: Int, _ A: [[Double]],
                   _ Y: [Double], _ X1: [Double],
		   _ var1: [Int], _ var2: [Int], _ var3: [Int],
		   _ uiter: Int, _ ub: Double, _ ueps: Double,
		   _ monitor: Int) -> [Double] { // X

  var eps = 1.0e-6 // possible default
  var b = 0.5 // stability convergence  factor  max 1.0
  var resid = 0.0  // residual from last iteration
  var presid = 0.0 // residual from prior to last iteration
  var maxiter = 10 // maximum number of iterations
  var Ja = [[Double]](repeating:[Double](repeating:0.0,count:n),count:n)
  var X       = [Double](repeating:0.0,count:n+nlin)
  var X_resid = [Double](repeating:0.0,count:n+nlin)
  var X_tmp2  = [Double](repeating:0.0,count:n+nlin)
  var X_next  = [Double](repeating:0.0,count:n+nlin)
  
  if monitor>0 { 
    print("simeq_newton3 running n=\(n), nlin=\(nlin)")
  } // end if
  eps = ueps
  maxiter = uiter
  b = ub
  // setup nonlinear terms
  for i in n..<n+nlin { // in nonlinear range
    X[i] = X[var1[i]] // must be valid
    if var2[i]>=0 {
      X[i] *= X[var2[i]] // -1 flag should not be used
    } // end if
    if var3[i]>=0 {
      X[i] *= X[var3[i]]
    } // end if
  } // end i
  for i in 0..<n+nlin {
    X[i] = X1[i]
    X_next[i] = X1[i]
  } // end i

  presid = 1.0e12
  // iterate to find solution
  for itr in 0..<maxiter {
    // compute residual
    resid = 0.0 
    for i in 0..<n {
      X_resid[i] = 0.0
      for j in 0..<n+nlin {
	X_resid[i] += A[i][j]*X[j]
      } // end j
      X_resid[i] -= Y[i] // X_resid used later if not solved
      resid = max(resid,abs(X_resid[i]))
    } // end i
    // debug print
    if monitor>2 {
      for i in 0..<n {
        print("residual X_resid[\(i)]=\(X_resid[i])")
      } // end i
      print(" ")
    } // end if

    // check for convergence
    if monitor>0 {
      print("simeq_newton3 itr \(itr), prev=\(presid) residual=\(resid)")
    } // end if
    if resid<eps {
      break
    } // end if
    // compute Jacobian
    for i in 0..<n { // equation
      for j in 0..<n { // variable
        Ja[i][j] = A[i][j] // each linear term
        for k in n..<n+nlin { // nonlinear
          // add derivative of nonlinear term
          if var1[k]==j && var2[k]==j && var3[k]==j {       /* X^3 */
            Ja[i][j] += A[i][k]*3.0*X[var2[k]]*X[var3[k]]   /* 3X^2 */
          } // end if
          else if var1[k]==j && var2[k]==j && var3[k]>=0 {  /* X^2 xa */
            Ja[i][j] += A[i][k]*2.0*X[var2[k]]*X[var3[k]]   /* 2X xa */
          } // end else if
          else if var1[k]==j && var2[k]==j && var3[k]<0 {    /* X^2 */
            Ja[i][j] += A[i][k]*2.0*X[var2[k]]               /* 2X */
          } // end else if
          else if var2[k]==j && var3[k]==j && var1[k]>=0 {   /* xa X^2 */
            Ja[i][j] += A[i][k]*2.0*X[var3[k]]*X[var1[k]]    /* 2X xa */
          } // end else if
          else if var2[k]==j && var3[k]==j && var1[k]<0 {    /*  X^2 */
            Ja[i][j] += A[i][k]*2.0*X[var3[k]]               /* 2X  */
          } // end else if
          else if var3[k]==j && var1[k]==j && var2[k]>=0 {   /* xa X^2 */
            Ja[i][j] += A[i][k]*2.0*X[var1[k]]*X[var2[k]]    /* 2X xa */
          } // end else if
          else if var3[k]==j && var1[k]==j && var2[k]<0 {    /*  X^2 */
            Ja[i][j] += A[i][k]*2.0*X[var1[k]]                /* 2X  */
          } // end else if
          else if var1[k]==j {
            if var2[k]>=0 && var3[k]>=0 { 
              Ja[i][j] += A[i][k]*X[var2[k]]*X[var3[k]]
            } // end if
            else if var2[k]>=0 {
	      Ja[i][j] += A[i][k]*X[var2[k]]
	    } // end else if
            else if var3[k]>=0 {
	      Ja[i][j] += A[i][k]*X[var3[k]]
	    } // end else if
	  } // end if var1
          else if var2[k]==j {
            if var1[k]>=0 && var3[k]>=0 { 
              Ja[i][j] += A[i][k]*X[var1[k]]*X[var3[k]]
	    } // end if
            else if var1[k]>=0 {
	      Ja[i][j] += A[i][k]*X[var1[k]]
	    } // end else if
            else if var3[k]>=0 {
	      Ja[i][j] += A[i][k]*X[var3[k]]
	    } // end else if
	  } // end if var2
          else if var3[k]==j {
            if var2[k]>=0 && var1[k]>=0 { 
              Ja[i][j] += A[i][k]*X[var2[k]]*X[var1[k]]
            } // end if
            else if var2[k]>=0 {
	      Ja[i][j] += A[i][k]*X[var2[k]]
	    } // else if
            else if var1[k]>=0 {
	      Ja[i][j] += A[i][k]*X[var1[k]]
	    } // end else if
          } // end if var3
        } // end k
      } // end j
    } // end i
    if monitor>4 {
      print("Ja computed ")
      for i in 0..<n {
        for j in 0..<n {
	  print("Ja[\(i)][\(j)]=\(Ja[i][j])")
	} // end j
      } // end i
      print(" ")
    } // end if

    // invert Ja
    Ja = matinv(Ja)
    if monitor>5 {
      print("Ja inverted ")
      for i in 0..<n {
        for j in 0..<n {
	  print("Ja[\(i)][\(j)]=\(Ja[i][j])")
	} // end j
      } // end i
      print(" ")
    } // end if

    // 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 i in 0..<n {
      X_tmp2[i] = 0.0
      for j in 0..<n {
	X_tmp2[i] += Ja[i][j]*X_resid[j] 
      } // end j
    } // end i
    // debug print
    if monitor>3 {
      for i in 0..<n {
        print("change X_tmp2[\(i)]=\(X_tmp2[i])")
      } // end i
      print(" ")
    } // end if
    for i in 0..<n {
      X_next[i] = X[i] - b*X_tmp2[i]
    } // end i
    for i in n..<n+nlin {
      X_next[i] = X_next[var1[i]] // must be valid
      if var2[i]>=0 {
        X_next[i] *= X_next[var2[i]]
      } // end if
      if var3[i]>=0 {
         X_next[i] *= X_next[var3[i]]
      } // end if
    } // end i

    if monitor>2 {
      for i in 0..<n+nlin {
        print("X_next[\(i)]=\(X_next[i])")
      } // end i
      print(" ")
    } // end if

    // iterate
    presid = resid 
    for i in 0..<n+nlin {
      X[i] = X_next[i]
    } // end i
    if monitor>2 {
      print(" ")
    } // end if

  } // end iteration    

  if monitor>0 {
    print("simeq_newton3.java finished")
  } // end if
  return X
} // end  simeq_newton3

// end simeq_newton3.swift


// matinv.swift  should have  import
func matinv(_ a: [[Double]]) -> [[Double]] {
  let n = a.count
  var inv = [[Double]](repeating:[Double](repeating:0.0,count:n),count:n)
  for i in 0..<n {
    for j in 0..<n {
      inv[i][j] = a[i][j]
    }
  }

  var row = [Int](repeating:0,count:n)
  var col = [Int](repeating:0,count:n)
  var temp = [Double](repeating:0,count:n)
  var hold = 0
  var I_pivot = 0
  var J_pivot = 0
  var pivot = 0.0
  var abs_pivot = 0.0

  // set up row and column interchange vectors
  for k in 0..<n {
    row[k] = k
    col[k] = k
  }
  // begin main reduction loop
  for k in 0..<n {
    // find largest element for pivot
    pivot = inv[row[k]][col[k]]
    I_pivot = k
    J_pivot = k
    for i in k..<n {
      for j in k..<n {
        abs_pivot = abs(pivot)
        if abs(inv[row[i]][col[j]]) > abs_pivot {
          I_pivot = i
          J_pivot = j
          pivot = inv[row[i]][col[j]]
        }
      }
    }
    if abs(pivot) < 1.0E-10 {
      print("Matrix is singular !")
      return [[Double]](repeating:[Double](repeating:0.0,count:n),count:n)
    }
    hold = row[k]
    row[k] = row[I_pivot]
    row[I_pivot] = hold
    hold = col[k]
    col[k] = col[J_pivot]
    col[J_pivot] = hold
    // reduce about pivot
    inv[row[k]][col[k]] = 1.0 / pivot
    for j in 0..<n {
      if j != k {
        inv[row[k]][col[j]] = inv[row[k]][col[j]] * inv[row[k]][col[k]]
      }
    }
    // inner reduction loop
    for i in 0..<n {
      if k != i {
        for j in 0..<n {
          if k != j {
            inv[row[i]][col[j]] = inv[row[i]][col[j]] - inv[row[i]][col[k]] * inv[row[k]][col[j]]
          }
        }
        inv[row[i]][col [k]] = -inv[row[i]][col[k]] * inv[row[k]][col[k]]
      }
    }
  }
  // end main reduction loop

  // unscramble rows
  for j in 0..<n {
    for i in 0..<n {
      temp[col[i]] = inv[row[i]][j]
    }
    for i in 0..<n {
      inv[i][j] = temp[i]
    }
  }
  // unscramble columns
  for i in 0..<n {
    for j in 0..<n {
      temp[row[j]] = inv[i][col[j]]
    }
    for j in 0..<n {
      inv[i][j] = temp[j]
    }
  }
  return inv
} // end matinv