%GSZRMEX  Solve Poisson equation
%
% This is the MATLAB equivalent implementation of libmeq/gszr.c.
% A more detailed help is available in GSZRMEX.
%
% [+MEQ MatlabEQuilibrium Toolbox+] Swiss Plasma Center EPFL Lausanne 2022. All rights reserved.
function Fx = gszrmexm(Fb,Iy,cx,cq,cr,cs,ci,co,dz)

[nz,nr] = size(Iy); nz = nz+2; nr = nr+2; nz1 = nz-1; nr1 = nr-1; nr2 = nr1-1;
p = zeros(nz1-1,nr2);
Fx = zeros(nz,nr);

i = 2:nz1; j = 2:nr1;
Fx(i,j) = repmat(cx',nz-2,1).*Iy;

i = 1:nz;
Fx(i,1 )  = Fb(i);
Fx(i,nr)  = Fb(i+nz+2*nr2);
Fx(i,2 )  = Fx(i,2  ) + ci*Fb(i);
Fx(i,nr1) = Fx(i,nr1) + co*Fb(i+nz+2*nr2);
j = 2:nr1;
Fx( 1,j) = Fb(nz+(1:2:2*nr2));
Fx(nz,j) = Fb(nz+(2:2:2*nr2));

l = nz1/2;
i = 3:2:nz1;
ip = i-1;
j = 2:nr1;
p(ip,:) = 2.*Fx(i,j);
Fx(i,j) = Fx(i+1,j)+Fx(i-1,j);
p = gsp(p,cq,cr,cs,nr1,l,ip);
Fx = gsu(Fx,p,nr1,i,ip);

lo = nz1/4; li = lo*2; id = 2; ih = 1; io = 4;
while lo > 1
  ii = id+ih;
  i = io+1:io:nz1;
  ip = i-1;
  j = 2:nr1;
  p(ip,:) = Fx(i,j)-Fx(i+ii,j)-Fx(i-ii,j);
  Fx(i,j) = Fx(i,j)-Fx(i+ih,j)-Fx(i-ih,j)+Fx(i+id,j)+Fx(i-id,j);
  p(ip,:) = p(ip,:)+Fx(i,j);
  for l = lo:li:nz1
    p = gsp(p,cq,cr,cs,nr1,l,ip);
  end
  Fx = gsu(Fx,p,nr1,i,ip);
  id = id*2; ih = ih*2; io = io*2; lo = lo/2; li = li/2;
end

ii = io; io = id;
while ih >= 1
  i = io+1:ii:nz1;
  ip = i-1;
  j = 2:nr1;
  p(ip,:) = 2*Fx(i,j)+Fx(i+id,j)+Fx(i-id,j);
  Fx(i,j) = Fx(i,j)-Fx(i+ih,j)-Fx(i-ih,j);
  for l = lo:li:nz1
    p = gsp(p,cq,cr,cs,nr1,l,ip);
  end
  Fx = gsu(Fx,p,nr1,i,ip);
  id = id/2; ih = ih/2; io = io/2; ii = ii/2; lo = lo*2; li = li*2;
end

i = 2:2:nz1;
ip = i-1;
j = 2:nr1;
p(ip,:) = 2.*Fx(i,j)+Fx(i+1,j)+Fx(i-1,j);
Fx(i,j) = 0.;
p = gsp(p,cq,cr,cs,nr1,lo,ip);
Fx = gsu(Fx,p,nr1,i,ip);

% Shift solution
Fx = Fx + (Fx([2:end end],:)-Fx([1 1:end-1],:)).*[dz;repmat(0.5*dz,nz-2,1);dz];
end
function p = gsp(p,cq,cr,cs,nr1,l,ip)
pp = p(ip,nr1-1);
for jj = nr1-2:-1:1
  pp = p(ip,jj) + cr(jj,l)*pp;
  p(ip,jj) = pp;
end
pp = cq(1,l)*p(ip,1);
p(ip,1) = pp;
for jj = 2:nr1-1
  pp = cq(jj,l)*p(ip,jj) + cs(jj,l)*pp;
  p(ip,jj) = pp;
end
end

function Fx = gsu(Fx,p,nr1,i,ip)
jj = 2:nr1;
Fx(i,jj) = Fx(i,jj) + p(ip,:);
end