function ForceBalance(M,it,norm,log,contourscale, zlims)
%ForceBalance Show the radial force balance
%   Plot the three radial forces for the given time-step it or averaged
%   over the range of time steps defined in it

if it==':'
    it=1:size(M.t2d);
end
switch nargin
    case 2
        contourscale=0.6;
        norm=false;
        log=false;
        zlims=[-inf inf];
    case 3
        log=false;
        contourscale=0.6;
        zlims=[-inf inf];
    case 4
        contourscale=0.6;
        zlims=[-inf inf];
    case 5
        zlims=[-inf inf];
    case 6
        
    otherwise
        error("Invalid number of arguments")
end
r=(M.rgrid(1:end-1)+M.rgrid(2:end))/2;
z=(M.zgrid(1:end-1)+M.zgrid(2:end))/2;
[Zc,Rc]=meshgrid(z,r);
[Z,R]=meshgrid(M.zgrid,M.rgrid);

contourcolor=[0 0 0];%[255 20 147]/255;

avgpresstens=M.Presstens(:,1:end-1,1:end-1,it);
uThet=M.fluidUTHET(1:end-1,1:end-1,it);
Ntmp=M.N(1:end-1,1:end-1,it);
dim3=size(Ntmp,3);
presstens=zeros(6,size(M.rgrid,1),size(M.zgrid,1),dim3);
fluiduThet=zeros(size(M.rgrid,1),size(M.zgrid,1),dim3);
N=zeros(size(M.rgrid,1),size(M.zgrid,1),dim3);
for j=1:size(avgpresstens,4)
    for i=1:6
        presstens(i,:,:,j)=interp2(Zc,Rc,squeeze(avgpresstens(i,:,:,j)),Z,R);
    end
    fluiduThet(:,:,j)=interp2(Zc,Rc,uThet(:,:,j),Z,R);
    N(:,:,j)=interp2(Zc,Rc,Ntmp(:,:,j),Z,R);
end
presstens(isnan(presstens))=0;
fluiduThet(isnan(fluiduThet))=0;
N(isnan(N))=0;
dr=[(M.rgrid(3)-M.rgrid(1))/2;(M.rgrid(3:end)-M.rgrid(1:end-2))/2];
dz=[(M.zgrid(3)-M.zgrid(1))/2;(M.zgrid(3:end)-M.zgrid(1:end-2))/2];
[dz, dr]=meshgrid(dz,dr);

Rinv=1./R;
%Rinv=padarray(Rinv, [1,1], 0, 'post');
Rinv(1,:)=0;

dpr=zeros(size(N,1)-1,size(N,2)-1,dim3);
dpz=zeros(size(N,1)-1,size(N,2)-1,dim3);
for j=1:size(presstens,4)
    dpr(:,:,j)=squeeze(presstens(1,2:end,1:end-1,j)-presstens(1,1:end-1,1:end-1,j))./dr;
    dpz(:,:,j)=squeeze(presstens(4,1:end-1,2:end,j)-presstens(4,1:end-1,1:end-1,j))./dz;
end
dpr=padarray(dpr, [1,1,0], 0, 'post');
dpz=padarray(dpz, [1,1,0], 0, 'post');
maxdens=max(N(:));

densitycontour=contour(M.zgrid,M.rgrid,mean(N,3),[contourscale contourscale]*maxdens);


Eforce=-M.qe*N.*M.Er(:,:,it);
Bforce=zeros(size(N,1),size(N,2),size(N,3));
inertforce=zeros(size(N,1),size(N,2),size(N,3));
Pressforce=zeros(size(N,1),size(N,2),size(N,3));
for j=1:size(N,3)
    Bforce(:,:,j)=-M.qe*N(:,:,j).*fluiduThet(:,:,j).*M.Bz';
    inertforce(:,:,j)=M.me*N(:,:,j).*fluiduThet(:,:,j).^2.*Rinv;
    Pressforce(:,:,j)=-( dpr(:,:,j)...
        + squeeze(presstens(1,:,:,j)).*Rinv + squeeze(-presstens(4,:,:,j)).*Rinv...
        + dpz(:,:,j) );
end

Fmax=max([max(Eforce(:)),max(Bforce(:)),max(inertforce(:)),max(Pressforce(:))]);
if log
    E=Eforce(Eforce~=0);
    B=Bforce(Bforce~=0);
    inert=inertforce(inertforce~=0);
    press=Pressforce(Pressforce~=0);
    Fmin=min([min(abs(E(:))),min(abs(B(:))),min(abs(inert(:))),min(abs(press(:)))]);
else
    Fmin=min([min(Eforce(:)),min(Bforce(:)),min(inertforce(:)),min(Pressforce(:))]);
end
rgridmax=sum(M.nnr(1:2));
f=figure();

ax1=subplot(4,1,1);
plotvalue(ax1,mean(Eforce,3),'Electric force density', 'F_E [Nm^{-3}]',false);

%zlim=[Fmin Fmax];
zlim=[-inf inf];

ax2=subplot(4,1,2);
plotvalue(ax2,mean(Bforce,3),'Magnetic force density','F_B [Nm^{-3}]',false,zlim)

ax3=subplot(4,1,3);
plotvalue(ax3,mean(inertforce,3),'Centrifugal force density','F_i [Nm^{-3}]',false,zlim)

ax4=subplot(4,1,4);
plotvalue(ax4,mean(Pressforce,3),'Pressure force density','F_p [Nm^{-3}]',false,zlim)

linkaxes([ax1 ax2 ax3 ax4],'xy')
sgtitle(sprintf('Radial forces t=[%1.2g-%1.2g]s n_e=%1.2gm^{-3}',M.t2d(min(it)),M.t2d(max(it)),double(maxdens)))

f.PaperOrientation='portrait';
[~, name, ~] = fileparts(M.file);
f.PaperUnits='centimeters';
papsize=[20 29];
f.PaperSize=papsize;
print(f,sprintf('%sforce_balance',name),'-dpdf','-fillpage')

%% Radial relative forces
f=figure();
old_log=log;
log=true;
ax1=subplot(4,1,1);
plotvalue(ax1,mean(Eforce./Bforce,3),'Relative Electric force density', 'F_{Er}/F_{Br} [Nm^{-3}]',false);

%zlim=[Fmin Fmax];
zlim=[-inf inf];

ax2=subplot(4,1,2);
plotvalue(ax2,mean(Bforce./Bforce,3),'Relative Magnetic force density','F_{Br}/F_{Br} [Nm^{-3}]',false)

ax3=subplot(4,1,3);
plotvalue(ax3,mean(inertforce./Bforce,3),'Relative Centrifugal force density','F_i/F_{Br} [Nm^{-3}]',false)

ax4=subplot(4,1,4);
plotvalue(ax4,mean(Pressforce./inertforce,3),'Relative Pressure force density','F_{pr}/F_{Br} [Nm^{-3}]',false)

linkaxes([ax1 ax2 ax3 ax4],'xy')
sgtitle(sprintf('Relative radial forces t=[%1.2g-%1.2g]s n_e=%1.2gm^{-3}',M.t2d(min(it)),M.t2d(max(it)),double(maxdens)))

f.PaperOrientation='portrait';
[~, name, ~] = fileparts(M.file);
f.PaperUnits='centimeters';
papsize=[20 29];
f.PaperSize=papsize;
print(f,sprintf('%s_relatforce_balance',name),'-dpdf','-fillpage')
log=old_log;

f2=figure();
ax5=subplot(2,1,1);
totforce=mean(Eforce+Bforce+inertforce+Pressforce,3);
Forcemax=cat(3,mean(Eforce,3),mean(Bforce,3),mean(inertforce,3),mean(Pressforce,3));
Forcemax=max(Forcemax,[],3);
plotvalue(ax5,totforce./Forcemax,'Total force density','F [Nm^{-3}]',false,zlims)

ax6=subplot(2,1,2);
surface(ax6,M.zgrid,M.rgrid,mean(N,3),'edgecolor','none');
hold on;
zlevel=interp2(M.zgrid, M.rgrid, mean(N,3), densitycontour(1,2:end), densitycontour(2,2:end));
plot3(ax6,densitycontour(1,2:end),densitycontour(2,2:end),zlevel,'-.','linewidth',2,'color',contourcolor)
xlim(ax6,[M.zgrid(1) M.zgrid(end)])
ylim(ax6,[M.rgrid(1) M.rgrid(rgridmax)])
colormap(ax6,'parula')
xlabel(ax6,'z [m]')
ylabel(ax6,'r [m]')
title(ax6,'Density')
c = colorbar(ax6);
c.Label.String= 'n[m^{-3}]';
view(ax6,2)

sgtitle(sprintf('Radial forces sum t=[%1.2g-%1.2g]s n_e=%1.2g m^{-3}',M.t2d(min(it)),M.t2d(max(it)),double(maxdens)))
f2.PaperOrientation='portrait';
[~, name, ~] = fileparts(M.file);
f2.PaperUnits='centimeters';
papsize=[20 29];
f2.PaperSize=papsize;
print(f2,sprintf('%sforce_dens',name),'-dpdf','-fillpage')

f3=figure();
subN=M.N(:,:,it);
t=M.t2d(it);
UR=M.fluidUR(:,:,it);
UZ=M.fluidUZ(:,:,it);
dndt=zeros(size(subN,1)-1,size(subN,2)-1,size(subN,3)-1);
ndivu=zeros(size(subN,1)-1,size(subN,2)-1,size(subN,3)-1);
ugradn=zeros(size(subN,1)-1,size(subN,2)-1,size(subN,3)-1);
for i=1:size(subN,3)-1
    dndt(:,:,i)=(subN(1:end-1,1:end-1,i+1)-subN(1:end-1,1:end-1,i))./(t(i+1)-t(i));
    
    ndivu(:,:,i)=subN(1:end-1,1:end-1,i).*(UR(2:end,1:end-1,i)-UR(1:end-1,1:end-1,i))./dr...
        + UR(1:end-1,1:end-1,i).*subN(1:end-1,1:end-1,i)./Rc...
        + subN(1:end-1,1:end-1,i).*(UZ(1:end-1,2:end,i)-UZ(1:end-1,1:end-1,i))./dz;
    
    ugradn(:,:,i)= UR(1:end-1,1:end-1,i).*(subN(2:end,1:end-1,i)-subN(1:end-1,1:end-1,i))./dr ...
        + UZ(1:end-1,1:end-1,i).*(subN(1:end-1,2:end,i)-subN(1:end-1,1:end-1,i))./dz;
end
dndt=padarray(dndt, [1,1], 0, 'post');
ndivu=padarray(ndivu, [1,1], 0, 'post');
ugradn=padarray(ugradn, [1,1], 0, 'post');

ax7=subplot(4,1,1);
plotvalue(ax7,mean(dndt,3),'\partialn/\partialt','\partialn/\partialt [m^{-3}s^{-1}]',true)


zlim=[max([min(min(mean(ndivu(2:end,:,:),3))),min(min(mean(ugradn(2:end,:,:),3)))]) min([max(max(mean(ndivu(2:end,:,:),3))),max(max(mean(ugradn(2:end,:,:),3)))])];

ax8=subplot(4,1,2);
plotvalue(ax8,mean(ndivu,3),'n\nabla(u)','n\nabla(u)[m^{-3}s^{-1}]',true,zlim)

ax9=subplot(4,1,3);
plotvalue(ax9,mean(ugradn,3),'u\nabla(n)','u\nabla(n)[m^{-3}s^{-1}]',true,zlim)

ax10=subplot(4,1,4);
plotvalue(ax10,mean(ndivu+ugradn+dndt,3),'Total','total[m^{-3}s^{-1}]',true,zlim)

sgtitle(sprintf('Continuity t=[%1.2g-%1.2g]s n_e=%1.2g m^{-3}',M.t2d(min(it)),M.t2d(max(it)),double(maxdens)))
f3.PaperOrientation='portrait';
[~, name, ~] = fileparts(M.file);
f3.PaperUnits='centimeters';
papsize=[20 19];
f3.PaperSize=papsize;
print(f3,sprintf('%scontinuity',name),'-dpdf','-fillpage')


    function plotvalue(axis,matrix,titl, clab, rm1strow, zlims)
        if nargin < 5
            rm1strow=false;
        end
        if nargin>5
            zmin=zlims(1);
            zmax=zlims(2);
        elseif log
            zmin=0;
            zmax=inf;
        else
            zmin=-inf;
            zmax=inf;
        end
        if rm1strow
            if log
                h=surface(axis,M.zgrid,M.rgrid(2:end),abs(matrix(2:end,:)));
            else
                h=surface(axis,M.zgrid,M.rgrid(2:end),matrix(2:end,:));
            end
        else
            if log
                h=surface(axis,M.zgrid,M.rgrid,abs(matrix));
            else
                h=surface(axis,M.zgrid,M.rgrid,matrix);
            end
        end
        set(h,'edgecolor','none');
        hold on;
        if log
            zpos=interp2(M.zgrid, M.rgrid, abs(matrix), densitycontour(1,2:end), densitycontour(2,2:end));
        else
            zpos=interp2(M.zgrid, M.rgrid, matrix, densitycontour(1,2:end), densitycontour(2,2:end));
        end
        plot3(axis,densitycontour(1,2:end),densitycontour(2,2:end),zpos,'-.','linewidth',2,'color',contourcolor)
        xlabel(axis,'z [m]')
        ylabel(axis,'r [m]')
        xlim(axis,[M.zgrid(1) M.zgrid(end)])
        ylim(axis,[M.rgrid(1) M.rgrid(rgridmax)])
        colormap(axis,'jet')
        c = colorbar(axis);
        c.Label.String=clab;
        caxis(axis,[zmin zmax])
        if norm
            caxis(axis,[Fmin Fmax]);
        end
        if log
            set(axis,'colorscale','log')
            caxis(axis,'auto')
        end
        title(axis,titl)
        grid(axis, 'on')
        view(axis,2)
    end
end