! equation2_nl.f90  handle reciprocal terms
!  f1(x1, x2, x3) = a1*x1 + a2*x2^2 + a3*x3^3 + a4/x4 + a5/x5^2  = Y
!  f1'(x1) = a1
!  f1'(x2) = a2*2*x2
!  f1'(x3) = a3*3*x3^2
!  f1'(x4) = -a4/x4^2
!  f1'(x5) = -a5*2/x5^3
!
!  the other coefficients are shown in matrix form
!
! in matrix form, the equations are: A * X = Y
! | 5  4  3  2  1 |   |  x1  |   | 205.1488 |
! | 4  5  3  2  1 |   | x2^2 |   | 217.6888 |
! | 3  3  3  2  1 | * | x3^3 | = | 177.3088 |
! | 1  2  2  2  1 |   | 1/x4 |   | 113.5318 |
! | 2  3  1  1  2 |   |1/x5^2|   | 100.3626 |
!         A         *    Xt    =      Y
!
! The Y values are for specific unknown values of x1, ... ,x5
!
! in matrix form, the analytic Jacobian is:
!     |    1    |
!     |  2*x2   |
! A * |  3*x3^2 | =  Ja
!     | -1/x4^2 |
!     | -2/x5^3 |
! A *     D      =  Ja
!     deriv of X
!     wrt x1,x2,x3,x4,x4
!
! A numeric Jacobian can be approximated if the analytic is unknown.
!
! Newton iteration:
! x_next = x-fctr*(A*Xt-Y)/Ja,  /Ja is times transpose inverse of Ja
! x is individual variables, unknowns. Xt is equation terms
! fctr may need to be less than 1 for stability
! fctr may be adjusted by heuristics.
! Ja is, in general, dependent on all variables
!
! The goal is to zero the sum of the absolute values of
! all of the equations. This is the computed total error.
! Thus, solving a system of nonlinear equations has become
! a problem of finding roots.
!
! Because: if anything can go wrong then it will go wrong -
! Automatic control of the change factor, fctr, needs to be added:
! x_next = x - fctr*(A*Xt-Y)/Ja,  /Ja is times transpose inverse of Ja
! Start with fctr = 1.0 and reduce fctr by half if there is no
! improvement. Then, increase fctr by sqrt(2) if two iterations
! show improvement. (Other heuristics may be used)
!
! This equation has all unique 'terms' and thus all are included
! in the iterative computation.
! An equation with an  x1 term,  an  x2 term  and any terms involving
! these two, would have nx=2.  x1^2 term , x2^3 term , x1*x2 term
! are completely defined by x1 and x2.
! The Jacobian is based on each variable and all the terms in
! the equations to be solved.

! functions representing Y, given a set of x's and a's
function f1(A, X) result(y)
  double precision, dimension(5) :: A, X
  double precision :: y

  y = A(1)*X(1) + A(2)*X(2)*X(2) + A(3)*X(3)*X(3)*X(3) + &
      A(4)/X(4) + A(5)/(X(5)*X(5)) 
end function f1

program equation2_nl

  implicit none
  interface
    subroutine inverse(n, sz, A, AI)
      integer, intent(in) :: n  ! number of equations
      integer, intent(in) :: sz ! dimension of arrays
      double precision, dimension(sz,sz), intent(in) :: A
      double precision, dimension(sz,sz), intent(out) :: AI
    end subroutine inverse
    function f1(A, X) result(y)
      double precision, dimension(5) :: A, X
      double precision :: y
    end function f1
  end interface

  integer :: nitr = 20
  integer :: Hitlimit = 0
  integer :: Improve = 0
  integer :: Improve_Count = 0
  integer :: Fctr_Count = 0
  integer, parameter :: nx = 5 ! number of unique unknowns
  double precision, dimension(nx,nx) :: A
  data A / 5.0, 4.0, 3.0, 1.0, 2.0, & ! row major 
           4.0, 5.0, 3.0, 2.0, 3.0, &
           3.0, 3.0, 3.0, 2.0, 1.0, &
           2.0, 2.0, 2.0, 2.0, 1.0, &
           1.0, 1.0, 1.0, 1.0, 2.0 /
  double precision, dimension(nx) :: Arow   ! temporary row of A
  double precision, dimension(nx) :: Xt     ! X terms, A * Xt = Y
                                            ! with X being nonlinear terms
  double precision, dimension(nx) :: Y      ! RHS of equations
  double precision, dimension(nx) :: D      ! derivative of X terms
  double precision, dimension(nx,nx) :: Ja  ! Jacobian to be inverted
  double precision, dimension(nx,nx) :: JI  ! Jacobian  inverted
  double precision, dimension(nx) :: F      ! A*Xt-Y error to be multiplied
                                            ! by Ja inverse transpose
  double precision, dimension(nx) :: X      ! initial guess and solution
  double precision :: fctr = 0.5D0          ! can be unstable
  double precision :: Perror, Terror        ! previous, total error,
                                            ! sum of absolute values
  double precision, dimension(nx) :: S      ! desired solution
  data S / 5.1, 4.2, 3.3, 2.4, 1.5 /
  character(len=6), dimension(nx) :: Xts
  data Xts / '  x1  ', ' x2^2 ', ' x3^3 ', ' 1/x4 ', '1/x5^2' /
  integer :: i, j, k, itr

  print *, 'equation2_nl.f90 running '
  print *, 'solve system of nonlinear equations '
  do i=1,nx ! compute Y accurately, given A and solution S
    do j=1,nx
      Arow(j) = A(i,j)
    end do
    Y(i) = f1(Arow,S)
  end do
  do i=1,nx
    print '("|",5f8.4," | | ",A6," | | ",f9.4," |")', &
          (A(i,j),j=1,nx), Xts(i),Y(i)
  end do
  print *, '                     A                    *    Xt    =     Y '
  print *, 'guess initial X(1,5) '
  print *, 'compute Xt = | x1  x2^2  x3^3  1/x4  1/x5^2 | '
  print *, 'compute derivative D = | 1  2*x2 3*x3^2 -1/x4^2 -2/x5^3  | '
  print *, 'compute the Jacobian Ja = A * D and invert into JI '
  print *, 'iterate X_next = X - fctr*(A*Xt-Y)*Ja^-1  inverse transpose'
  print *, 'no guarantee of solution or unique solution '
  print *, 'sum absolute value A*Xt-Y should go to zero '
  print *, 'initial fctr, for stability = ',Fctr
  print *, ' '

  do i=1,nx
    X(i) = 1.0 ! S(i)+0.1 ! variable initial guess
  end do

  do itr=1,nitr
    do i=1,nx
      print '("X(",I1,")=",f11.7)', i, X(i)
    end do
    Print *, ' '
    Xt(1) = X(1) ! terms based on this iteration
    Xt(2) = X(2)*X(2)
    Xt(3) = X(3)*X(3)*X(3)
    Xt(4) = 1.0/X(4)
    Xt(5) = 1.0/(X(5)*X(5))

    D(1) = 1.0 ! derivatives based on this iteration
    D(2) = 2.0*X(2)
    D(3) = 3.0*X(3)*X(3)
    D(4) = -1.0/(X(4)*X(4))
    D(5) = -2.0/(X(5)*X(5)*X(5))

    do i=1,nx
      F(i) = 0.0
      do j=1,nx
        F(i) = F(i) + A(i,j)*Xt(j)
        Ja(i,j) = A(i,j)*D(j) ! build Jacobian analytically
      end do
      F(i) = F(i) - Y(i) ! residual  A*Xt-Y for row i
    end do
    Terror = 0.0
    do i=1,nx
      Terror = Terror + abs(F(i))
    end do

    if(itr.eq.1) then
      Perror = Terror
    else
      if(Terror.lt.Perror) then
        Improve = Improve + 1
        Perror = Terror
        if(Improve.gt.0) then
          Improve_Count = Improve_Count + 1
          Fctr = Fctr*1.1
          if(Fctr > 1.0) Fctr = 1.0 
        end if
      else
        Fctr_Count = Fctr_Count + 1
        Fctr = Fctr*0.9
        Improve = 0
      end if
    end if
    print *, 'iteration ',itr,', total error=',Terror
    Print *, ' '

    call inverse(nx, nx, Ja, Ji)

    ! compute X for next iteration
    do j=1,nx  ! loop through variables
      do i=1,nx  ! loop through equations
        X(j) = X(j) - fctr*F(i)*JI(j,i) ! in place update
      end do
    end do
  end do ! end iteration

  print *, 'final results             should be'
  do i=1,nx
    print '("X(",I1,")=",f11.7,"      ",f11.7)', i, X(i), S(i)
  end do
  Terror = 0.0
  do i=1,nx
    F(i) = 0.0
    do j=1,nx
      F(i) = F(i) + A(i,j)*Xt(j)
    end do
    F(i) = F(i) - Y(i) ! A*X-Y for row i
    Terror = Terror + abs(F(i))
  end do
  if(Fctr_Count.gt.0) then
    print *, 'increased error reduced Fctr ',Fctr_Count,' times'
  end if
  if(Improve_Count.gt.0) then
    print *, 'Fctr increased due to reduced error ',Improve_Count,' times'
  end if
  print *, 'final fctr=', Fctr
  print *, 'final total error=', Terror

  print *, 'equation2_nl.f90 finished '
end program ! equation2_nl