function varargout = Hmode_gradient(varargin)
%#codegen
% Hmode_gradient  Hmode module specifying required gradients at edge.
% Calculates the heat flux in the pedestal from the power density,
% and calculates the associated chi_i|e from the requested gradient
% assuming diffusive transport. For particle transport, assumes pinch +
% diffusion with a tweakable vneped. Blends the transport given by
% transport model and requested transport together dependent on
% 'pp.transition'
%
% Options:
% pp.transition   Transition between requested and original transport
%                 0: hard, instantaneous transpot change [default]
%                 1: linear, linearly go from 100% original to requested
%                 2: smooth, use a sigmoid to blend original to requested
% pp.wsmooth      Width of transition region (for 1, 2) [default: 0.05]
% pp.vneped       Assumed pinch for particle transport [default: -1]
% pp.active       Used (outside the module itself) to activate/deactivate it.
% pp.doplot       Plot debugging plots of transport blend in pedestal [default: false]
%
% Example calls: 
%  chi = Hmode_gradient(stap, geop, trap, model, pp, species, equation, chi0, dchi0_dx, dchi0_dxdot, dchi0_du)
%  [chi, dchi_dx, dchi_dxdot, dchi_du] = Hmode_gradient(stap, geop, trap, model, pp, species, equation, chi0, dchi0_dx, dchi0_dxdot, dchi0_du)
%  [dne, vne] = Hmode_gradient(stap, geop, trap, model, pp, species, equation, dne0, ddne0_dx, vne0, dvne0_dx)
%  [dne, vne, ddne_dx, dvne_dx] = Hmode_gradient(stap, geop, trap, model, pp, species, equation, dne0, ddne0_dx, vne0, dvne0_dx)
%
% See also hmode

if nargin == 0,
  %% DEFAULT PARAMETERS
  default_params = struct( ...
    'active',true,...
    'wsmooth',0.05,...
    'vneped',-1, ...
    'transition',0,...
    'doplot',false...
    );
  
  varargout{1} = default_params;
  return %empty call, probably to get default structures
else
  stap = varargin{1};
  geop = varargin{2};
  trap   = varargin{3};
  
  model = varargin{4};
  pp = varargin{5}; % distribute inputs
  species = varargin{6}; % choose 'i' or 'e'
  equation = varargin{7};
end

switch equation
  case 'temperature'
    chi0 = varargin{8};
    dchi0_dx = varargin{9};
    dchi0_dxdot = varargin{10};
    dchi0_du = varargin{11};
    if nargout == 1
      varargout{1} = Hmode_te_gradient(stap,geop,trap,model,pp,species,chi0,dchi0_dx);
    elseif nargout==4
      [varargout{1},varargout{2},varargout{3},varargout{4}] = Hmode_te_gradient(stap,geop,trap,model,pp,species,chi0,dchi0_dx,dchi0_dxdot,dchi0_du);
    else
      error('something is wrong')
    end
    
  case 'density'
    dne0 = varargin{8};
    ddne0_dx = varargin{9};
    vne0 = varargin{10};
    dvne0_dx = varargin{11};
    if nargout==2
      [varargout{1},varargout{2}] = Hmode_ne_gradient(stap,geop,trap,model,pp,species,dne0,ddne0_dx,vne0,dvne0_dx);
    elseif nargout==4
      [varargout{1},varargout{2},varargout{3},varargout{4}] = ...
        Hmode_ne_gradient(stap,geop,trap,model,pp,species,dne0,ddne0_dx,vne0,dvne0_dx);
    else
      error('something is wrong');
    end
end
end

function varargout = Hmode_te_gradient(stap,geop,trap,model,pp,species,chi0,dchi0_dx,dchi0_dxdot,dchi0_du)
% #codegen
coder.extrinsic('debug_plot_te');

switch species
  case 'e'
    nx     = stap.ne;
    tx = stap.te; txp = stap.tep;
    dnx_dx = stap.dne_dx;
    pe     = trap.petot;
    dpe_dx     = trap.dpetot_dx;
    dpe_dxdot  = trap.dpetot_dxdot;
    dpe_du     = trap.dpetot_du;
    
    gtx_target = stap.gteped;
    
  case 'i'
    nx     = stap.ni;
        tx = stap.ti; txp=stap.tip;

    dnx_dx = stap.dni_dx;
    pe      = trap.pitot;
    dpe_dx     = trap.dpitot_dx;
    dpe_dxdot  = trap.dpitot_dxdot;
    dpe_du     = trap.dpitot_du;
    
    gtx_target = stap.gtiped;
  otherwise
    error('invalid species')
end

qelectron = 1.6e-19;


%%
if ~isempty(model.hmode.vind) && stap.lamH > 1e-3; % tolerance for filtered v
  inHmode = true;
else
  inHmode = false;
end

%%
qe = zeros(model.rgrid.nrhogauss,1);

if inHmode
  assert(gtx_target<0,'RAPTOR:HmodeGradientError','Hmode gradient request must be negative');
   
  frho = transition(stap,pp,model);
  
  iped = frho>0;
  % heat flux
  qe(iped) = int_gauss(pe.*geop.Vp,model,iped);
  % chie = chie_adapt*f(rho) + chie_original*(1-f(rho));
  chiadapt = -(geop.Vp.*qe./(qelectron*geop.g1.*nx.*gtx_target));
  chi = chiadapt.*frho + chi0.*(1-frho);
else % no Hmode
  % no chie modification
  frho = zeros(model.rgrid.nrhogauss,1);
  iped = false(model.rgrid.nrhogauss,1);
  chi = chi0;
  chiadapt = zeros(model.rgrid.nrhogauss,1);
end

if pp.doplot && inHmode
  debug_plot_te(stap,model,species,tx,txp,gtx_target,qe,chi0,chiadapt,chi)
end

varargout{1} = chi;

%%
if nargout>1
  if inHmode
    dqe_dx    = zeros(model.rgrid.nrhogauss,model.dims.nx);
    dqe_dxdot = zeros(model.rgrid.nrhogauss,model.dims.nx);
    dqe_du    = zeros(model.rgrid.nrhogauss,model.dims.nu);
    
    dqe_dx(iped,:)    = int_gauss(bsxfun(@times,geop.Vp,dpe_dx),model,iped);
    dqe_dxdot(iped,:) = int_gauss(bsxfun(@times,geop.Vp,dpe_dxdot),model,iped);
    dqe_du(iped,:)    = int_gauss(bsxfun(@times,geop.Vp,dpe_du),model,iped);  
    
    chiaoqe = zeros(model.rgrid.nrhogauss,1);
    chiaoqe(iped) = chiadapt(iped)./qe(iped);
    dchiadapt_dx    = bsxfun(@times,chiaoqe,dqe_dx) + ...
      bsxfun(@times,-chiadapt./nx,dnx_dx);
    
    dchi_dx = bsxfun(@times,frho,dchiadapt_dx) + ...
      bsxfun(@times,(1-frho),dchi0_dx);
    
    dchi_dxdot = bsxfun(@times,frho.*chiaoqe,dqe_dxdot) + ...
      bsxfun(@times,(1-frho),dchi0_dxdot);
    
    dchi_du = bsxfun(@times,frho.*chiaoqe,dqe_du) + ...
      bsxfun(@times,(1-frho),dchi0_du);
    
  else
    dchi_dx = dchi0_dx;
    dchi_dxdot = dchi0_dxdot;
    dchi_du = dchi0_du;
  end
  
  varargout{2} = dchi_dx;
  varargout{3} = dchi_dxdot;
  varargout{4} = dchi_du;
end
end

function varargout = Hmode_ne_gradient(stap,geop,trap,model,pp,species,dne0,ddne0_dx,vne0,dvne0_dx)
% #codegen

coder.extrinsic('debug_plot_ne');

if nargout>2
  calcder=true;
else
  calcder=false;
end

switch species
  case 'e'
    nx     = stap.ne;
    nxp    = stap.nep;
    dnx_dx = stap.dne_dx;
    
    snx     = trap.sne;
    dsnx_dx     = trap.dsne_dx;
    
    nxp_target = stap.gneped;
  case 'i'
    error('not done yet')
  otherwise
    error('invalid species')
end

%%
if ~isempty(model.hmode.vind) && stap.lamH > 1e-3; % tolerance for filtered v
  inHmode = true;
else
  inHmode = false;
end

%%
if inHmode
  vneped = pp.vneped; % fixed vne in pedestal region
  frho = transition(stap,pp,model);
  iped = frho>0;
  
  % particle flux
  Gammax = zeros(model.rgrid.nrhogauss,1);
  Gammax(iped) = int_gauss(snx.*geop.Vp,model,iped);
  
  vne = vneped.*frho + vne0.*(1-frho);
  vneadapt = vneped*ones(size(vne0));
  % solve particle balance equation for De
  % Gamma = -g1/Vp*De*nep + g0*vep*ne
  kk = iped.*(geop.Vp./(geop.g1.*nxp_target));
  dneadapt = kk.*(geop.g0.*(vne.*nx) - Gammax);
  
  dne = dneadapt.*frho + dne0.*(1-frho);
  
else % no Hmode
  dneadapt = zeros(size(dne0)); % for plot
  vneadapt = zeros(size(vne0)); % for plots
  dne = dne0;
  vne = vne0;
  % no modifications
end

if pp.doplot && inHmode
  debug_plot_ne(stap,model,species,nx,nxp,nxp_target,Gammax,dne0,dneadapt,dne,vne0,vneadapt,vne)
end

varargout{1} = dne;
varargout{2} = vne;

%%
if calcder
  if inHmode
    %dnxped_dx = interp1(model.rgrid.rhogauss,dnx_dx,pp.rhoped);
    
    dGamma_dx = zeros(model.rgrid.nrhogauss,model.dims.nx);
    dGamma_dx(iped,:) = int_gauss(bsxfun(@times,geop.Vp,dsnx_dx),model,iped);
    dvne_dx   = bsxfun(@times,(1-frho),dvne0_dx);
    dvnenx_dx = bsxfun(@times,nx,dvne_dx) + bsxfun(@times,vne,dnx_dx); % d(ve*nx)/dx
    
    ddneadapt_dx = bsxfun(@times,kk,bsxfun(@times,geop.g0,dvnenx_dx)-dGamma_dx);
    
    ddne_dx = bsxfun(@times,frho,ddneadapt_dx) + ...
      bsxfun(@times,(1-frho),ddne0_dx);
    
  else
    ddne_dx = ddne0_dx;
    dvne_dx = dvne0_dx;
    % pass unperturbed
  end
  varargout{3} = ddne_dx;
  varargout{4} = dvne_dx;
end

end

function frho = transition(stap,pp,model)
switch pp.transition
  case 0; %'hard'
    frho = zeros(size(model.rgrid.rhogauss));
    frho(model.rgrid.rhogauss > stap.rhoped) = 1;
  case 1 %'linear'
    frho = max(0,min((1)/(pp.wsmooth)*(model.rgrid.rhogauss-stap.rhoped),1)+0);
  case 2 %'smooth'
    frho = 1./(1.+exp((stap.rhoped+pp.wsmooth/2 - model.rgrid.rhogauss)./(pp.wsmooth/2)));
  otherwise
    error('unknown transition type')
end
end

function debug_plot_te(stap,model,species,tx,txp,gtx_target,qe,chi0,chiadapt,chi)

figname = 'Hmode energy equation';
hf = findobj('name',figname);
if isempty(hf)
  hf = figure('name',figname);
else
end
if strcmp(species,'e');
  ii=1;
else
  ii=2;
end
iped = model.rgrid.rhogauss>stap.rhoped;
% Gradient is requested, not t directly, so assume t = dt_target/drho * rho + 0
ax=subplot(4,2,ii,'parent',hf); 
hp=plot(ax,model.rgrid.rhogauss,tx,model.rgrid.rhogauss(iped),(model.rgrid.rhogauss(iped)-1)*gtx_target,'k--');
title(ax,sprintf('%s',species)); ylabel(ax,'temperature')
hleg=legend(hp,{'profile','target gradient'}); set(hleg,'box','off')
ax=subplot(4,2,ii+2,'parent',hf); 
plot(ax,model.rgrid.rhogauss,txp,'r',model.rgrid.rhogauss(iped),gtx_target*ones(sum(iped),1),'k--');
ylabel(ax,'dtemp/drho');


ax=subplot(4,2,ii+4,'parent',hf); plot(ax,model.rgrid.rhogauss,qe); ylabel(ax,'flux');
ax=subplot(4,2,ii+6,'parent',hf); plot(ax,model.rgrid.rhogauss,[chi0,chiadapt,chi]); ylabel(ax,'chi');
legend(ax,{'orig','patch','modif'},'location','north','box','off','orientation','hor')
drawnow;
end

function debug_plot_ne(stap,model,species,nx,nxp,nxp_target,Gammax,dne0,dneadapt,dne,vne0,vneadapt,vne)
  figname = 'Pedestal fix particle equation';
  hf = findobj('name',figname);
  if isempty(hf)
    hf = figure('name',figname);
  else
  end
  if strcmp(species,'e');
    ii=1;
  else
    ii=2;
  end
  iped = model.rgrid.rhogauss>stap.rhoped;
  % Gradient is requested, not n directly, so assume n = dn_target/drho * rho + 0
  ax=subplot(5,2,ii,'parent',hf); plot(ax,model.rgrid.rhogauss,nx,model.rgrid.rhogauss(iped),(model.rgrid.rhogauss(iped)-1)*nxp_target,'k--'); 
  title(ax,sprintf('%s',species)); ylabel(ax,'density')
  ax=subplot(5,2,ii+2,'parent',hf); plot(ax,model.rgrid.rhogauss,nxp,'r',model.rgrid.rhogauss(iped),nxp_target*ones(sum(iped),1),'k--');
  ylabel(ax,'dn/drho')
  ax=subplot(5,2,ii+4,'parent',hf); plot(ax,model.rgrid.rhogauss,Gammax); ylabel(ax,'Particle flux');
  ax=subplot(5,2,ii+6,'parent',hf); plot(ax,model.rgrid.rhogauss,[dne0,dneadapt,dne]); ylabel(ax,'D_{ne}');
  ax=subplot(5,2,ii+8,'parent',hf); plot(ax,model.rgrid.rhogauss,[vne0,vneadapt,vne]); ylabel(ax,'V_{ne}');
  legend(ax,{'orig','patch','modif'},'location','north','box','off','orientation','hor')
  drawnow; 
end
