function u = solveStationaryNonlinearEM(self,femodel,u0)
% solveStationaryNonlinearEM - Internal function for solving nonlinear
% magnetostatic problems
%
% Similar to usual solveStationaryNonlinear, but with updates specific to
% the EM workflow for robustness and performance reasons.

% Copyright 2022 The MathWorks, Inc.

p = femodel.p;
t = femodel.t;
if ~self.ParentModel.IsTwoD
    self.PDESystemSize = 3;
end
N = self.PDESystemSize;
np = size(p,2);
coefstruct = femodel.coefstruct;

pseudocoefstruct = coefstruct;
for i = 1:numel(pseudocoefstruct.Coefficients.c)
    pseudocoefstruct.Coefficients.c{i} = 1;
    pseudocoefstruct.Coefficients.a{i} = ones(N*N,1);
    pseudocoefstruct.Coefficients.m{i} = ones(N*N,1);
    pseudocoefstruct.Coefficients.f{i} = zeros(N,1);
end

if self.ParentModel.IsTwoD
    dummyResidual=@(u)compResidual2D(u,self,p,t,pseudocoefstruct,np*N,femodel);
    getResidual=@(u)compResidual2D(u,self,p,t,coefstruct,np*N,femodel);
    numjac_fh=@(~,u)compResidual2D(u,self,p,t,coefstruct,np*N,femodel);
    numdofs = np*N;
else
    dummyResidual=@(u)compResidual3D(u,self,p,t,pseudocoefstruct,np*N,femodel);
    getResidual=@(u)compResidual3D(u,self,p,t,coefstruct,np*N,femodel);
    numjac_fh=@(~,u)compResidual3D(u,self,p,t,coefstruct,np*N,femodel);
    numdofs = femodel.numEdges;
end

% Initial residual 
[r,pdenonK,pdenonF,u]=getResidual(u0);

if 1/condest(pdenonK) > 1e-10
    u=pdenonK\pdenonF;
    [r,pdenonK]=getResidual(u);
end

if ~allfinite(u) || ~allfinite(r)
    error(message('pde:pdenonlin:InvalidInitGuess'))
end

Knorm=pdenonK;
nrm=resNorm(r,self,Knorm);

if strcmp(self.SolverOptions.ReportStatistics,'on')
    disp('Iteration     Residual     Step size  ')
    fprintf('%4i%20.4e\n',0,nrm);
end

% For initial numjac call
fac=[];
G=[];
iter=0;
% Pad out S
[~,pseudoK]=dummyResidual(ones(size(u0)));
S=(pseudoK~=0);
if ~self.ParentModel.IsTwoD
    % in 3D, augment sparsity pattern with nodal connections. Perturbing an
    % edge leads to a nonlocal change in adjacent edges which must be
    % accounted for in the sparsity pattern for the Jacobian to be correct.
    S(1:femodel.numEdges,1:femodel.numEdges)...
        =(pseudoK(1:femodel.numEdges,1:femodel.numEdges)+femodel.M*femodel.M.')~=0;
end

% Do not display output of linesearch optimizer, and set maximum number of
% function evaluations. Each function evaluation requires reassembly.
opts = optimset("Display","off","MaxFunEvals",20);

while (nrm > self.SolverOptions.ResidualTolerance)  % Gauss-Newton iteration loop

    iter=iter+1;
    if iter>self.SolverOptions.MaxIterations
        error(message('pde:pdenonlin:TooManyIter'));
    end

    thresh = calcThresh(u, numdofs);
    S=S|(pdenonK~=0);
    [J,fac,G]=numjac(numjac_fh,0,u,r,thresh,fac,0,S,G); 
    lastwarn('','');
    warnState = warning('off', 'MATLAB:singularMatrix');

    cor = -J\r;
    [~,wid] = lastwarn;
    if(strcmp(wid,'MATLAB:singularMatrix'))
        warning(warnState);
        error(message('pde:pdenonlin:SingularJacobian'))
    end
    warning(warnState);

    % line search
    f = @(alpha) resNorm( getResidual(u + alpha*cor),self,Knorm );
    % alfa is correction, nrm is the residual norm. 
    [alfa,nrm] = fminbnd(f,0,1.25,opts);
    if alfa<self.SolverOptions.MinStep
        error(message('pde:pdenonlin:SmallStepSize'))
    end
    % update u based on line search step size
    u=u+alfa*cor;
    [rr,pdenonK]=getResidual(u);

    r=rr;
    Knorm=pdenonK;
    if strcmp(self.SolverOptions.ReportStatistics,'on')
        fprintf('%4i%20.4e%12.7f\n',iter,nrm,alfa);
    end
end % END  - Gauss-Newton iteration loop

if self.ParentModel.IsTwoD
    u=full(u(1:np*N));
else
    u = -u(1:femodel.numEdges);
    self.PDESystemSize = 1;
end
end

function [K, M, F, Q, G, H, R] = assemblePdeMatrices(thePde, p, t, coefstruct, u, femodel)
% only recompute potentially nonlinear quantities, otherwise use matrices
% passed in from original computation of femodel
femat = thePde.assembleSelectedFEMatrices(p, t, coefstruct, u, [], 'KFA', false);
K = femat.K;
F = femat.F;
M = femat.A;
Q = femodel.Q;
G = femodel.G;
H = femodel.H;
R = femodel.R;
if startsWith(thePde.ParentModel.AnalysisType,'magnetostatic','IgnoreCase',true) && ...
        isfield(coefstruct.Coefficients,'fc')
    % compute magnetization if present
    curlF = addCurlComponent(thePde,p,t,coefstruct,u); 
    F = F + curlF;
end
end

function [r, pdenonK, pdenonF, u]=compResidual2D(u, pdenonTempPde, pdenonp, pdenont, pdenoncoef, pdenonn, femodel)
% Residual for 2D nonlinear solver

u = full(u);
[K, M, F, Q, G, H, R] = assemblePdeMatrices(pdenonTempPde, pdenonp, pdenont, pdenoncoef, u(1:pdenonn), femodel);

pdenonK=K+M+Q;
pdenonF=F+G;
if any(any(H))
    if size(u,1)<(size(pdenonK,2)+size(H,1))
        u=[u;H'\(pdenonF-pdenonK*u)];
    end
    pdenonK=[pdenonK H';H zeros(size(H,1))];
    pdenonF=[pdenonF;R];
end
r=pdenonK*u-pdenonF;

end

function [r,pdenonK, pdenonF, u]=compResidual3D(u, pdenonTempPde, pdenonp, pdenont, pdenoncoef, ~, femodel)
% Residual for 3D nonlinear solver

u = full(u);
if size(u,2) == 1
    % if u has edge dofs, extrap to edges. 1:numEdges are the edge dofs
    pI = pdeInterpolant(pdenonp,pdenont,u(1:femodel.numEdges),femodel.edgeElem);
    [Xvals,Yvals,Zvals] = pI.extrapolateEdgeToNodes();
    uNodal = [Xvals;Yvals;Zvals];
else
    uNodal = u;
end
% set up new discretization matrices
[~,K,F,~,~] = discretizedModelNedelec(pdenonTempPde,pdenoncoef,femodel.edges,femodel.edgeElem,femodel.numEdges,uNodal,true);
M = femodel.M;
H = femodel.H;
R = femodel.R;
numNodalDOF = size(pdenonp,2);
faceIDs = 1:pdenonTempPde.ParentModel.Geometry.NumFaces;
nBC = findNodes(pdenonTempPde.Mesh,'region','Face',faceIDs(femodel.facesWithBC));
iS = 1:numel(nBC);
jS = nBC;
v = ones(size(nBC));
C = sparse(iS,jS,v,numel(nBC), numNodalDOF);
width = size(K,2) + size(H,1) + size(M,2) + size(C,1);
if size(u,2) ~= 1
    u = ones(width,1);
end
% zero block matrices
Z14 = sparse(size(K,1),size(C,1));
Z22_4 = sparse(size(H,1),width-size(H,2));
Z32_3 = sparse(size(M,2),size(H,1)+size(M,2));
Z41_2 = sparse(size(C,1),size(K,2) + size(H,1));
Z44 = sparse(size(C,1),size(C,1));
% define K + constraints
% -K and -H for the signs to work out with constraints, undone when writing
% solution out at end of iteration
pdenonK = [-K -H' M Z14 ; 
           -H Z22_4 ;
           M' Z32_3 C' ;
           Z41_2 C Z44];
pdenonF = [F;R;zeros(size(M,2),1);zeros(size(C,1),1)];

r = pdenonK*u - pdenonF;

end

function nrm = resNorm(r, self, K)
if ischar(self.SolverOptions.ResidualNorm) ...
        && strcmp(self.SolverOptions.ResidualNorm, 'energy') % energy norm
    nrm=sqrt(abs(r'*K*r));
else
    nrm=norm(r,self.SolverOptions.ResidualNorm)/length(r)^(1/self.SolverOptions.ResidualNorm);
end
end


function thresh = calcThresh(u, numDofs)
nu = size(u,1);
thresh=zeros(nu,1);
threshFact = 1e-6;
if all(u(1:numDofs)==0)
    thresh(1:numDofs)=sqrt(eps);
else
    thresh(1:numDofs)=threshFact*max(abs(u(1:numDofs)));
end
if numDofs<nu
    if all(u(numDofs+1:nu)==0)
        thresh(numDofs+1:nu)=sqrt(eps);
    else
        thresh(numDofs+1:nu)=threshFact*max(abs(u(numDofs+1:nu)));
    end
end
end
