%FGST  Forward Grad-Shafranov solver
% LY = FGST(L,LX) solves the Grad-Shafranov equation
% for a given conductor current distribution.
%
% parameters in L as obtained by FGSC and the inputs in LX 
% as obtained by FGSP<TOK> 
% returns a structure LY with the computed equilibrium.
%
% FGS-specific outputs:
% 
% .mkryl    Number of Krylov iterations carried out
% .nfeval   Number of function evaluations
% .rese     Residual of coil equation (0)
% .resy     Residual norm of Iy
% .resp     Residual norm of plasma current equation
% .resg     Residual norm of ag constraints
% .resFx    Residual norm of Fx
%
% [+MEQ MatlabEQuilibrium Toolbox+] Swiss Plasma Center EPFL Lausanne 2022. All rights reserved.

function [LY] = fgst(L,LX,varargin)

% optional parameters
for k = 1:2:length(varargin), L.P.(varargin{k}) = varargin{k+1}; end

%% Checks
agconstr = cellfun(@func2str,L.agconc(:,1),'UniformOutput',false);
assert(~any(contains(lower(agconstr),'cde')),...
  'Can not run FGS with cde constraint, use Ip constraint instead');
assert(~isequal(L.P.algoNL,'Picard') || ~L.P.usepreconditioner,...
  'can not use preconditioner for algoNL=''Picard'', please set usepreconditioner=0')

%% Debug
if L.P.debug
  meqinfo(L,LX)
  tstart=tic;
end

if L.P.usepreconditioner && any(LX.Ip~=0)
  % preconditioner
  LX1 = meqxk(LX,1);
  if isfield(L,'lin')  && ...
      all(L.lin.Iy(:)== LX1.Iy(:))  && ...
      all(L.lin.Ie   == [LX1.Ia;LX1.Iu])
    % re-use
    dFdIy = L.lin.dFdIy; dFdag = L.lin.dFdag;
    if L.P.debug, fprintf('FGS:computing preconditioner based on L.lin...'); end
  else
    % compute linearization
    if L.P.debug, fprintf('FGS:computing preconditioner, evaluate fgsFlin...'); end
    [dFdIy,dFdag] = fgsFlin(L,LX1);
  end
  ii = [L.ind.iy,L.ind.ig];
  JJ(ii,ii) = [dFdIy,dFdag];
  assert(~any(isnan(JJ(:))),'NaNs in Jacobian')
   assert(rank(JJ)==size(JJ,2),'Jacobian is not full rank')
  Prec = inv(JJ);
  if L.P.debug; fprintf('done\n'); end
else
  Prec = 1;
end

%% Time loop
for it = 1:numel(LX.t)
  LXt = meqxk(LX,it); % select inputs at this time
  
  if strcmpi(L.P.algoNL,'Picard')
    % Function handle to compute equation residual.
    F = @(x,doplot) fgsFpicard(x,L,LXt,doplot);
    x0 = [LXt.rA;LXt.zA;LXt.ag]./L.xscal; % initial values
  else
    % Function handle to compute equation residual.
    F = @(x,doplot) fgsF(x,L,LXt,LXt,doplot);
    % initial state guess
    if L.icde, IpD = LXt.IpD; else, IpD=[]; end
    x0([L.ind.iy,L.ind.ig,L.ind.ip],1) = [LXt.Iy(:);LXt.ag(:);IpD(:)]; % init NL state
    
  end
  % Solve the equation F(x)=0
  [X1, solverinfo] = solveF(...
    F,x0,'algoF', L.P.algoF,'mkryl', L.P.mkryl, 'Prec', Prec,...
    'relax', L.P.relax,'debug', L.P.debug, 'kmax', L.P.kmax, ...
    'tolF', L.P.tolF, 'algoGMRES', L.P.algoGMRES);
  
  % Compute outputs
  if solverinfo.isconverged
    [~, LYt] = F(X1,L.P.debug>3);
  else
    [~, LYt] = F(x0,L.P.debug>3);
    meqmsge('i',mfilename,L.P.tokamak,LXt.t,solverinfo.niter,...
    L.P.shot,'Newton solver did not converge','FGS:nosolution');
  end
   
  % Add other output info
  LYt.t = LXt.t; LYt.shot = LXt.shot;
  LYt.res    = solverinfo.res;
  LYt.mkryl  = solverinfo.mkryl;
  LYt.nfeval = solverinfo.nfeval;
  LYt.niter  = solverinfo.niter;
  LYt.isconverged = solverinfo.isconverged;
  
  if ~LYt.isconverged
    meqmsge('w',mfilename,L.P.tokamak,LYt.t,LYt.niter,L.P.shot,'fgs: No convergence','fgstFail');
  end
 
  % debug output
  if L.P.debug
    if ~rem(it-1,10)
      fprintf(  '\n  %7s | %3s | %4s | %4s | %8s | %6s | %6s | %6s | %6s | %6s\n', ...
        't','it', 'niter','conv', 'res', 'Ip [kA]' ,'zA[m]','rA[m]','bp','qA');
      fprintf('  %s\n',repmat('=',85,1));
    end
    fprintf(  '  %7.5f | %3d | %5d | %4d | %8.2e | %6.2f | %6.3f | %6.3f | %6.3f | %6.3f\n',...
      LYt.t,it,LYt.niter,LYt.isconverged, LYt.res,...
      LYt.Ip/1e3,LYt.zA(1),LYt.rA(1),LYt.bp,LYt.qA(1));
  end
  if L.P.debug>1, meqplott(L,LYt); end
  
  % store in LY
  LY(it) = LYt; %#ok<AGROW>
end

% pack outputs
LY = meqlpack(LY, {'FW'});

% mask non-converged
if ~L.P.LYall
  LY = meqxk(LY,LY.isconverged);
end

if L.P.debug, fprintf('FGS done: time elapsed: %2.2f[s]\n',toc(tstart)); end
end
