function [x, solverinfo] = solveF(F,x0,varargin)
% function [x, solverinfo] = solveF(F,x0,varargin)
% General solver to find solution of F(x)=0
% Implements JFNK (Jacobian-Free Newton-Kryov Method
%    or Newton Method with Finite-Difference approximation of gradients
%
% [+MEQ MatlabEQuilibrium Toolbox+] Swiss Plasma Center EPFL Lausanne 2022. All rights reserved.

%% Parse inputs
ip = inputParser; % initialize parser
default.mkryl = 200; % default value of m (size of Krylov subspace)
default.tolF = 1e-10; % default value of tolF (minimum tolerance max(abs(F(x))))
default.kmax = 50; % default value of kmax (maximum number of Newton steps)
default.epsilon_d = sqrt(eps); % default values of epsilon_d (variation used to approximate the Jacobian)
default.algoF = 'jfnk'; % 'jfnk' 'Newton' Specify the algorithm to be used
default.epsilon_d_tol = 10*eps; % min value of epsiol_d when robust decrease on
default.relax = 0; % Apply relaxation on solution
default.debug = 0; % Show convergenve info during iterations.
default.Prec = 1; % preconditionig matrix
default.workers = 0; % number of workers for parallel computation
default.robust_on = 1;
default.nrestarts = 5;
default.epsilon_res = 1e-12; % Residual tolerance for inner residual of MGS methods.
% 'cgs' = Classical Gram-Schmidt
% 'mgsgiv' = Modified Gram-Schmidt with Givens rotations
% 'mgsaya' = Modified Gram-Schmidt with Ayachour's method (no Givens rotations)
default.algoGMRES = 'mgsgiv';
for ifield=fieldnames(default)'
  addParameter(ip,ifield{1},default.(ifield{1}));
end
parse(ip,varargin{:}); % parse the inputs, input xyz is now available in P.xyz
P = ip.Results;

%% debug/verbosity options
dispprogress   = (P.debug>1); % textual display of each iteration results
ploteachnewton = (P.debug>4); % each call F in newton iteration produces a plotf
ploteachkryl   = (P.debug>5); % each call of F during krylov iteration produces a plot

%% Initialization
n = length(x0); % size of solution sought
m = min(P.mkryl,n); % size of the Krylov space
H = zeros(m+1,m); % Hessenberg matrix (memory allocation)
V = zeros(n,m+1); % matrix with columns containing the Krylov vectors
switch P.algoGMRES
  case 'cgs'
    use_modified_gram_schmidt = false;
  case 'mgsgiv'
    use_modified_gram_schmidt = true;
    C = zeros(m+1,1); % Cosine component of Givens rotations
    S = zeros(m+1,1); % Sine component of Givens rotations
  case 'mgsaya'
    use_modified_gram_schmidt = true;
    u = zeros(m,1); % Residual vector.
    R = zeros(m,m); % Inverse Hessenberg sub-matrix (Hk).
end

x = x0; % Initial solution
b = -F(x,ploteachnewton); % Initial residual

% initialize
kjfnk = 0; % iteration counter
nfeval = 1; % Counter for number of function call steps
resmax = max(abs(b)); resnorm = norm(b);
res = resmax; accept=true; % init
alpha_step = 1-P.relax;
restarts = 0; % number of restarts

if dispprogress
  fprintf(   '\n%12s | %15s | %15s | %20s | %7s | %4s | %5s |\n', 'k_step', '||dx||/||x||)', 'max(residual)', 'norm(residual)', 'alpha','mkry','accept')
  fprintf(   '%12d | %15.3e | %15.3e | %20.17f | %2.1e | %4d | %5d\n', kjfnk, NaN, resmax, resnorm, alpha_step, NaN, accept)
end

%% Start Newton iteration
tstart=tic;
while (res >= P.tolF) && (kjfnk < P.kmax) % iterate until tolerance is reached
  % Relaxation parameter update
  m = min(P.mkryl,n); % Reset size of the Krylov space for this iteration.
  if accept
    % compute new candidate Newton direction
    switch P.algoF
      case 'Newton'
        % Compute gradient
        J = jacfd(F,x);
        dx = - J\F(x,ploteachnewton);
        nfeval = nfeval + 2*numel(x) + 1;

      case 'jfnk'
        % JFNK with preconditioning
        % Krylov solution of J delta x = -F(xkm1)
        beta = norm(b);
        V(:,1) = b/beta; % First vector of krylov space
        Pinv = P.Prec; % Set preconditioner
        switch P.algoGMRES
          case 'cgs' % Classical Gram-Schmidt.
            for j=1:m % Applied on all Krylov space dimension
              % One step in the Arnoldi algorithm
              [V,H,~,nfeval,iscompleted] = one_step_arnoldi(V,H,j,P,Pinv,F,x,b,nfeval,ploteachkryl,use_modified_gram_schmidt);
              if ~iscompleted
                break
              end

              if (j<m) && (H(j+1,j) < P.epsilon_res) % stop increasing the krylov space size
                if P.debug
                  fprintf('Sufficiently large Krylov space m=%d\n',j)
                end
                m = j;
                break
              end
            end
            e1 = zeros(m,1); e1(1) = 1;
            ym = H(1:m,1:m)\(beta*e1); % solve for the minimum in the least square sense
          case 'mgsgiv' % Modified Gram-Schmidt with Givens rotations.
            j = 1;
            G = beta*eye(m+1,1);
            while (beta > P.epsilon_res) && (j <= m)  % Applied on all Krylov space dimension
              % One step in the Arnoldi algorithm
              [V,H,~,nfeval,iscompleted] = one_step_arnoldi(V,H,j,P,Pinv,F,x,b,nfeval,ploteachkryl,use_modified_gram_schmidt);
              if ~iscompleted
                break
              end

              % Make the Hessenberg matrix upper triangular via Givens rotations.
              if j > 1
                H(1:j,j) = givens_rotation(C(1:j-1),S(1:j-1),H(1:j,j),j-1);
              end
              % Calculate new rotations.
              normH = norm(H(j:j+1,j));
              if normH ~= 0
                C(j) = H(j,j)/normH;
                S(j) = -H(j+1,j)/normH;
                H(j,j) = C(j)*H(j,j) - S(j)*H(j+1,j);
                H(j+1,j) = 0;
                G(j:j+1) = givens_rotation(C(j),S(j),G(j:j+1),1);
              end

              beta = abs(G(j+1));
              j = j+1;
            end
            m = j-1;
            ym = H(1:m,1:m)\G(1:m); % solve for the minimum in the least square sense
          case 'mgsaya' % Modified Gram-Schmidt with Ayachour method.
            c = beta;
            % This follows the implemention of Theorem 3.1 of
            % A fast implementation for GMRES method, by E.H. Ayachour,
            % Journal of Computational and Applied Mathematics 159 (2003), 269 - 283
            for j=1:m  % Applied on all Krylov space dimension
              % One step in the Arnoldi algorithm
              [V,H,Hinv,nfeval,iscompleted] = one_step_arnoldi(V,H,j,P,Pinv,F,x,b,nfeval,ploteachkryl,use_modified_gram_schmidt);
              if ~iscompleted
                break
              end

              % Build next part of R=inv(H_k') matrix, see
              % GMRES implementations and residual smoothing techniques for solving ill-posed linear systems,
              % by Matinfar, Zareamoghaddam, Eslami, and Saeidy,
              % Computers and Mathematics with Applications 63 (2012) 1-13
              % Algorithm 7, step 2.c.
              % R_j = (R_j-1 -R_j-1*g*Hinv)
              %       (         Hinv      )
              % From steps 2.c to 2.e, if H(j+1,j)=0, then the method
              % exits, or equivalently applies Hinv=1
              if j > 1
                R(1:j-1,j) = -R(1:j-1,1:j-1)*H(2:j,j)*Hinv;
              end
              R(j,j) = Hinv;
              u(j) = R(1:j,j)'*H(1,1:j)';
              c = 1 / sqrt(1 + u(1:j)'*u(1:j));
              if (beta*c < P.epsilon_res)
                break;
              end
            end
            m = j;
            ym = (beta*c*c)*R(1:m,1:m)*u(1:m);
          otherwise
            error('Not a valid GMRES method: "%s"', P.algoGMRES)
        end
        dx = Pinv*(V(:,1:m)*ym); % compute the delta x
      otherwise
        error('Not a valid solution method: "%s"', P.algoF)
    end
  end

  % x step in dx direction
  xtmp = x + alpha_step*dx;

  % Compute the residual
  b = -F(xtmp,ploteachnewton);

  nfeval = nfeval + 1;

  % Check acceptance of this candidate solution
  resnormj = sqrt(sum(b.^2));
  resmaxj = max(abs(b));

  if any(isnan(b)) || resnormj>resnorm
    accept = false;
  else
    accept = true;
  end

  %% Update iteration diagnostics
  kjfnk = kjfnk+1; % iteration counter

  dx_rel_norm = alpha_step*norm(dx)/norm(xtmp); % the maximum incremental percentage
  if dispprogress
    fprintf(   '%12d | %15.3e | %15.3e | %20.17f | %2.1e | %4d | %5d\n', kjfnk,dx_rel_norm, resmaxj, resnormj, alpha_step, m, accept)
  end

  if accept
    %% Update the solution x
    x = xtmp;
    % residuals
    resnorm = resnormj; % total residual
    resmax  = resmaxj; % maximum residual
    res = resmax; % this one is checked
    alpha_step = min(1-P.relax,alpha_step*2);
  else
    if (dx_rel_norm/res < 1e-5 || alpha_step < 1e-3) && (restarts < P.nrestarts)
      % local minimum, restart with randomized new dx
      % limit the number of restarts
      restarts = restarts + 1;
      % reset the step counter
      nfeval = 0;
      %dx = dx+1e-3*randn(size(x)).*dx;
      alpha_step = 1;
      %x = x + alpha_step*dx;
      x = x + 1e-3*randn(size(x)).*x;
      resnorm = realmax; resmax = realmax;
      res = resmax; accept = true;
      if dispprogress; fprintf('restart due to likely local minimum\n'); end
    else
      %just try smaller step
      alpha_step = alpha_step/4; % decrease alpha step for next iteration
    end
  end
end
if any(isnan(b))
  x = x*NaN;
end

% Output statistics
solverinfo.isconverged = (res <= P.tolF) && (kjfnk < P.kmax) && ~any(isnan(x));
solverinfo.niter = kjfnk;
solverinfo.residual_max = resmax;
solverinfo.residual_2norm = resnorm;
solverinfo.res = res;
solverinfo.restarts = restarts;
solverinfo.mkryl = m;
solverinfo.nfeval = nfeval;

if P.debug > 1
  if solverinfo.isconverged
    fprintf( ' \n Method: %s converged in %d iterations with %d F evaluations \n',...
      P.algoF,kjfnk, nfeval);
  elseif ~any(isnan(b))
    fprintf( ' \n Method %s did not converge \n ', P.algoF );
  else
    warning('No valid solution found using algo %s\n', P.algoF)
  end
  fprintf(' Elapsed time: %f[s]\n',toc(tstart));
end

end


function [V,H,Hinv,nfeval,iscompleted] = one_step_arnoldi(V,H,j,P,Pinv,F,x,b,nfeval,ploteachkryl,modified_gram_schmidt)
% Perform a new step in the Arnoldi algorithm by approximating the Jacobian,
% computing a Gram-Schmidt step and filling the matrices V and H with the specified function handle.
% Hinv is an optional output and if required it is used to normalize the basis.
iscompleted = true; % Initialize flag
Hinv = 1;           % Default Hinv

% Right side preconditioner as in eq 28, D.A.Knoll  JCP 193 (2004) 357-397
dxstep = P.epsilon_d*(Pinv*V(:,j));
Fstep = F(x + dxstep,ploteachkryl);
nfeval = nfeval + 1;

% Reduce variational step
epsilon_d = P.epsilon_d; % Initial step to approximate the Jacobian
while P.robust_on && any(~isfinite(Fstep)) && epsilon_d > P.epsilon_d_tol
  epsilon_d = epsilon_d/10; % smaller step
  Fstep = F(x + epsilon_d*(Pinv*V(:,j)),ploteachkryl);
  nfeval = nfeval + 1;
  if P.debug > 1
    fprintf('\n *** Krylov space F(x + eps) = NaN try to reduce dx *** \n\n\n');
  end
end
if epsilon_d < P.epsilon_d_tol
  warning('Unable to compute directional derivative');
  iscompleted = false;
  return
end

% Modified Gram-Schmidt procedure for constructing Krylov vectors.
V(:,j+1) = ( Fstep + b )/epsilon_d; % Compute new basis
if modified_gram_schmidt
  for i = 1:j
    H(i,j) = V(:,i)'*V(:,j+1); % Compute Hessenberg matrix with updated V(:,j+1)
    V(:,j+1) = V(:,j+1) - H(i,j)*V(:,i); % Arnoldi: subtract the component to make it orthogonal
  end
else
  H(1:j,j) = V(:,1:j)'*V(:,j+1); % Compute Hessenberg matrix
  V(:,j+1) = V(:,j+1) - V(:,1:j)*H(1:j,j); % Arnoldi: subtract the component to make it orthogonal
end
H(j+1,j) = norm(V(:,j+1));
if H(j+1,j) ~= 0
  Hinv = 1 / H(j+1,j);
  V(:,j+1) = V(:,j+1) * Hinv; % normalize basis
end
end

function rotated = givens_rotation(cos_rot,sin_rot,mat_in,k)
% Applies Givens rotations on k columns to perform essentially a QR decomposition.
%
% These rotations are noted in the original GMRES paper:
%   GMRES: A generalized minimal residual algorithm for solving nonsymmetric linear systems
%   SIAM Journal on scientific and statistical computing, 1986
% but later papers refer to them as "Givens rotations".
rotated = mat_in;
for i = 1:k
  w1 = cos_rot(i)*rotated(i) - sin_rot(i)*rotated(i+1);
  w2 = sin_rot(i)*rotated(i) + cos_rot(i)*rotated(i+1);
  rotated(i:i+1) = [w1,w2];
end
end
