%% Electron-ion power loss
%   This module calculates the electron-ion equipartition power
%   Since ions are not modeled self-consistently at this stage, a siplified
%   model is used where there is a (prescribed) ratio between electron and
%   ion temperatures.
%

%% Default Parameters
function [varargout] = pei(varargin) %#codegen
%   params = pei                                    Gets default parameters
%   [Pei] = pei(stap,geop,model,module_params)
%   [Pei, ~, dPei_dx, ~] = pei(stap,geop,model,module_params)

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%% Input processing

if nargin==0
  varargout{1} = 'src'; % type
  varargout{2} = [];
  return
elseif nargin == 2
  module_params = struct(     ...
    'active', false ,       ... % Module is not-active as default
    'DTfraction',0    ... % nT/(nD+nT) fraction of Tritium vs Deuterium
    );
  
  mm.name = mfilename;
  varargout{1} = mm;
  varargout{2} = module_params;
  return %empty call, probably to get default structures
elseif nargin==5 % change this depending on number of inputs to module
  stap = varargin{1};
  model = varargin{4};
  pei_params = varargin{5}; % distribute inputs
else
  error('must call with 0 or 5 inputs');
end

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% Initialize
Pei       = zeros(size(model.rgrid.rhogauss));
dPei_dx   = zeros(length(model.rgrid.rhogauss),model.dims.nx);

if pei_params.active
    calcder = (nargout>=3);
    %%
    te = stap.te;
    dte_dx = stap.dte_dx;
    ne     = stap.ne;
    dne_dx = stap.dne_dx;
    ni     = stap.ni;
    dni_dx = stap.dni_dx;
    ti     = stap.ti;
    dti_dx = stap.dti_dx;
    
    % avoid negative sqrt(imag numbers)
    te = abs(te); ne = abs(ne); ni = abs(ni); ti = abs(ti);
    
    % For now, only main ions (T,D) are considered, no impurities
    sum_ions    = (1-pei_params.DTfraction)/2 + pei_params.DTfraction/3;
    [lnLam,dlnLam_dte,dlnLam_dne] = lnLambda(te,ne);
    
    tekeV = 1e-3*te;
    nuequi_over_ni = (1./(0.99e19).*sum_ions.*lnLam)./(tekeV.*sqrt(tekeV));
    nuequi      = nuequi_over_ni .* ni;
    
    kpei        = 2.*1.602e-19.*ne;
    kpeinuequi  = kpei.*nuequi;
    Pei         = kpeinuequi.*(te-ti);
    
    %% Derivatives
    if calcder
        % dPei_dx
        if derivative_is_required('te',model)
            dnuequi_dte = nuequi .* (-3./(2*te) + dlnLam_dte./lnLam);
            dPei_dte = kpeinuequi + kpei.*dnuequi_dte.*(te-ti);
            dPei_dx   = dPei_dx + bsxfun(@times,dPei_dte,dte_dx);
        end
        
        if derivative_is_required('ti',model)
            dPei_dti = -kpeinuequi;
            dPei_dx = dPei_dx + bsxfun(@times,dPei_dti,dti_dx);
        end
        
        if derivative_is_required('ne',model)
            dnuequi_dne = nuequi.*(dlnLam_dne./lnLam);
            dkpei_dne = 2.*1.602e-19;
            % Over ne
            dPei_dne = (dkpei_dne.*nuequi + kpei.*dnuequi_dne).*(te-ti);
            % Over x
            dPei_dx = dPei_dx + bsxfun(@times,dPei_dne,dne_dx);
        end
        
        if derivative_is_required('ni',model)
            dnuequi_dni = nuequi_over_ni;
            dPei_dni = kpei.*dnuequi_dni.*(te-ti);
            % Over x
            dPei_dx = dPei_dx + bsxfun(@times,dPei_dni,dni_dx);
        end

    end
end

varargout{1} = Pei;
varargout{2} = zeros(size(Pei));

if nargout>=3
    varargout{3} = dPei_dx;
    varargout{4} = zeros(size(dPei_dx));

end

end