% FGSF Physics operator for FGS
% [res, LY] = fgsF(x,L,LX,LYp,doplot)
% Returns the (scaled) residual of
% * GS equation res(L.ind.iy)
% * Constraint equations res(L.ind.ig)
% * CDE equation (optional when running FGE with CDE) (res(L.ind.ip))
%
% x is the nonlinear state vector 
% [Iy(:);ag(:)] or [Iy(:);ag(:);Ip(:)]
%
% L is the ancilliary data structure
% LX contains input data
% LYp contains the previous equilibrium.
% for FGS, LYp=LX, for FGE, LYp and LX will be different
%
% [+MEQ MatlabEQuilibrium Toolbox+] Swiss Plasma Center EPFL Lausanne 2022. All rights reserved.

function [res, LY] = fgsF(x,L,LX,LYp,doplot)
if nargin<5
  doplot = false;
end

%% init
res = zeros(L.ny+L.ng+L.np,1);
if nargout>1, LY = LYp; end

%% Extract NL unknowns
Iy  = reshape( x(L.ind.iy) , L.nzy, L.nry);
ag  = reshape( x(L.ind.ig) , L.ng ,     1);

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

% Sign of Ip
Ia = LX.Ia;
Iu = LX.Iu;
sIp = sign(LX.Ip(1));
rBt = LX.rBt;

%% Solve GS to get new flux map
Fx = meqFx(L,Iy,[Ia;Iu]);

%% Plasma domain and current distribution
[rA,zA,FA,dr2FA,dz2FA,drzFA,rX,zX,FX,~,~,~,~,~,...
  rB,zB,FB,lB,lX,Opy] = meqpdom(Fx,sIp,L.P.isaddl,L);
nA=numel(rA); nB = numel(FB); % number of boundaries = number of active domains

if any(IpD), pla=true; else, pla=false; end 
%% Various domain checks
if pla && ~lB 
  meqmsge('w',mfilename,L.P.tokamak,LX.t,0,L.P.shot,'could not find LCFS','NoLCFS'); 
  res(:)=NaN; return; 
end
if pla && nA==0
  meqmsge('w',mfilename,L.P.tokamak,LX.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,LX.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,LX.t,0,L.P.shot,'Tilted Hessian at magnetic axis','MagAxTilt'); 
  res(:)=NaN; return; 
end
if nB > L.nD
  meqmsge('w',mfilename,L.P.tokamak,LX.t,0,L.P.shot,...
    sprintf('number of found domains (%d) exceeds maximum number (%d)',nB,L.nD),'TooManyDomains');
  res(:)=NaN; return
end

%% Plasma current distribution from basis function coefficients
[F0,F1] = meqpdomlim(FA,FB,L.nD); % flux limits of various domains

residualIy = Iy(:);
[Tyg,TpDg,ITpDg] = L.bfct(1,L.bfp,Fx,F0,F1,Opy,L.ry,L.iry);
Iy = reshape(Tyg*ag,L.nzy,L.nry);
residualIy = Iy(:) - residualIy; % this gives residual = Iy-Iy(previous)

%% CDE constraint residuals (FGE only)
residualcde = zeros(L.np,1);
SQ = zeros(L.nQ,L.nD);

if L.np
  IpQ = zeros(L.nQ,L.nD); FtPQ = zeros(L.nQ,L.nD); TQ = zeros(L.nQ,L.nD); Fy0 = zeros(L.ny,1);
  % Post processing needed for CDE
  switch L.P.cde
    % Prevent using meqintQ (CPU intensive) when not needed
    case {'cde0D','cde1D','cde0Dss'}
      [IpQ,~,~,~,FtQ,Ft0Q] = meqintQ(L,F0,F1,rBt,ag,Fx,Opy,sIp);
      FtPQ = FtQ+Ft0Q; % total Plasma toroidal flux on pQ grid
      [~,~,~,TQ] = meqprof(L.fPg,L.fTg,ag,L.pQ(:).^2,F0,F1,LX.rBt,L.bfct,L.bfp,L.idsx);
    case {'OhmTor_rigid'}
      Fy0  = [L.G.Mxa(L.lxy(:),:),L.G.Mxu(L.lxy(:),:)]*[Ia; Iu];  % Current time step
      if ~L.P.iLpext
        Fyp = Fx(L.lxy(:)) - Fy0; % Cheaper than computing Myy*Iy
        LX.Lp = Fyp'*Iy(:)/IpD^2; % This is positive definite since it is equivalent to (Iy Myy Iy)/Ip^2
      end
  end
  % Compute residual CDE
  indcde = find(L.icde);
  for ip=1:L.np
    iD = indcde(ip);
    if iD<=nB
      [residualcde(iD),SQ(:,iD)] = ...
        meqcde(L,LX,LYp,F0,F1,Opy,Iy,Fx,Fy0,IpD(iD),iD,IpQ(:,iD),FtPQ(:,iD),TQ(:,iD));
    else
      residualcde(iD) = IpD(iD) - LYp.IpD(iD); % if CDE domain is not present, keep previous Ip
    end
  end
end

%% ag constraint residuals
residualag = meqagcon(L,LX,LYp,...
    F0,F1,rA,dr2FA,dz2FA,drzFA,IpD,ag,Fx,...
    Opy,L.fPg,L.fTg,TpDg,ITpDg);

%% Collect residuals
iF = [L.ind.iy,L.ind.ig,L.ind.ip]; % index into F
res(iF) = L.resscal(iF) .* [residualIy;residualag(:);residualcde(:)];

%% Add outputs
if nargout>1 || doplot
 
  LY = struct('shot',LX.shot,'t',LX.t,'aq',LYp.aq,'aW',LYp.aW);
  LY = meqpost(L,LY,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(residualIy);
  resg = norm(residualag);
  resp = norm(residualcde);
  rese = 0;
  resFx = abs(max((Fx-meqFx(L,Iy,[Ia;Iu])))); % Poloidal flux residual w.r.t. new Iy
  
  LY = meqlarg(LY,resy,resg,resp,rese,resFx);
end

%% debug
if doplot
  %%
  clf;
  subplot(121)
  plot(res,'.'); title('residual values');
  title(sprintf('max(|res|)=%4.2e',max(abs(res))));

  subplot(122)
  imagesc(L.ry,L.zy,Opy,[0 3]); axis xy; hold on;
  contour(L.rx,L.zx,Fx,21,'k--'); hold on;
  if nA==1, FF=[FB,FB]; else, FF=FB; end
  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',LY.Ip));
  axis equal tight; drawnow;
end
