-- equation2_nl.adb  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)
--
-- Due to the possibility of multiple solutions and just plain wild
-- oscillations of x values, bound the minimum and maximum of
-- all variables when any information is known.
--
-- 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.

with Ada.Text_IO;
use  Ada.Text_IO;
with Real_Arrays;
use  Real_Arrays;
with Inverse;
with Rderiv;

procedure equation2_nl is

  -- functions representing Y, given a set of x's and a's
  function f1(A : Real_Vector; X : Real_Vector) return Real is
  begin
     return 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 f1;

  type Funct is access function (A : Real_Vector; X : Real_vector) return Real;

  procedure Build_Jacobian(A : in Real_Matrix; X : in Real_Vector;
                           F : Funct; Jb : out Real_Matrix) is
    nx : Integer := A'last(1);
    np : Integer := 7; -- odd number 5, 7, 9
    Cx : Real_Vector(1..np);
    H : Real := 0.1;
    Xd : Real_Vector(1..np) := (others => 0.0);
    Hx : Real_Vector(1..np) := (others => 0.0);
    Deriv : Real := 0.0;
    Arow : Real_Vector(1..nx) := (others => 1.0);
  begin
    if X'Last /= A'Last(2) or A'Last(1) /= A'Last(2) then
      Put_Line("dimension error in call to Build_Jacobian");
    end if;
    Rderiv(1, np, (np+1)/2, H, Cx);
    for i in 1..np loop
      Put_Line("Cx("&Integer'Image(i)&")="&Real'Image(Cx(i)));
    end loop;
    Hx(1) := -Real((np-1)/2)*H;
    for i in 2..np loop
       Hx(i) := Hx(i-1)+H;
    end loop;
    for i in 1..np loop
      Put_Line("Hx("&Integer'Image(i)&")="&Real'Image(Hx(i)));
    end loop;

    for i in A'Range(1) loop -- outer loop to Jb
      for k in A'Range(2) loop
        Arow(k) := A(i,k); -- i_th equation
      end loop;
      for j in A'Range(2) loop -- inner loop to Jb
        Jb(i,j) := 0.0;
        for k in X'range loop
          Xd(k) := X(k);
        end loop;
        for k in 1..Np loop
          Xd(j) := X(j) + Hx(k); -- derivative wrt X(j)
          Jb(i,j) := Jb(i,j) + F(Arow,Xd)*Cx(k);
        end loop;
      end loop;
    end loop;
  end Build_Jacobian;

  package Real_IO is new Float_IO(Real);
  use Real_IO;

  nitr : Integer := 20;
  Iflimit : Boolean := false; -- apply Xmin and Xmax constraints
  Hitlimit : Boolean := False;
  Improve, Improve_Count, Limit_Count, Fctr_Count : Integer := 0;
  nx : Integer := 5; -- number of unique unknowns
  A : Real_Matrix(1..nx,1..nx) := ((5.0, 4.0, 3.0, 2.0, 1.0),
                                   (4.0, 5.0, 3.0, 2.0, 1.0),
                                   (3.0, 3.0, 3.0, 2.0, 1.0),
                                   (1.0, 2.0, 2.0, 2.0, 1.0),
                                   (2.0, 3.0, 1.0, 1.0, 2.0));
  Arow : Real_Vector(1..nx);
  Xt : Real_Vector(1..nx); -- X terms, A * Xt = Y  with X being nonlinear terms
  Y : Real_Vector(1..nx); -- RHS of equations
  D : Real_Vector(1..nx); -- derivative of X terms
  Ja : Real_Matrix(1..nx,1..nx); -- Jacobian to be inverted
  Jb : Real_Matrix(1..nx,1..nx); -- Jacobian, alternate build
  JI : Real_Matrix(1..nx,1..nx); -- Jacobian  inverted
  F : Real_Vector(1..nx);     -- A*X-Y error to be multiplied by Ja inverse
  X : Real_Vector(1..nx);     -- initial guess and solution
  Xd : Real_Vector(1..nx) := (others => 0.0); -- delta to subtract
                                              -- from X for next iteration
  fctr : Real := 0.5; -- can be unstable
  Perror, Terror : Real;      -- previous, total error, sum of absolute values
  Tmp, Sum : Real;
  S : Real_Vector(1..nx) := (5.1, 4.2, 3.3, 2.4, 1.5); -- desired solution
  Xts : array(1..nx) of String(1..6) := ("  x1  ",
                                         " x2^2 ",
                                         " x3^3 ",
                                         " 1/x4 ",
                                         "1/x5^2");
  Xmax : Real_Vector(1..Nx) := (9.0, 9.0, 9.0, 9.0, 9.0); -- maximum value,
                                                          -- optional
  Xmin : Real_Vector(1..Nx) := (0.1, 0.1, 0.1, 0.1, 0.1); -- minimum value,
                                                          -- optional

begin

  Put_Line("equation2_nl.adb running ");
  Put_Line("solve system of nonlinear equations ");
  for i in 1..nx loop -- compute Y accurately, given A and solution S
    for j in 1..nx loop
      Arow(j) := A(i,j);
    end loop;
    Y(i) := f1(Arow,S);
  end loop;
  for i in 1..nx loop
    Put("|");
    for j in 1..nx loop
      Put(A(i,j),3,4,0);
    end loop;
    Put("|   |"&Xts(i)&"|   | "); Put(Y(i),4,4,0); Put_Line("|");
  end loop;
  Put_Line("                     A                     *    Xt    =     Y ");
  Put_Line("guess initial X(1..5) ");
  Put_Line("compute Xt = | x1  x2^2  x3^3  1/x4  1/x5^2 | ");
  Put_Line("compute derivative D = | 1  2*x2 3*x3^2 -1/x4^2 -2/x5^3  | ");
  Put_Line("compute the Jacobian Ja = A * D and invert into Jb ");
  Put_Line("iterate X_next = X - fctr*(A*Xt-Y)*Ja^-1  inverse transpose");
  Put_Line("no guarantee of solution or unique solution ");
  Put_Line("sum absolute value A*Xt-Y should go to zero ");
  Put_Line("initial fctr, for stability = "&Real'Image(Fctr));
  New_Line;

  for i in 1..nx loop
    X(i) := 1.0; -- S(i)+0.1; -- variable initial guess
  end loop;

  Build_Jacobian(A, X, F1'access, Jb); -- numerically, a test

  for itr in 1..nitr loop
    for i in 1..nx loop
      Put("X("&Integer'Image(i)&")="); Put(X(i),3,7,0);
      Put("     -Xd="); Put(-Xd(i),3,7,0); New_Line;
    end loop;
    New_Line;
    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));
    if Itr = 1 then
      for i in 1..nx loop
        Put("Xt("&Integer'Image(i)&")="); Put(Xt(i),3,5,0); New_Line;
      end loop;
      New_Line;
    end if;
    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));
    if Itr = 1 then
      for i in 1..nx loop
        Put("D("&Integer'Image(I)&")="); Put(D(i),3,5,0); New_Line;
      end loop;
      New_Line;
    end if;
    for i in 1..nx loop
      F(i) := 0.0;
      for j in 1..nx loop
        F(i) := F(i) + A(i,j)*Xt(j);
        Ja(i,j) := A(i,j)*D(j); -- build Jacobian analytically
      end loop;
      F(i) := F(i) - Y(i); -- residual  A*Xt-Y for row i
    end loop;
    Terror := 0.0;
    for i in 1..nx loop
      Terror := Terror + abs(F(i));
    end loop;
    if Itr = 1 then
      for i in 1..nx loop
        Put("F("&Integer'Image(i)&")="); Put(F(i),3,5,0); New_Line;
      end loop;
      New_Line;
    end if;
    if itr = 1 then
      Perror := Terror;
    else
      if Terror < Perror then
        Improve := Improve + 1;
        Perror := Terror;
        if Improve>0 then
          Improve_Count := Improve_Count + 1;
          Fctr := Fctr*1.1;
          if Fctr > 1.0 then Fctr := 1.0; end if;
          Put_Line("Fctr improved="&Real'Image(Fctr));
        end if;
      else
        Fctr_Count := Fctr_Count + 1;
        Fctr := Fctr*0.9;
        Improve := 0;
        Put_Line("Terror Fctr reduced="&Real'Image(Fctr));
      end if;
    end if;
    Put_Line("iteration "&Integer'Image(itr)&", total error="&
             Real'Image(Terror));
    New_Line;
    if Itr = 1 then
      for i in 1..nx loop
        for j in 1..nx loop
          put("Ja("&Integer'Image(i)&","&Integer'Image(j)&")=");
          Put(Ja(i,j),3,5,0); Put("   Jb="); Put(Jb(i,j),3,5,0); New_Line;
        end loop;
      end loop;
      New_Line;
    end if;
    JI := inverse(Ja);
    if Itr = 1 then
      for i in 1..nx loop
        for j in 1..nx loop
          put("JI("&Integer'Image(i)&","&Integer'Image(j)&")=");
          Put(JI(i,j),3,5,0); New_Line;
        end loop;
      end loop;
      New_Line;
      Sum := 0.0;
      for i in 1..nx loop
        for j in 1..nx loop
          Tmp := 0.0;
          for K in 1..nx loop
            Tmp := Tmp + Ja(i,k)*Ji(k,j);
          end loop;
          if i = j then Tmp := Tmp -1.0; end if;
          Sum := Sum + abs(Tmp);
        end loop;
      end loop;
      Put_Line("inversion error="&Real'Image(Sum));
      New_Line;
    end if; -- Itr=1

    -- compute X for next iteration
    Hitlimit := True;
    while Hitlimit loop
      for i in 1..nx loop -- loop through variables
      Xd(I) := 0.0;
        for j in 1..nx loop -- loop through equations
          Xd(i) := Xd(i) + fctr*F(j)*JI(i,j); -- transpose?
        end loop;
      end loop;
      if Itr = 1 then
        for i in 1..nx loop
          Put("Xd("&Integer'Image(i)&")="); Put(Xd(i),3,5,0); New_Line;
        end loop;
      end if;
      for i in 1..nx loop
        X(I) := X(I)-Xd(I);
      end loop;

      Hitlimit := false;
      if Iflimit then
        for i in 1..nx loop
          if X(i)>Xmax(i) then Hitlimit := true; exit; end if;
          if X(i)<Xmin(i) then Hitlimit := true; exit; end if;
        end loop;
      end if; -- limit
      if Hitlimit then
        Limit_Count := Limit_Count + 1;
        Fctr := Fctr*0.9;
        Put_Line("limit Fctr reduced="&Real'Image(Fctr));
      end if;
    end loop; -- while
  end loop; -- iteration

  New_Line;
  Put_Line("final results             should be");
  for i in 1..nx loop
     Put("X("&Integer'Image(i)&")="); Put(X(i),3,7,0);
     Put("      "); Put(S(I),3,7,0); New_Line;
  end loop;
  Terror := 0.0;
  for i in 1..nx loop
    F(i) := 0.0;
    for j in 1..nx loop
      F(i) := F(i) + A(i,j)*Xt(j);
    end loop;
    F(i) := F(i) - Y(i); -- A*X-Y for row i
    Terror := Terror + abs(F(i));
  end loop;
  if Fctr_Count>0 then
    Put_Line("increased error reduced Fctr "&
             Integer'Image(Fctr_Count)&" times");
  end if;
  if Improve_Count>0 then
    Put_Line("Fctr increased due to reduced error "&
             Integer'Image(Improve_Count)&" times");
  end if;
  if Limit_Count>0 then
    Put_Line("iterations hit a max/min limit "&
             Integer'Image(Limit_Count)&" times");
  end if;
  Put_Line("final fctr="&Real'Image(Fctr));
  Put_Line("final total error="&Real'Image(Terror));

  Put_Line("equation2_nl.adb finished ");
end equation2_nl;