function h = meqplotfancy(L,LY,varargin)
%MEQPLOTFANCY Fancy plot a single timeslice with coils and vv currents
%
% h = meqplotfancy(L,LY,...)
% grid than x.
% parameter-value pairs:
%       'hax': axis handle for plot re-use between time slices
%       'plotFz': flag to plot flux surfaces outside grid (default: true)
%             need to run with izgrid=true for this
%       'vacuum',     if true plot the conductor-only flux as well, needs ivacuum=true;
%       'plotArrows', if true plot also the vacuum field lines (needs ivacuum=ifield=true)
%       'openFieldlineStyle', style for open field lines
%       'openFieldlineColor', color for open field lines
%
% [+MEQ MatlabEQuilibrium Toolbox+] Swiss Plasma Center EPFL Lausanne 2022. All rights reserved.

p=inputParser;
p.addParameter('parent',gca,@(x) ishandle(x));
p.addParameter('plotFz',true,@(x) islogical(x)); % plot on extended RZ grid
p.addParameter('openFieldlineStyle','-');
p.addParameter('openFieldlineColor',0.5*[1 1 1]);
p.addParameter('plasma',true,@(x) islogical(x)); % plot plasma flux and shape
p.addParameter('vacuum',false,@(x) islogical(x)); % plot vacuum contribution
p.addParameter('plotArrows',false); % plot arrows indicating vacuum field direction
p.addParameter('FxLevels',[]);
p.addParameter('Ivlim',500);
p.addParameter('fillPlasma',true);

parse(p,varargin{:}); P = p.Results;

hax = P.parent;
hold(hax,'on');
%% Find old axis handles to reuse or delete them
hves    = flipud(findobj(hax,'Tag','vessel')); % LIFO
hcoils  = flipud(findobj(hax,'Tag','coils')); % LIFO
hoco    = findobj(hax,'Tag','opencontours');
hlcfs   = findobj(hax,'Tag','LCFS');
hlim    = findobj(hax,'Tag','limiter');
hxpts   = findobj(hax,'Tag','xpoints');
hplasma = findobj(hax,'Tag','plasma');

if nargin == 1; LY = []; end % empty default LY

if ~isempty(LY)
  assert(numel(LY.t)==1,'can plot only one time slice, use meqxk() to select one slice')
end

hasvessel = ~isequal(L.P.selu,'n');

%% Parameters
assert(L.P.izgrid~=0,'Must set L.P.izgrid=true for fancy plot');

% plasma or no plasma
noplasma = isempty(LY) || ~any(LY.lB) || ~any(LY.Iy(:));

%% Parameters
assert(logical(L.P.izgrid),'Must set L.P.izgrid=true for fancy plot');

iselc = find(any(L.G.Tca,2)); % select nonzero rows of Tca;
nc = numel(iselc);

%% Data
if ~isempty(LY) % if no LY, just plot geometry
  if ~noplasma
    if LY.nA>1 || LY.nB>1
      sIp = sign(LY.Ip);
      FB = sIp*min(sIp*LY.FB); FA = sIp*max(sIp*LY.FA);
    else
      FB = LY.FB(1); FA = LY.FA(1);
    end
    
    assert(~any(isnan(LY.Fx(:))),'NaNs in Fx, can not plot')
    FxN = (LY.Fx-FB)/(FA-FB);
    FzN = (LY.Fz-FB)/(FA-FB);
    
    if isempty(P.FxLevels) || isscalar(P.FxLevels)
      if isempty(P.FxLevels), n = 40;
      else,                   n = P.FxLevels;
      end
      Fc = linspace(-3,1,n); Fc = Fc(1:end-1);
    else % direct specification
      Fc = P.FxLevels;
    end
    
    % choose on which vacuum flux to plot
    if P.plotFz
      rr = L.rz; zz=L.zz;
      FluxN = FzN;
    else
      rr = L.rx; zz=L.zx;
      FluxN = FxN;
    end
    
    if isempty(hoco)
      contour(rr,zz,FluxN,Fc,...
        'Color',[0.5 0.5 0.5],'parent',hax,'Tag','opencontours');
    else
      hoco.Visible = 'on';
      hoco.ZData = FluxN;
      hoco.LevelList = Fc;
    end
  end
  
  %% X points
  if all(isfield(LY,{'rX','zX'}))
    if isempty(hxpts)
      plot(hax,LY.rX,LY.zX,'xr','Tag','xpoints');
    else
      hxpts.XData = LY.rX; hxpts.YData=LY.zX;
    end
  end
  
end

%% structures
if isempty(hcoils)
  % plot only at init, otherwise just change colors
  
  %% plot coils etc
  % add width for G coils
  L.G.wc(isnan(L.G.wc))=0.02;
  L.G.hc(isnan(L.G.hc))=0.02;
  
  for ic=1:nc
    icoil = iselc(ic);
    xx = L.G.rc(icoil) + L.G.wc(icoil)*[-0.5 0.5 0.5 -0.5];
    yy = L.G.zc(icoil) + L.G.hc(icoil)*[-0.5 -0.5 0.5 0.5];
    hcoils(ic) = patch(xx,yy,'r','Parent',hax,'Tag','coils');
  end
end

if isempty(hlim)
  %% limiter
  % Patch to exclude ugly flux outside limiter
  if ~P.plotFz
    maxrx = max(L.G.rx); minrx = min(L.G.rx);
    minzx = min(L.G.zx); maxzx = max(L.G.zx);
    
    rr = [L.G.rl;minrx;minrx;maxrx;maxrx;minrx];
    zz = [L.G.zl;0    ;maxzx;maxzx;minzx;minzx];
    patch(rr,zz,'w','EdgeColor','None','Tag','limiter');
  end
  % limiter contour
  plot(hax,[L.G.rl;L.G.rl(1)],[L.G.zl;L.G.zl(1)],'Color',[0.4 0.4 0.4],'LineWidth',2,...
    'Tag','limiter')
end

% Vessel
if isempty(hves) && hasvessel
  % Categorize vessel segments by name
  vlabel = cell(L.G.nv,1);
  for ives = 1:L.G.nv
    inumbers = regexp(L.G.dimv{ives},'_*\d');
    vlabel{ives} = L.G.dimv{ives}(1:(inumbers(1)-1));
  end
  [vlabels,iappear] = unique(vlabel,'first');
  [~,isort] = sort(iappear);
  vlabels=vlabels(isort);
  iv = 1; % total vessel label index
  
  % plot each vessel segment separately
  for ives = vlabels'
    iives = startsWith(L.G.dimv,ives{:}); % elements with this label
    rvv = L.G.rv(iives); zvv = L.G.zv(iives);
    wvv = L.G.wv(iives); hvv = L.G.hv(iives);
    nvv = numel(rvv);
    
    for ivv=1:nvv
      % local patch shape
      xx = wvv(ivv)*[-0.5 0.5 0.5 -0.5]+rvv(ivv);
      yy = hvv(ivv)*[-0.5 -0.5 0.5 0.5]+zvv(ivv);
      
      hves(iv) = patch(xx,yy,'r','edgecolor','none',...
        'Parent',hax,'Tag','vessel'); hold on;
      iv = iv+1;
    end
  end
end

if ~isempty(hplasma), delete(hplasma); end
if ~isempty(hlcfs), delete(hlcfs); end

%% Plasma
if ~isempty(LY)
  if ~noplasma
    if P.fillPlasma
      nB = sum(~isnan(LY.FB));
      %%% Fill inner contours
      % flux contours to plot: due to contourc accuracy reasons must space many
      % points close to lcfs (=0 value)
      Fplasma = unique([logspace(-4,-1,5),linspace(0.1,1,11)]);
      switch nB
        case 1
          FxN = (LY.Fx-LY.FB)/(LY.FA-LY.FB);
          plot_plasma(hax,L,FxN,Fplasma);
          plot_LCFS(hax,L,LY.Fx-LY.FB(1));
        case 2  % 2 disjoint domains, scale colors separately
          for idom=1:2
            FxN = (LY.Fx-LY.FB(idom))/(LY.FA(idom)-LY.FB(idom));
            plot_plasma(hax,L,FxN,Fplasma,LY.rA(idom),LY.zA(idom))
            plot_LCFS(hax,L,FxN);
          end
        case 3 % doublet
          plot_plasma(hax,L,FxN,Fplasma);
          for idom=1:3
            FB = LY.FB(idom);
            plot_LCFS(hax,L,LY.Fx-FB);
          end
      end
    else
      contour(L.rx,L.zx,LY.Fx,LY.FB(end)*[1,1],'r','linewidth',2);
    end
  end
  
  %% Colors
  % current limits for color scaling
  nolim = isinf(L.G.Ialim);
  Ialim( nolim,1) = max(abs(LY.Ia(:)));
  Ialim(~nolim,1) = L.G.Ialim(~nolim);
  Iclim = diag(L.G.Tca(iselc,:)*diag(Ialim)*(L.G.Tca(iselc,:))');
  Ivlim = P.Ivlim; % vessel current limits for plotting
  Ic = L.G.Tca(iselc,:)*LY.Ia;
  
  % blue-red colormap
  nC    = 101; % number of colors
  mymap = redwhitebluemap(nC);
  levels = linspace(-1,1,nC-1); % level - to - color mapping
  
  iCcoils = ifloor(levels,Ic./Iclim,1)+1; % colormap index per coil
  for ic=1:nc
    hcoils(ic).FaceColor = mymap(iCcoils(ic),:);
  end
  
  if hasvessel
    % vessel colors
    mymap = redgreybluemap(nC);
    iCves = ifloor(levels,LY.Iv./Ivlim,1)+1;
    for iv=1:L.G.nv
      hves(iv).FaceColor = mymap(iCves(iv),:);
    end
  end
else
  for iv = 1:L.G.nv
    hves(iv).EdgeColor = 'k';
  end
end
%%
h=hax;

%% Vacuum flux surfaces
if ~isempty(LY) && P.vacuum
  plot_vacuum_FS(hax,L,LY,P)
end
axis(hax,'equal');
hold(hax,'on');

set(gca,'XLim',[min(L.G.rz),max(L.G.rz)])
set(gca,'YLim',[min(L.G.zz),max(L.G.zz)])

end

function mymap = redwhitebluemap(nC)
nc2 = ceil(nC/2);
red   = [linspace(0,1,nc2) linspace(1,1,nc2)]';
gre   = [linspace(0,1,nc2) linspace(1,0,nc2)]';
blu   = [linspace(1,1,nc2) linspace(1,0,nc2)]';
mymap = [red gre blu];
% delete central color (equal)
mymap = mymap([1:nc2,(nc2+2):(2*nc2)],:);
end


function mymap = redgreybluemap(nC)
nc2 = ceil(nC/2);
dw = 0.7; % whiteness level
red   = [linspace(0,dw,nc2) linspace(dw,1,nc2)]';
gre   = [linspace(0,dw,nc2) linspace(dw,0,nc2)]';
blu   = [linspace(1,dw,nc2) linspace(dw,0,nc2)]';
mymap = [red gre blu];
% delete central color (equal)
mymap = mymap([1:nc2,(nc2+2):(2*nc2)],:);
end

function plot_plasma(hax,L,FxN,Fplasma,rA,zA)

C=contourc(L.G.rx,L.G.zx,FxN,Fplasma);
levels = linspace(-0.1,1,64);
mymap = colormap(hax,'hot');
while ~isempty(C)
  k = C(2,1)+1;
  isclosed =  (C(:,2) == C(:,k)); % closed contours
  if isclosed % Draw only closed contours
    rc = C(1,2:k); zc = C(2,2:k); Fc = C(1,1);
    if nargin==4 || inpolygon(rA,zA,rc,zc)
      col = mymap(ifloor(levels,Fc),:);
      patch(rc,zc,col, ...
        'EdgeColor','none','UserData','P',...
        'parent',hax,'Tag','plasma');
    end
  end
  C(:,1:k) = [];
end
end
function plot_LCFS(hax,L,FxN)
% LCFS
contour(L.G.rx,L.G.zx,FxN,[0 0],'k','linewidth',2,...
  'parent',hax,'Tag','LCFS');
end

function plot_vacuum_FS(hax,L,LY,P)
if LY.Ip~=0
  assert(L.P.ivacuum,'should run with ivacuum=true to compute vacuum fields when |Ip|>0')
  Fz  = LY.F0z;
  Fx  = LY.F0x;
  if P.plotArrows
    assert(L.P.ifield,'need ifield=true to plot fields')
    Brz = LY.Br0z; Bzz = LY.Bz0z;
    Brx = LY.Br0x; Bzx = LY.Bz0x;
  end
else
  % fields/fluxes are already in vacuum
  Fz = LY.Fz;
  Fx = LY.Fx;
  if P.plotArrows
    Brz = LY.Brz; Bzz = LY.Bzz;
    Brx = LY.Brx; Bzx = LY.Bzx;
  end
end

% Determine levels
if isempty(P.FxLevels)
  dF = 0.02;
else
  dF = 1;
end

% Use Fz as general scaling
Fzmin = floor(min(Fz(:))/dF)*dF;
Fzmax = ceil( max(Fz(:))/dF)*dF;
FzN = (Fz-Fzmin)/(Fzmax-Fzmin);
FxN = (Fx-Fzmin)/(Fzmax-Fzmin);

% Choose which flux levels to plot
if isempty(P.FxLevels)
  Fc = 0:dF/(Fzmax-Fzmin):1;
else
  Fc = P.FxLevels;
end

% choose on which vacuum flux to plot
if P.plotFz
  rr = L.rz; zz=L.zz;
  FluxN = FzN;
else
  rr = L.rx; zz=L.zx;
  FluxN = FxN;
end

hvco    = findobj(hax,'Tag','vacuumcontours');
harrows = findobj(hax,'Tag','arrows');

% plot vacuum flux
if isempty(hvco)
  contour(rr,zz,FluxN,Fc,...
    'parent',hax,'Tag','vacuumcontours',...
    'LineStyle',P.openFieldlineStyle,'Color',P.openFieldlineColor);
else
  hvco.Visible = 'on';
  hvco.ZData = FluxN;
  hvco.LevelList = Fc;
end
% plot field direction arrows
if P.plotArrows
  if ~isempty(harrows); delete(harrows); end
  if P.plotFz, rr=L.rz; zz=L.zz; Br = Brz; Bz=Bzz; 
  else,        rr=L.rx; zz=L.zx; Br = Brx; Bz=Bzx; 
  end
  meqplotarrows(rr,zz,FluxN,Br,Bz,'FxLevels',Fc,'plotContour',false);
end

% Put these at the top
uistack(hvco,'top');

end