-- simeq_newton3.adb  solve nonlinear system of equations
--                    method: newton iteration using Jacobian
--
--  Given problem  A X = Y  where X may have terms x1, x2, x3, x4, ...
--                 and nonlinear such as: x1*x2, x1*x1*x3, x4*x4*x4, ...
--                 A sparse matrix may be used, coefficients given
--                 Y is vector of real, given
--                 independent unknowns are x1, ... , xn
--
--  Solve by initial guess at values of x1, x2, x3, x4 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
--
--  There may be no solution
--  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     x4     x1*x2 ... x3*x4
--
-- | A1,1   A1,2   A1,3   A1,4   A1,5 ...  A1,10 | | x1    | |y1 |
-- | A2,1   A2,2   A2,3   A2,4   A2,5 ...  A2,10 | | x2    | |y2 |
-- | A3,1   A3,2   A3,3   A3,4   A3,5 ...  A3,10 | | x3    | |y3 |
-- | A4,1   A4,2   A4,3   A4,4   A4,5 ...  A4,10 |*| x4    |=|y4 |
-- | A5,1   A5,2   A5,3   A5,4   A5,5 ...  A5,10 | | x1*x2 | |y5 |
-- | A6,1   A6,2   A6,3   A6,4   A6,5 ...  A6,10 | | x1*x3 | |y6 |
-- | ...                                         | |       | |   |
-- | A2,1   A2,2   A2,3   A2,4   A2,5 ...  A2,10 | | x3*x4 | |y10|
--
--  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,5*x2 + A1,6*x3 + A1,7*x4   deriv Row 1 wrt x1
--  J1,2 = A1,2 + A1,5*x1 + A1,8*x3 + A1,9*x4   deriv Row 1 wrt x2
--  J1,3 = A1,3 + A1,6*x1 + A1,8*x2 + A1,10*x4  deriv Row 1 wrt x3
--  J1,4 = A1,4 + A1,7*x1 + A1,9*x2 + A1,10*x3  deriv Row 1 wrt x4
--  J1,5 = A1,5                                 deriv Row 1 wrt x1, wrt x2
--  J1,6 = A1,6                                 deriv Row 1 wrt x1, wrt x3
--  J1,7 = A1,7                                 deriv Row 1 wrt x1, wrt x4
--  J1,8 = A1,8                                 deriv Row 1 wrt x2, wrt x3
--  J1,9 = A1,9                                 deriv Row 1 wrt x2, wrt x4
--  J1,10 = A1,10                               deriv Row 1 wrt x3, wrt x4
--
--  J2,1 = A2,1 + A2,5*x2 + A2,6*x3 + A2,7*x4   deriv Row 2 wrt x1
--  J2,2 = A2,2 + A2,5*x1 + A2,8*x3 + A2,9*x4   deriv Row 2 wrt x2
--
--  etc.
--
--  J1,4 ... J4,10  should last ones be needed? may not be independent
--
--  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.

with Ada.text_io;
use  Ada.text_io;
with real_arrays;
use  real_arrays;
with integer_arrays;
use  integer_arrays;
with inverse;

procedure simeq_newton3(n : integer; -- number of independent variables 1..n
                        nlin : integer; -- number of nonlinear terms
                        A : real_matrix; -- A(1..n, 1..n+nlin) coefficients
                        Y : real_vector; -- right hand side
                        v1 : integer_vector; -- independent 1..n, nonlinear
                        v2 : integer_vector; -- 0 for 1..n, nonlinear
                        v3 : integer_vector; -- 0 for 1..n, nonlinear or 0
                        X : in out Real_Vector; -- initial guess and solution
                        Maxitr : Integer := 6;  -- maximum iterations
                        Eps : Real := 1.0e-6;   -- convergence epsilon
                        Monitor : Integer := 6) -- output monitor
                                                is
  Ja : real_matrix(1..n,1..n);
  JaInv : real_matrix(1..n,1..n);
  X_prev : real_vector(1..n+nlin);
  X_resid : real_vector(1..n+nlin);
  X_tmp2 : real_vector(1..n+nlin);
  X_next : real_vector(1..n+nlin);
  var1 : integer_vector(1..n+nlin);
  var2 : integer_vector(1..n+nlin);
  var3 : integer_vector(1..n+nlin);

  b : real := 1.00;     -- stability factor
  resid : real;         -- residual from last iteration
  presid : real;        -- residual from prior to last iteration
  k, k1, k2, k3 : Integer;
  debug : boolean := false;
  badv : boolean := false;
  v_error : exception;


  package fltio is new float_io(real);
  use fltio;
  package intio is new integer_io(integer);
  use intio;

  procedure print_v is
  begin
    if not badv then
      badv := true;
      put("v1=(");
      for i in 1..n+nlin loop
        put(v1(i),2);
        if i = n+nlin then put(")"); else put(","); end if;
      end loop;
      new_line;
      put("v2=(");
      for i in 1..n+nlin loop
        put(v2(i),2);
        if i = n+nlin then put(")"); else put(","); end if;
      end loop;
      new_line;
      put("v3=(");
      for i in 1..n+nlin loop
        put(v3(i),2);
        if i = n+nlin then put(")"); else put(","); end if;
      end loop;
      new_line;
    end if; -- not badv
  end print_v;

  procedure errormsg(msg : string) is
  begin
     badv := true;
     put_line(Msg);
  end errormsg;


  procedure print_var is
  begin
    if debug or monitor>1 then
      put("var1=(");
      for i in 1..n+nlin loop
        put(var1(i),2);
        if i = n+nlin then put(")"); else put(","); end if;
      end loop;
      new_line;
      put("var2=(");
      for i in 1..n+nlin loop
        put(var2(i),2);
        if i = n+nlin then put(")"); else put(","); end if;
      end loop;
      new_line;
      put("var3=(");
      for i in 1..n+nlin loop
        put(var3(i),2);
        if i = n+nlin then put(")"); else put(","); end if;
      end loop;
      new_line;
    end if; -- debug
  end print_var;

  procedure check_input is
  begin
    if A'length(1) < n then
      errormsg("A first dimension small, minimum n");
    end if;
    if A'length(2) < n+nlin then
      errormsg("A second dimension small, minimum n+nlin");
    end if;
    if X'length < n+nlin then
      errormsg("X dimension small, minimum n+nlin");
    end if;
    if Y'length < n then
      errormsg("Y dimension small, minimum n");
    end if;
    if v1'length < n+nlin then
      errormsg("v1 dimension small, minimum n+nlin");
    end if;
    if v2'length < n+nlin then
      errormsg("v2 dimension small, minimum n+nlin");
    end if;
    if v3'length < n+nlin then
      errormsg("v3 dimension small, minimum n+nlin");
    end if;
    if badv then raise v_error; end if;
    for j in 1..n loop -- check and clean up v input
      X_prev(j) := X(j);
      var1(j) := v1(j);
      if var1(j) /= j then
        var1(j) := j;
        put("call error, v1(");
        put(j,2);
        put(") should be");
        put(j,2);
        new_Line;
        print_v;
      end if;
      var2(j) := v2(j);
      if var2(j) /= 0 then
        var2(j) := 0;
        put("call error, v2(");
        put(j,2);
        put_line(") should be 0");
        print_v;
      end if;
      var3(j) := v3(j);
      if var3(j) /= 0 then
        var3(j) := 0;
        put("call error, v3(");
        put(j,2);
        put_line(") should be 0");
        print_v;
      end if;
    end loop;
    for j in n+1..n+nlin loop
      var1(j) := v1(j);
      var2(j) := v2(j);
      var3(j) := v3(j);
      if var2(j) > var1(j) then -- sort biggest first
        k := var1(j);
        var1(j) := var2(j);
        var2(j) := k;
      end if;
      if var3(j) > var1(j) then
        k := var1(j);
        var1(j) := var3(j);
        var3(j) := k;
      end if;
      if var3(j) > var2(j) then
        k := var2(j);
        var2(j) := var3(j);
        var3(j) := k;
      end if;
      if var1(j)<1 or var1(j)>n then
        put("call error, v1(");
        put(j,2);
        put_line(") must be in (1..n)");
        print_v;
      end if;
      if var2(j)<1 or var2(j)>n then
        put("call error, v2(");
        put(j,2);
        put_line(") must be in (1..n)");
        print_v;
      end if;
      if var3(j)<0 or var3(j)>n then
        put("call error, v2(");
        put(j,2);
        put_line(") must be in (0..n)");
        print_v;
      end if;
      if badv then raise v_error; end if;
    end loop;
    print_var;
  end check_input;

begin
   if debug or monitor>0 then
      put_line("simeq_newton3.adb running");
   end if;
   check_input; -- also set X_prev and var1, var2, var3

  for j in n+1..n+nlin loop -- set non linear X_prev
    X_prev(j) := X_prev(var1(j)) * X_prev(var2(j));
    -- var1 must not be empty, var2 must not be empty for j>n
    if(var3(j) /= 0) then
      X_prev(j) := X_prev(j)* X_prev(var3(j));
    end if;
  end loop;

  if debug or monitor>4 then
    for i in 1..n+nlin loop
      put("X_prev(");
      put(i,2);
      put(")=");
      put(X_prev(i),3,6);
      new_line;
    end loop;
    new_line;
  end if; -- debug

  presid := 1.0e12;
  -- iterate to find solution
  for itr in 1..maxitr loop -- could be more iterations
    -- compute residual
    resid := 0.0;
    for i in 1..n loop
      X_resid(i) := 0.0;
      for j in 1..n+nlin loop
        X_resid(i) := X_resid(i) + A(i,j)*X_prev(j);
      end loop; -- j
      X_resid(i) := X_resid(i) - Y(i); -- X_resid used later if not solved
      resid := resid + abs(X_resid(i));
    end loop; -- i
    -- debug/monitor print
    if Debug or Monitor>1 then
      for i in 1..n loop
        put("residual X_resid(");
        put(i,2);
        put(")= ");
        put(X_resid(i),3,6);
        new_line;
      end loop; -- i
      new_line;
    end if;
    if Debug or Monitor>0 then
      put("itr ");
      put(itr,2);
      put(", prev= ");
      put(presid);
      put(", residual= ");
      put(resid);
      new_line;
    end if;
    -- check for convergence
    if resid<eps then exit; end if;

    -- compute derivatives to load Ja
    -- compute Jacobian
    for i in 1..n loop
      for j in 1..n loop
        Ja(i,j) := A(i,j); -- always have linear terms
      end loop; -- j
    end loop; -- i

    -- add values from nonlinear terms
    for i in 1..n loop
      for j in 1..n loop
        for k in n+1..n+nlin loop
          k1 := var1(k); -- sorted largest first
          k2 := var2(k); -- sorted largest first
          k3 := var3(k); -- can be zero
          if k1=j and k2=j and k3=j then -- cube
            Ja(i,j) := Ja(i,j) + 3.0*A(i,k)*X_prev(k2)*X_prev(k3);
          elsif k1=j and k2=j and k3=0 then -- square
            Ja(i,j) := Ja(i,j) + 2.0*A(i,k)*X_prev(k2);
          elsif k2=j and k3=j then -- square k1/=0
            Ja(i,j) := Ja(i,j) + 2.0*A(i,k)*X_prev(k3)*X_prev(k1);
          elsif k1=j and k3/=0 then
            Ja(i,j) := Ja(i,j) + A(i,k)*X_prev(k2)*X_prev(k3);
          elsif k1=j and k3=0 then
            Ja(i,j) := Ja(i,j) + A(i,k)*X_prev(k2);
          elsif k2=j and k3/=0 then
            Ja(i,j) := Ja(i,j) + A(i,k)*X_prev(k1)*X_prev(k3);
          elsif k2=j and k3=0 then
            Ja(i,j) := Ja(i,j) + A(i,k)*X_prev(k1);
          elsif k3=j then
            Ja(i,j) := Ja(i,j) + A(i,k)*X_prev(k1)*X_prev(k2);
          end if;
        end loop; -- k
      end loop; -- j
    end loop; -- i

    -- debug print
    if debug or monitor>4 then
      for i in 1..n loop
        for j in 1..n loop
          put("Ja(");
          put(i,2);
          put(",");
          put(J,2);
          put(")=");
          put(Ja(i,j),3,6);
          new_line;
        end loop; -- j
      end loop; -- i
      new_line;
    end if;
    -- invert Ja
    JaInv := inverse(Ja);
    for i in 1..n loop
      for j in 1..n loop
        Ja(i,j) := JaInv(i,j);
        -- debug print
        if debug or monitor>5 then
          put("JaInv(");
          put(i,2);
          put(",");
          put(J,2);
          put(")=");
          put(JaInv(i,j),3,6);
          new_line;
        end if;
      end loop; -- j
    end loop; -- i

    -- X_next = X_prev - (J_prev^-1 * (A X_prev - Y))*b
    -- X_next = X_prev - (J_prev^-1 * (X_resid))*b
    -- X_next = X_prev - (X_tmp2)*b

    for i in 1..n loop
      X_tmp2(i) := 0.0;
      for j in 1..n loop
        X_tmp2(i) := X_tmp2(i) + JaInv(i,j)*X_resid(j);
      end loop; -- j
    end loop; -- i
    -- debug print
    if debug or monitor>5 then
      for i in 1..n loop
        put("change X_tmp2(");
        put(i,2);
        put(")= ");
        put(X_tmp2(i),3,6);
        new_line;
      end loop; -- i
      new_line;
    end if; -- debug
    for i in 1..n loop
      X_next(i) := X_prev(i) - b*X_tmp2(i);
    end loop; -- i
    -- compute non linear terms
    for j in n+1..n+nlin loop -- set non linear X_prev
      X_next(j) := X_next(var1(j)) * X_next(var2(j));
      -- var1 must not be empty, var2 must not be empty for j>n
      if(var3(j) /= 0) then
        X_next(j) := X_next(j)* X_next(var3(j));
      end if;
    end loop;

    if debug or monitor>2 then
      for i in 1..n+nlin loop
        put("X_next(");
        put(i,2);
        put(")= ");
        put(X_next(i),3,6);
        new_line;
      end loop; -- i
      new_line;
    end if;
    -- iterate
    presid := resid;
    for i in 1..n+nlin loop
      X_prev(i) := X_next(i);
    end loop; -- i

    new_line;
  end loop; -- itr    end iteration

  for i in 1..n+nlin loop
    X(i) := X_next(i);
  end loop; -- i
  if resid<eps then
    put_line("simeq_newton3 found solution");
  else
    put_line("simeq_newton3 did not converge <"&long_float'image(Eps));
  end if;
exception
  when v_error =>
    put_line("simeq_newton3 had errors, X unchanged");
end simeq_newton3;