-- pde_nl21.adb  non linear PDE, second order, second degree, one dimension
--
-- The PDE to be solved for u(x) is:
-- Uxx(x) + 2*Ux(x) + U(x)^2  = f(x)
--
-- for testing code:
--        f(x) = x^4 + 2x^2 + 4x + 3
--        u(xmin=0)=1  u(xmax=1)=2
-- analytic solution and boundary  ub(x) = x^2+1
--        thus  Ux(x) = 2*x
--              Uxx(x) = 2
--

with Ada.Text_IO;
use  Ada.Text_IO;
with Ada.Numerics.Long_Elementary_Functions;
use  Ada.Numerics.Long_Elementary_Functions;
with Real_Arrays; -- types real, real_vector, real_matrix
use  Real_Arrays;
with Integer_Arrays; -- types integer_vector
use  Integer_Arrays;
with deriv;
with rderiv;
with Simeq_Newton5;

procedure pde_nl21 is
  nx : Integer := 5;      -- number of coefficients for discrete derivative
  dof : Integer := nx-2;  -- number of equations and number of linear unknowns
  nlin : Integer := nx-2; -- number of nonlinear terms, specific this PDE, nx-2
  ntot : Integer := dof+nlin; -- total number of terms, includes nonlinear
  A  : Real_Matrix(1..dof,1..ntot); -- nonlinear system of equations
  xg : Real_Vector(1..nx);      -- x grid
  ug : Real_Vector(1..nx);      -- computed solution at x grid, boundary needed
  Y  : Real_Vector(1..dof);     -- RHS of equations
  Us : Real_Vector(1..ntot);    -- solution being computed, not include bound
  var1 : Integer_Vector(1..ntot);  -- for nonlinear simeq
  var2 : Integer_Vector(1..ntot);  -- for nonlinear simeq
  var3 : Integer_Vector(1..ntot);  -- for nonlinear simeq
  vari1 : Integer_Vector(1..ntot); -- for nonlinear simeq
  vari2 : Integer_Vector(1..ntot); -- for nonlinear simeq

  Ua : Real_Vector(1..nx);    -- known includes boundary
  cx  : Real_Vector(1..nx);   -- numerical first derivative
  cxx : Real_Vector(1..nx);   -- numerical second derivative
  hx : real;
  err, avgerr, Maxerr, x : real;
  eps : real := 1.0e-4;
  xmin : Real := 0.0;
  xmax : Real := 1.0;

  -- layout
  -- xmin   <--  dof       -->   xmax
  -- xg(1)  xg(2)  xg(3)  xg(4)  xg(nx)  x grid (generally cx, cxx etc)
  -- Ua(1)  Ua(2)  Ua(3)  Ua(4)  Ua(nx)  known solution for testing
  -- U (1)  U (2)  U (3)  U (4)  U (nx)  being computed
  --        Us(1)  Us(2)  Us(3)          ... from nonlinear simeq

  -- definition of PDE to be solved
  function f(x:real) return real is -- forcing function
  begin
    return x*x*x*x + 2.0*x*x + 4.0*x + 3.0;
  end f;

  function ub(x:real) return real is -- analytic solution and boundary
  begin
    return x*x+1.0;
  end ub;

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

  package intio is new integer_io(integer);
  use intio;

  procedure Build_Vars is -- based on PDE
  begin
    -- var1()  = ( 1, 2, 3, 1, 2, 3);
    -- var2()  = ( 0, 0, 0, 1, 2, 3);
    -- var3()  = ( 0, 0, 0, 0, 0, 0);
    -- vari1() = ( 0, 0, 0, 0, 0, 0);
    -- vari2() = ( 0, 0, 0, 0, 0, 0);
    -- j index     1  2  3  4  5  6
    for i in 1..dof loop
      var1(i)  := i; -- linear terms, required
      var2(i)  := 0;
      var3(i)  := 0;
      vari1(i) := 0;
      vari2(i) := 0;
    end loop;
    for i in Dof+1..ntot loop
      -- set var1, var2 for nonlinear terms
      var1(i)  := i-dof; -- x1^2  x2^2  x3^2
      var2(i)  := i-dof;
      var3(i)  := 0;
      vari1(i) := 0;
      vari2(i) := 0;
    end loop;
  end Build_Vars;

  procedure Build_Equations is
  begin
    -- initialize A and Y
    for i in 1..dof loop
      for j in 1..ntot loop
        A(i,j) := 0.0;
      end loop;
      Y(i) := 0.0;
    end loop;

    -- compute entries in nonlinear equation matrix A*Us=Y
    -- method
    --  u' + 2 u' + u*u = f(x)  becomes system of non linear equations
    --  at xg(2), xg(3), ... , xg(nx-1)
    --  U(2), U(3) and U(4) are linear DOF, U(1) and U(nx) are boundary
    --  U(2)^2, U(3)^2, U(4)^2 are nonlinear terms
    --  renumbered for Us nonlinear system of quations solver:
    --  Us(1), Us(2) and Us(3) are linear terms  1..dof
    --  Us(4), Us(5) and Us(6) are squared nonlinear terms dof+1..ntot
    --
    -- discretized equations become:
    -- Y(1) := f(x2)
    -- with cx and cxx for point 2
    -- cxx(1)*U(1)+cxx(2)*Us(1)+cxx(3)*Us(2)+cxx(4)*Us(3)+cxx(5)*U(5)+
    --    2*cx(1)*U(1)+2*cx(2)*Us(1)+2*cx(3)*Us(2)+2*cx(4)*Us(3)+2*cx(5)*U(5)+
    --         Us(1)*Us(1) + Us(2)*Us(2) + Us(3)*Us(3) = Y(1) = f(x2)
    -- except, U(1) and U(5) are known boundary values, thus subtract for Y(1)
    -- Y(1) := Y(1) - cxx(1)*U(x1) - cxx(5)*U(x5)
    --
    -- The linear and non linear unknown terms are:
    -- Us(x2)=Us(1)        Us(x3)=Us(2)        Us(x4)=Us(3)
    -- Us(x2)*Us(x2)       Us(x3)*Us(x3)       Us(x4)*Us(x4)
    --
    -- Thus A*Us=Y
    -- A(1,1) := cxx(2) + 2*cx(2)  using cx and cxx at point 2
    -- A(1,2) := cxx(3) + 2*cx(3)
    -- A(1,3) := cxx(4) + 2*cx(3)
    -- A(1,4) = 1.0  just Us^2 terms
    -- similar for A(2,*) and A(3,*) using cx and cxx at points 3 and 4

    Put_line("compute non linear A matrix");
    for ia in 1..dof loop -- xg and point are one greater
      x := xg(ia+1);
      Y(ia) := f(x); -- RHS
      rderiv(1, nx, ia+1, hx, cx); -- discrete derivative coefficients
      rderiv(2, nx, ia+1, hx, cxx);

      -- term  1 * Uxx
      for j in 1..nx loop
        if j=1 or j=nx then
          Y(ia) := Y(ia) - cxx(j)*ug(j);
        else
          A(ia,j-1) := A(ia,j-1) + cxx(j);
        end if;
      end loop;

      -- term  2 * Ux
      for j in 1..nx loop
        if j=1 or j=nx then
          Y(ia) := Y(ia) - 2.0*cx(j)*ug(j);
        else
          A(ia,j-1) := A(ia,j-1) + 2.0*cx(j);
        end if;
      end loop;

      -- linear section done
      -- nonlinear section
      -- term  U^2
      A(ia,dof+ia) := A(ia,dof+ia) + 1.0; -- just U at this x
    end loop; -- ia
  end Build_Equations;

  procedure Print_AY is
  begin
    Put_Line("nonlinear matrix A");
    for i in 1..dof loop
      for j in 1..ntot loop
        Put_Line("i="&Integer'Image(I)&", j="&Integer'Image(J)&", A(i,j)="&
                  Real'Image(A(i,j)));
      end loop;
    end loop;
    New_Line;
    Put_Line("Y computed RHS ");
    for i in 1..dof loop
      Put_Line("Y("&Integer'Image(i)&")="&Real'Image(Y(i)));
    end loop;
    New_Line;
  end Print_AY;

  procedure check_soln(Ux : Real_Vector) is -- check when solution unknown
    -- u'' + 2 u' + u^2 = f(x) check for each xg inside boundary
    xi, err, Smaxerr : real;
  begin
    Put_line("check_soln");
    smaxerr := 0.0;
    for i in 2..nx-1 loop -- xg index
      xi := xg(i);
      rderiv(1, nx, i, hx, cx);
      rderiv(2, nx, i, hx, cxx);
      err := Ux(i)*Ux(i)-f(xi); -- term and RHS
      for j in 1..nx loop
        err := err + cxx(j)*Ux(j); -- u'' term
        err := err + 2.0*cx(j)*Ux(j); -- 2*u' term
      end loop; -- j
      if abs(err)>smaxerr then smaxerr := abs(err); end if;
      Put_line("i="&Integer'Image(I)&", err="&Real'Image(err));
    end loop; -- i
    Put_line("check_soln max error="&Real'Image(smaxerr));
  end Check_Soln;

  procedure Print_Equations(n : integer; nlin : integer;
                            var1 : integer_vector; var2 : integer_vector;
                            var3 : integer_vector; vari1 : integer_vector;
                            vari2 : Integer_Vector) is
    Xnames : array(1..5) of String(1..2) := ("X1", "X2", "X3", "X4", "X5");
  begin
    Put_Line("system of equations to be solved, i=1.."&Integer'Image(n));
    for j in 1..n loop
      Put_line("A(i,"&Integer'Image(j)&")*"&Xnames(j)&"+");
    end loop;
    for j in n+1..n+nlin loop
      Put("A(i,"&Integer'Image(J)&")");
      if Var1(J)>0 then
         Put("*"&Xnames(Var1(J)));
      end if;
      if Var2(J)>0 then
         Put("*"&Xnames(Var2(J)));
      end if;
      if Var3(J)>0 then
         Put("*"&Xnames(Var3(J)));
      end if;
      if Vari1(J)>0  then
         Put("/("&Xnames(Vari1(J)));
      end if;
      if Vari2(J)>0  then
         Put("*"&Xnames(Vari2(J)));
      end if;
      if Vari1(J)>0 or Vari2(J)>0 then
         Put(")");
      end if;
      if J /= n+nlin then
        Put_Line("+");
      else
        New_Line;
      end if;
    end loop;
    Put_Line(" = Y(i)");
    New_Line;
  end Print_Equations;

  procedure Check_Equations(n : integer; nlin : integer;
                            A : Real_Matrix; Y : Real_Vector;
                            var1 : integer_vector; var2 : integer_vector;
                            var3 : integer_vector; vari1 : integer_vector;
                            vari2 : Integer_Vector; X : in out Real_vector) is
    resid : Real := 0.0;
    Xnames : array(1..5) of String(1..2) := ("X1", "X2", "X3", "X4", "X5");
  begin
    for i in 1..n loop
      Put_Line(Xnames(i)&" = "&Real'Image(X(i)));
    end loop;
    for j in n+1..n+nlin loop -- set non linear X
      X(j) := 1.0;
      if var1(j) > 0 then
        X(j) := X(j)*X(var1(j));
      end if;
      if var2(j) > 0 then
        X(j) := X(j)*X(var2(j));
      end if;
      if var3(j) > 0 then
        X(j) := X(j)*X(var3(j));
      end if;
      if vari1(j) > 0 then
        X(j) := X(j) / X(vari1(j));
      end if;
      if vari2(j) > 0 then
        X(j) := X(j) / X(vari2(j));
      end if;
    end loop;
    -- compute residual, abs sum of all equations
    for i in 1..n loop
      for j in 1..n+nlin loop
        resid := resid + A(i,j)*X(j);
      end loop; -- j
      resid := resid - Y(i);
      resid := abs(resid);
    end loop; -- i
    Put_Line("Check_Equation residual = "&Real'Image(resid));
    New_Line;
  end Check_Equations;

begin
  Put_Line("pde_nl21.adb running ");
  Put_Line("The PDE to be solved for u(x) is:");
  Put_Line(" Uxx(x) + 2*Ux(x) + U(x)^2  = f(x)");
  Put_Line(" ");
  Put_Line(" for testing code:");
  Put_Line("       f(x) = x^4 + 2x^2 + 4x + 3");
  Put_Line("       u(xmin=0)=1  u(xmax=1)=2");
  Put_Line(" ");
  Put_Line("analytic solution and boundary  ub(x) = x^2+1");
  Put_Line("       thus  Ux(x) = 2*x    Uxx(x) = 2");
  Put_Line("xmin="&Real'Image(xmin)&", xmax="&Real'Image(xmax));

  Put_Line("nx="&Integer'Image(nx)&" points for numeric derivative");
  hx := (xmax-xmin)/real(nx-1);
  Put_Line("x grid, analytic solution, f(x) ");
  for i in 1..nx loop
    xg(i) := xmin+Real(i-1)*hx;
    ug(I) := ub(xg(i));
    Ua(i) := ub(xg(i));
    Put("i="&Integer'Image(i)&", Ua(");
    Put(xg(i),3,5,0);
    Put(")=");
    Put(Ua(i),3,5,0);
    Put(", f(x)=");
    Put(f(xg(i)),3,5,0);
    New_Line;
  end loop;
  New_Line;

  Build_Vars;

  Build_Equations; -- discrete
  Print_AY;
  Print_Equations(dof, nlin, Var1, Var2, Var3, Vari1, Vari2);

  for i in 1..dof loop
    Us(i) := ug(i+1); -- set up exact solution
  end loop;
  Put_Line("test, giving exact solution");
  Check_Equations(dof, nlin, A, Y, Var1, Var2, Var3, Vari1, Vari2, Us);

  Simeq_Newton5(dof, nlin, A, Y, Var1, Var2, Var3, Vari1, Vari2,
                Us, 20, 0.001, 1);
  New_Line;

  -- guess at initial conditions, boundary not included
  for I in 1..dof loop
    Us(i) := 1.0;
  end loop;
  Put_Line("initial guess, all 1.0");

  Simeq_Newton5(dof, nlin, A, Y, Var1, Var2, Var3, Vari1, Vari2,
                Us, 20, 0.001, 1);
  for i in 2..nx-1 loop
    ug(i) := Us(i-1); -- ug is now boundary with computed solution
  end loop;

  avgerr := 0.0;
  maxerr := 0.0;
  Put_Line("U computed from nonlinear equations, Ua analytic, error ");
  for i in 1..nx loop
    err := abs(ug(i)-Ua(i));
    if err>maxerr then  maxerr := err; end if;
    avgerr := avgerr + err;
    Put("ug("&Integer'Image(I)&")=");
    Put(ug(i),3,7,0);
    Put(", Ua="); Put(Ua(i),3,7,0);
    Put(", err="); Put(ug(i)-Ua(i),3,7,0);
    New_Line;
  end loop;
  Put_Line("                   maxerr="&
           Real'Image(maxerr)&", avgerr="&
           Real'image(avgerr/Real(dof)));
  New_Line;

  Put_line("check_soln on computed solution against PDE");
  check_soln(ug);  -- for use when analytic solution, Ua, is not known
  New_Line;
  Put_line("just checking check_soln when given correct solution");
  check_soln(Ua); -- checking on checker
  New_Line;

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