! equation_nl.f90  handle reciprocal terms
!  f1(x1, x2, x3) = x1 + 2 x2^2 + 3/x3 = 10
!  f1'(x1) = 1
!  f1'(x2) = 4 x2
!  f1'(x3) = -3/x3^2
!
!  f2(x1, x2, x3) = 2 x1 + x2^2 + 1/x3 = 6.333
!  f2'(x1) = 2
!  f2'(x2) = 2 x2
!  f2'(x3) = -1/x3^2
!
!  f3(x1, x2, x3) = 3 x1 + x2^2 + 4/x3 = 8.333
!  f3'(x1) = 3
!  f3'(x2) = 2 x2
!  f3'(x3) = -4/x3^2
!
! in matrix form, the equations are: A * X = Y
! | 1  2  3 |   |  x1  | = | 10.000 |
! | 2  1  1 | * | x2^2 | = |  6.333 |
! | 3  1  4 |   | 1/x3 | = |  8.333 |
!      A      *    Xt    =     Y
!
! in matrix form, the Jacobian is:
!     |   1    |   | 1   4*x2  -6/x3^2 |
! A * | 2*x2   | = | 2   2*x2  -2/x3^2 | = Ja
!     |-2/x3^2 |   | 3   2*x2  -8/x3^2 |
! A *     D      =        Ja
!     deriv of X
!     wrt x1,x2,x3
!
! Newton iteration:
! x_next = x-fctr*(A*Xt-Y)/Ja,  /Ja is times transpose inverse of Ja
! Ja is, in general, dependent on all variables

! three functions representing Y
function f1(x1, x2, x3) result(y)
  double precision, intent(in) :: x1, x2, x3
  double precision :: y
  y = x1 + 2.0*x2*x2 + 3.0/x3  ! = 10.0
end function f1

function f2(x1, x2, x3) result(y)
  double precision, intent(in) :: x1, x2, x3
  double precision :: y
  y = 2.0*x1 + x2*x2 + 1.0/x3 ! = 6.333
end function f2

function f3(x1, x2, x3) result(y)
  double precision, intent(in) :: x1, x2, x3
  double precision :: y
  y = 3.0*x1 + x2*x2 + 4.0/x3 ! = 8.333
end function f3

program equation_nl

  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
  end interface

  integer :: nitr = 7
  double precision, dimension(3,3) :: A
  data A / 1.0D0, 2.0D0, 3.0D0,  & ! need row major, OK, symmetrick
           2.0D0, 1.0D0, 1.0D0,  &
           3.0D0, 1.0D0, 4.0D0 /
  integer :: i, j, itr
  double precision, dimension(3) :: Xt ! X terms, A * X = Y  with X being terms
  double precision, dimension(3) :: Y ! RHS of equations
  double precision, dimension(3) :: D ! derivative of X terms
  double precision, dimension(3,3) :: Ja ! Jacobian to be inverted
  double precision, dimension(3,3) :: JI ! Jacobian  inverted
  double precision, dimension(3) :: F ! A*X-Y error to be multiplied by Ja
  double precision :: x1, x2, X3      ! initial guess and solution
  double precision :: fctr = 1.0      ! can be unstable

  Print *, 'equation_nl.f90 running '
  Print *, 'solve   x1 + 2 x2^2 + 3/x3 = 10 '
  Print *, '      2 x1 +   x2^2 + 1/x3 = 6.333 '
  Print *, '      3 x1 +   x2^2 + 4/x3 = 8.333 '
  Print *, ' | 1  2  3 |   |  x1  | = | 10.000 | '
  Print *, ' | 2  1  1 | * | x2^2 | = |  6.333 | '
  Print *, ' | 3  1  4 |   | 1/x3 | = |  8.333 | '
  Print *, '      A      *    Xt    =     Y      '
  Print *, 'guess initial x1, x2, x3 '
  Print *, 'compute Xt = | x1  x2^2  1/x3 | '
  Print *, 'compute derivative D = | 1  2*x2 -1/x3^2 | '
  Print *, 'compute the Jacobian Ja = A * D and invert '
  Print *, 'iterate x_next = x - fctr*(A*Xt-Y)*Ja^-1 '
  Print *, 'no guarantee of solution or unique solution '
  Print *, 'sum absolute value A*Xt-Y should go to zero '

  Y(1) = f1(1.0D0, 2.0D0, 3.0D0) ! desired solution 1, 2, 3
  Y(2) = f2(1.0D0, 2.0D0, 3.0D0) ! compute Y accurately
  Y(3) = f3(1.0D0, 2.0D0, 3.0D0)

  print '("| ",f7.4," ",f7.4," ",f7.4," | |  x1  | |",f7.4,"|")', &
        A(1,1), A(1,2), A(1,3), Y(1)
  print '("| ",f7.4," ",f7.4," ",f7.4," | |x2*x2 | |",f7.4,"|")', &
        A(2,1), A(2,2), A(2,3), Y(2)
  print '("| ",f7.4," ",f7.4," ",f7.4," | | 1/x3 | |",f7.4,"|")', &
        A(3,1), A(3,2), A(3,3), Y(3)
  print *, ' '

  x1 = 1.0D0 ! variable initial guess
  x2 = 1.0D0
  x3 = 1.0D0
  do itr=1,nitr
    print '("x1=",f8.5,", x2=",f8.5,", x3=",f8.5)',x1, x2, x3
    Xt(1) = x1 ! terms based on this iteration
    Xt(2) = x2*x2
    Xt(3) = 1.0/x3
    print '("Xt(1)=",f8.5,", Xt(2)=",f8.5,", Xt(3)=",f8.5)',Xt(1), Xt(2), Xt(3)
    D(1) = 1.0 ! derivatives based on this iteration
    D(2) = 2.0*x2
    D(3) = -1.0/(x3*x3)
    print '("D(1)=",f8.4,", D(2)=",f8.4,", D(3)=",f8.4)',D(1), D(2), D(3)
    do i=1,3
      F(i) = 0.0D0
      do j=1,3
        F(i) = F(i) + A(i,j)*Xt(j)
        Ja(i,j) = A(i,j)*D(j)
      end do
      F(i) = F(i) - Y(i) ! residual  A*X-Y for row i
    end do
    print '("F(1)=",f8.5,", F(2)=",f8.5,", F(3)=",f8.5)',F(1), F(2), F(3)
    Print *, 'iteration ',itr,', total error=',abs(F(1))+abs(F(2))+abs(F(3))
    print *, ' '
    do i=1,3
      print '("Ja(",i1,",1)=",f8.5,", Ja(",i1,",2)=",f8.5,", Ja(",i1",3)=",f8.5)', &
      i, Ja(i,1), i, Ja(i,2), i, Ja(i,3)
    end do
    call inverse(3, 3, Ja, JI)
    do i=1,3
      print '("JI(",i1,",1)=",f8.5,", JI(",i1,",2)=",f8.5,", JI(",i1",3)=",f8.5)', &
      i, Ja(i,1), i, Ja(i,2), i, Ja(i,3)
    end do
    do i=1,3
      x1 = x1 - fctr*F(i)*JI(1,i) ! transpose
      x2 = x2 - fctr*F(i)*JI(2,i) ! transpose
      x3 = x3 - fctr*F(i)*JI(3,i) ! transpose
    end do
  end do

  print *, ' '
  print '("final x1=",f8.4,", x2=",f8.4,", x3=",f8.4)',x1, x2, x3
  Xt(1) = x1 ! terms based on initial guess
  Xt(2) = x2*x2
  Xt(3) = 1.0/x3
  print '("terms Xt(1)=",f8.4,", Xt(2)=",f8.4,", Xt(3)=",f8.4)', &
        Xt(1), Xt(2), Xt(3)
  do i=1,3
    F(i) = 0.0D0
    do j=1,3
      F(i) = F(i) + A(i,j)*Xt(j)
    end do
    F(i) = F(i) - Y(i) ! A*X-Y for row i
  end do
  Print *, 'final total error=',abs(F(1))+abs(F(2))+abs(F(3))
  print *, 'equation_nl.f90 finished'
end program ! equation_nl