function  [res, LYt] = fgeFpicard(x,L,LXt,LYp,doplot)
% [res, LYt] = fgeFpicard(x,L,LXt,LYp,doplot)
% FGE operator using Picard method for convergence between flux and plasma
% The method is inspired from the FEQIS code by E. Fable, IPP Garching [To be published]
%
%  This function can called by an external Newton solver which renders res
%  requal to zero.
%
% Inputs:
%  * x is [rC;zC;ag] (divided by scaling L.xscal)
%    where rC,zC is the location to which the magnetic axis is forced to
%    converge by applying an additional flux 
%    Fx = Fx(plasma) + Fx(ext) + FR*R^2 + FZ*Z
%    and ag are the basis function coefficients
%  * L is the usual L structure
%  * LXt contains the inputs, in particular Ia,Iu,Ip,constraints, 
%    and intial guess for Iy.
%  * LYp contains the previous time step equilibrium (for fge only)
%  * doplot = true forces plotting
%
% Outputs:
%  * res is the residual [FR,FZ,rescon] (with scaling L.resscal)
%    Where FR,FZ are the terms of the required stabilizing flux
%    and rescon is the residual of the meqagcon constraints
%  * Optional LYt output triggers post-processing and returns full
%    equilibirum structure
%
% [+MEQ MatlabEQuilibrium Toolbox+] Swiss Plasma Center EPFL Lausanne 2022. All rights reserved.

debuglevel = L.P.debug;
postprocess = nargout>1 || doplot;

tolPicard = L.P.tolPicard;
tolInner  = max(tolPicard/1e2,1e-14);

% variables from x
x = L.xscal.*x; % rescale x
rC = x(1);
zC = x(2);
ag = x(2+(1:L.ng));

if any(L.icde)
  IpD = x(L.ind.ip); % Ip is part of state
else
  IpD = LXt.(L.agconc{1,3}); % LX.Ip or LX.IpD as source
end

% get get initial guess from LX
Ie  = [LXt.Ia;LXt.Iu];
Iy  = LXt.Iy;
sIp = sign(LXt.Ip(1));
rBt = LXt.rBt;

if any(IpD), pla=true; else, pla=false; end % flag for plasma case or not

%% First flux calculations
Fx = meqFx(L,Iy,Ie); % Total flux
Fxe = meqFx(L,zeros(L.nzy,L.nry),Ie); % Flux from external coils only

%% Picard loop, converge Iy and Fx
for ipicard=1:L.P.kpic
  FR = 0;  FZ = 0;  dFR = NaN; dFZ = NaN;   
  zFx = Fx; zIy = Iy;
  %% Inner loop
  % Find stabilizing terms FR,FZ required so that
  % (rA,zA) of Fx+Fstab at (rC,zC)
  % where Fstab = FR*R^2 + FZ*Z
  
  for iinner = 1:L.P.kinner
    
    %% new rA,zA and Hessians
    [zA,rA,FA,dz2FA,dr2FA,drzFA,~,~,~,~,~,~,stat] = ...
      asxymex(Fx,L.G.zx,L.G.rx,L.P.dasm,L.dzx,L.drx,L.idzx,L.idrx,L.Oasx,L.dimw{:});
    if L.P.icsint
      [zA,rA,FA,dz2FA,dr2FA,drzFA,~,~,~,~,~,~] = ...
        asxycs(Fx,zA,rA,[],[],L);
    end
    nA = numel(rA);
    if ~stat
      meqmsge('w',mfilename,L.P.tokamak,LXt.t,0,L.P.shot,'too many A/X points','dimWoverflow');
      res(:)=NaN; return;
    end
    if pla && (nA == 0)
      meqmsge('w',mfilename,L.P.tokamak,LXt.t,0,L.P.shot,'No magnetic axis','NoMagAx');
      res(:)=NaN; return;
    end
    if ~L.P.idoublet && (nA>1)
      meqmsge('w',mfilename,L.P.tokamak,LXt.t,0,L.P.shot,'Multiple magnetic axes but not doublet','MultMagAx');
      res(:)=NaN; return;
    end
    if any(dr2FA.*dz2FA<drzFA.^2)
      meqmsge('w',mfilename,L.P.tokamak,LXt.t,0,L.P.shot,'Tilted Hessian at magnetic axis','MagAxTilt');
      res(:)=NaN; return;
    end
    
    if debuglevel>3 && iinner==1
      % plot this Fx
      if ipicard == 1, clf;
        ax1=subplot(121); meqgplot(L.G,gca,'vl'); hold on;
        ax2=subplot(122); meqgplot(L.G,gca,'vl'); hold on;
      else
        delete(hc1); delete([hC,hA]); delete(hc3);
      end
      [~,hc1] = contour(ax1,L.rx,L.zx,Fx-FA,21);
      hA = plot(ax1,rA,zA,'sr'); hC = plot(ax1,rC,zC,'xk');
      
      Fxstab = FR*L.rrx.^2 + FZ*L.zzx; % stabilization field
      [~,hc3] = contour(ax2,L.rx,L.zx,Fxstab,21);
      
      drawnow;
    end
    
    % Check axis and exit inner loop if close enough
    drAC = (rC-rA); dzAC = (zC-zA);
    
    % optional debugging display
    if debuglevel>4
      if ~rem(iinner-1,10)
        fprintf('\n%-5s.%-9s %-9s %-9s %-9s %-9s\n',...
          'ipic','ii','dFR','dFZ','rAC','zAC');
      end
      fprintf('%2d.%2d:    %+8.3e %+8.3e %+8.3e %+8.3e\n',...
        ipicard,iinner,dFR,dFZ,drAC,drAC);
    end
    
    if abs(drAC)< tolInner*L.P.r0 && abs(dzAC)< tolInner*L.P.r0, break; end

    %% Find stabilizing field to put axis on desired place.
    % take F = Fx + FR*R^2 + FZ*Z
    % Locally at rC,zC close to rA,zA
    % F = FA + 0.5*dr2FA*(rC-rA)^2 + 0.5*dz2FA*(zC-zA)^2 + drzFA*(zC-zA)*(rC-rA) + dFR*rC^2 + dFZ*zC
    % dF/dR = dr2FA*(rC-rA) + drzFA*(zC-zA) + 2*rC*dFR
    % dF/dZ = dz2FA*(zC-zA) + drzFA*(rC-rA) + dFZ
    % solve for dFR,dFZ such that dF/dR,dF/dZ=0 at rC,zC (local extremum)
    
    KR=1; KZ=1; % inner loop gain (1 = no damping)
    dFR = -KR*(dr2FA*drAC + drzFA*dzAC)/(2*rC);
    dFZ = -KZ*(dz2FA*dzAC + drzFA*drAC);
  
    % add new contribution
    dFstab = dFR.*L.rrx.^2 + dFZ*L.zzx;
    Fx = Fx + dFstab;
    
    % keep track of total FR,FZ
    FR = FR + dFR;
    FZ = FZ + dFZ;
  end % end inner loop for FR,FZ
  
  %% Plasma current distribution from basis function coefficients
  % rerun meqpdom with X points
  [rA,zA,FA,dr2FA,dz2FA,drzFA,rX,zX,FX,~,~,~,~,~,...
      rB,zB,FB,lB,lX,Opy] = meqpdom(Fx,sIp,L.P.isaddl,L);
  
  [F0,F1] = meqpdomlim(FA,FB,L.nD); % flux limits of various domains
  
  [Tyg,TpDg,ITpDg] = L.bfct(1,L.bfp,Fx,F0,F1,Opy,L.ry,L.iry);
  
  % new Iy
  Iyscal = sum(IpD)/(TpDg*ag); % scale Iy to maintain Ip
  Iy = Iyscal*reshape(Tyg*ag,L.nzy,L.nry);
  % NB: This scaling is necessary because otherwise the Picard algorithm 
  % is unstable, continuing to increase or decrease Ip
  % Scaling ag instead makes the Picard algorithm oscillate the axis position.
  % This scaling indeed introduces an inconsistency between ag and Ip (and Iy)
  % However this inconsistency is resolved in the outer Newton loop which makes
  % rescon=0 hence enforcing consistency between Ip and ag.

  if isequal(L.code,'fge')
    % for FGE, update new Ie given new Iy
    % Implicit time stepping
    dtFepdot = L.P.mey*L.Mey*(Iy(:)-LYp.Iy(:)); % Voltage at conductors due to plasma change
    Ie = (L.Mee + diag(LXt.dt*L.Re))\...
      ([LXt.dt*LXt.Va;zeros(L.G.nu,1)] + L.Mee*[LYp.Ia;LYp.Iu] - dtFepdot);

    Fx  = meqFx(L,Iy,Ie); % update flux with new plasma + conductor part
  else
    % for FGS, keep conductors fixed
    FxP = meqFx(L,Iy); % update only plasma part of the flux
    Fx = FxP + Fxe;    % total flux
  end

  % Residual
  resFx = abs(max(Fx(:)-zFx(:)));
  
  if debuglevel>2
    if ~rem(ipicard-1,10) || debuglevel>4
      fprintf('%-4s %-6s %-10s %-10s %-10s %-10s %-10s %-10s\n',...
        'ipic','ninner','FR','FZ','resFx','rA','zA','Ip');
    end
    fprintf('%3d:  %-5d %+8.3e %+8.3e %+8.3e %+8.3e %+8.3e %+8.3e\n',...
      ipicard,iinner,FR,FZ,resFx,rA,zA,sum(Iy(:)));
  end
  
  % Exit condition for Picard loop
  if resFx < tolPicard
    if debuglevel>2, fprintf('Picard iteration converged\n\n'); end
    break 
  end
end % end Picard iteration loop

% Constraint residual
resagcon = meqagcon(L,LXt,LYp,...
    F0,F1,rA,dr2FA,dz2FA,drzFA,IpD,ag,Fx,...
    Opy,L.fPg,L.fTg,TpDg,ITpDg);

% stack all residuals that we want to be equal to zero
res = L.resscal([L.ind.iy,L.ind.ig]).*[FR;FZ;resagcon];

%% Add outputs
if postprocess
  Ia = Ie(1:L.G.na); Iu = Ie(L.G.na+(1:L.G.nu));

  LYt = struct('shot',LXt.shot,'t',LXt.t,'aq',LYp.aq,'aW',LYp.aW);
  LYt = meqpost(L,LYt,TpDg,ITpDg,ag,...
    Fx,FA,FB,rA,zA,dr2FA,dz2FA,drzFA,rB,zB,lB,lX,rX,zX,FX,...
    rBt,Ia,Iu,Iy,Opy);

  % collect residuals
  resy = norm(Iy-zIy);
  resg = norm(res(L.ind.ig));
  resp = norm(res(L.ind.ip));
  if isequal(L.code,'fge')
    rese = (L.Mee*(Ie-[LYp.Ia;LYp.Iu]) + dtFepdot)/LXt.dt + L.Re.*Ie - [LXt.Va;zeros(L.G.nu,1)]; % conductor equation residual
  else
    rese = 0;
  end
  LYt = meqlarg(LYt,resy,resg,resp,rese,resFx);
end

%% debug
if doplot
  %%
  clf;
  subplot(121)
  bar(1:numel(res),res);
  title(sprintf('max(|res|)=%4.2e',max(abs(res))));
  
  subplot(122)
  if nA==1, FF=[FB,FB]; else, FF=FB; end
  imagesc(L.ry,L.zy,Opy,[0 3]); axis xy; hold on;
  contour(L.rx,L.zx,Fx-F0(1),21); hold on;
  contour(L.rx,L.zx,Fx,FF,'k','linewidth',2,'color','k')
  plot(L.G.rl,L.G.zl,'w','linewidth',2);
  plot(rA,zA,'oc',rX,zX,'xc');
  title(sprintf('Ip=%5.2e\nzA=%7.5e\nrA=%7.5e\nFB=%7.5e',LYt.Ip,LYt.zA,LYt.rA,LYt.FB));
  axis equal tight; drawnow;
end

end
