function varargout = qnze(x,g,v,model,gauss,varargin)
% #codegen
% function [o1name,o2name,o1,o2,do1_dx,do2_dx] = qnze(x,g,v,model,gauss)
%
% Determines two unknown quantities out of the set ne, ni, n1, n2, n3, ze
% by solving equations for quasineutrality and Zeff.
%
% Which species we solve for is automatically detected from e.g. model.nxx.method
%
% o1name, o2name: string names of outputs
% o1,o2: values of outputs
% do1_dx,do2_dx: jacobians of outputs w.r.t. states
%
% x,g,v: RAPTOR vectors
% model: RAPTOR model
% gauss: true or false to get data on gauss grid or not
%

% Equations for quasineutrality and zeff
% ne = sum(n_j*Z_j) => -ne + sum(n_j*Z_j) = 0
% Zeff = sum(n_j*Z_j^2/ne) => -Zeff*ne + sum(n_j*Z_j^2) = 0
% model.ze.method = 'direct': known Zeff => system of two linear equations
% model.ze.method = 'qnze': unknown Zeff => solve quasineurality eq., then for Zeff

%% Define known/unknown species
% Initialize right-hand side

if ~isempty(varargin{1})
    prime = varargin{1}; % flag to compute derivatives of profile instead
    assert(any(prime==[0,1,2])) % only 0,1,2nd derivative supported
else
    prime = 0; % default: ask value
end

if nargout==6
    der = true; % flag to compute derivatives w.r.t x or not
else
    der = false;
end

% Loads ADAS data used to compute ze for partial ionisation
switch model.prad.modeltype
    case 'adas'
        tabmat = Lz_zave_HDT0_default;

        fn = fieldnames(tabmat);

        for k=1:numel(fn)
            if(tabmat.(fn{k}).Z == model.atom.Z1)
               ind1 = k;
            elseif(tabmat.(fn{k}).Z == model.atom.Z2)
               ind2 = k;
            elseif(tabmat.(fn{k}).Z == model.atom.Z3)
               ind3 = k;
            end
         end

        assert(~isnan(ind1) && ~isnan(ind2) && ~isnan(ind3),...
	       'No ADAS data for one or several impurities')
end
	  
namestr = (['  ';'  ']); % init empty

if gauss
    profiles = zeros(model.rgrid.nrhogauss,2); % storage for profiles
    njZj2  = zeros(model.rgrid.nrhogauss,1); % right hand side for zeff equation
    njZj   = zeros(model.rgrid.nrhogauss,1); % right hand side for qn equation
    zis = zeros(model.rgrid.nrhogauss,2); % vector to store charges for unknown species
else
    profiles = zeros(model.rgrid.nrho,2); % storage for profiles
    njZj2  = zeros(model.rgrid.nrho,1); % right hand side for zeff equation
    njZj   = zeros(model.rgrid.nrho,1); % right hand side for qn equation
    zis = zeros(model.rgrid.nrho,2); % vector to store charges for unknown species
end
dnjZj_dx  = zeros(size(njZj ,1),numel(x));
dnjZj2_dx = zeros(size(njZj2,1),numel(x));

% Index
ip = 1;

% Ion density)
zj = model.atom.Zi;
switch model.ni.method
    case 'qnze'
        zis(:,ip) = zj;
        % Parameter name
        namestr(ip,:) = 'ni';
        % Increase index for next unknown parameter
        ip = ip + 1;
    otherwise
        if prime == 0
            [nj,dnj_dx] = eval_ni(x,g,v,model,gauss);
        elseif prime == 1
            [nj,dnj_dx] = eval_nip(x,g,v,model,gauss);
        elseif prime == 2
            [nj,dnj_dx] = eval_nipp(x,g,v,model,gauss);
        end
        [njZj,njZj2] = known_ion(njZj,njZj2,nj,zj);
        if der
            dnjZj_dx  = dnjZj_dx  + zj*dnj_dx;
            dnjZj2_dx = dnjZj2_dx + zj.^2*dnj_dx;
        end
end

% Minor species N1 density
% When ADAS radiation model is used, the averaged
% charge state is considered instead of full ionisation

switch model.prad.modeltype
    case 'adas'
        te = eval_te(x,g,v,model,gauss);
        % Zj is now an array (depends on Te -> should also include corresponding derivatives)
	  zj=interp1(tabmat.(fn{ind1}).data(:,1),tabmat.(fn{ind1}).data(:,3),te/1e3);
    otherwise
        zj = model.atom.Z1;
end

switch model.n1.method
    case 'qnze'
        zis(:,ip) = zj;
        % Parameter name
        namestr(ip,:) = 'n1';
        % Increase index for next unknown parameter
        ip = ip + 1;
    otherwise
        if ~der
            nj = eval_nx(1,x,g,v,model,gauss,prime);
            [njZj,njZj2] = known_ion(njZj,njZj2,nj,zj);
        else
            [nj, dnj_dx] = eval_nx(1,x,g,v,model,gauss,prime);
            [njZj,njZj2] = known_ion(njZj,njZj2,nj,zj);
            dnjZj_dx  = dnjZj_dx  + bsxfun(@times,zj,dnj_dx);
            dnjZj2_dx = dnjZj2_dx + bsxfun(@times,zj.^2,dnj_dx);
        end
end

% Minor species N2 density

switch model.prad.modeltype
    case 'adas'
        % Zj is now an array (depends on Te -> should also include corresponding derivatives)
        zj=interp1(tabmat.(fn{ind2}).data(:,1),tabmat.(fn{ind2}).data(:,3),te/1e3);
    otherwise
        zj = model.atom.Z2;
end

switch model.n2.method
    case 'qnze'
        zis(:,ip) = zj;
        % Parameter name
        namestr(ip,:) = 'n2';
        % Increase index for next unknown parameter
        ip = ip + 1;
    otherwise
        if ~der
            nj = eval_nx(2,x,g,v,model,gauss,prime);
            [njZj,njZj2] = known_ion(njZj,njZj2,nj,zj);
        else
            [nj, dnj_dx] = eval_nx(2,x,g,v,model,gauss,prime);
            [njZj,njZj2] = known_ion(njZj,njZj2,nj,zj);
            dnjZj_dx  = dnjZj_dx  + bsxfun(@times,zj,dnj_dx);
            dnjZj2_dx = dnjZj2_dx + bsxfun(@times,zj.^2,dnj_dx);
        end
end

% Minor species N3 density

switch model.prad.modeltype
    case 'adas'
        % Zj is now an array (depends on Te -> should also include corresponding derivatives)
        zj=interp1(tabmat.(fn{ind3}).data(:,1),tabmat.(fn{ind3}).data(:,3),te/1e3);        
    otherwise
        zj = model.atom.Z3;
end

switch model.n3.method
    case 'qnze'
        zis(:,ip) = zj;
        % Parameter name
        namestr(ip,:) = 'n3';
        % Increase index for next unknown parameter
        ip = ip + 1;
    otherwise
        if ~der
            nj = eval_nx(3,x,g,v,model,gauss,prime);
            [njZj,njZj2] = known_ion(njZj,njZj2,nj,zj);
        else
            [nj, dnj_dx] = eval_nx(3,x,g,v,model,gauss,prime);
            [njZj,njZj2] = known_ion(njZj,njZj2,nj,zj);
            dnjZj_dx  = dnjZj_dx  + bsxfun(@times,zj,dnj_dx);
            dnjZj2_dx = dnjZj2_dx + bsxfun(@times,zj.^2,dnj_dx);
        end
end

assert(ip<=3,'Too many unknown species. This can not happen, please check')

%% Solve for unknowns
switch model.ze.method
    case 'direct'
        % Zeff is known
       ze = eval_ze(x,g,v,model,gauss,prime);

        % Electron density
        switch model.ne.method
            case 'qnze' % ze='direct' and ne='qnze'
                % electrons are one of the unknown species, with another
                % ion species nx.
                % Substitute nx from quasineutrality into Zeff equation,
                % this gives an equation for ne, then fill into qn again to
                % get nx.
                zx = zis(:,1); % charge of the nx ion species
                ne = (njZj2 + zx.*njZj)./(ze+zx);
                nx = (njZj-ne) ./ zx;
                
                if der
                    dne_dx = bsxfun(@times,1./(ze+zx),dnjZj2_dx + zx.*dnjZj_dx);
                    dnx_dx = (dnjZj_dx - dne_dx)./zx;
                end
                
                % Output name
                namestr(ip,:) = 'ne';
                % Output values
                profiles(:,1) = nx;
                profiles(:,2) = ne;
                
                if der
                    varargout{5} = dnx_dx;
                    varargout{6} = dne_dx;
                end
            otherwise % ze='direct' and ne is not 'qnze'
                % then solve for two unknown ion species
                
                % Output names were already filled in earlier
                
                % Output values
                if ~der % no derivatives
                    [na,nb] = solve_twoions(ze,zis,njZj,njZj2,dnjZj_dx,dnjZj2_dx,...
                        x,g,v,model,gauss,prime);
                    profiles(:,1) = na;
                    profiles(:,2) = nb;
                else % also state derivatives
                    [na,nb,dna_dx,dnb_dx] = ...
                        solve_twoions(ze,zis,njZj,njZj2,dnjZj_dx,dnjZj2_dx,...
                        x,g,v,model,gauss,prime);
                    profiles(:,1) = na;
                    profiles(:,2) = nb;
                    varargout{5} = dna_dx;
                    varargout{6} = dnb_dx;
                end
        end
    case 'qnze'
        % Zeff is unknown and is part of the quantities to be solved for
        % Electron density
        switch model.ne.method
            case 'qnze' % ze='qnze' and ne='qnze'
                % Unknown electron density
                ne = njZj; % get ne directly from quasineutrality
                % Zeff
                ze = njZj2./ne; % zeff directly from Zeff equation
                if der
                    dne_dx = dnjZj_dx;
                    dze_dx = bsxfun(@times,1./ne,dnjZj2_dx) ...
                        + bsxfun(@times,-njZj2./ne.^2,dne_dx);
                end
                % Output names
                namestr(1,:) = 'ne';
                namestr(2,:) = 'ze';
                % Output values
                profiles(:,1) = ne;
                profiles(:,2) = ze;
                if der
                    varargout{5} = dne_dx;
                    varargout{6} = dze_dx;
                end
            otherwise % ze='qnze' and ne is not 'qnze'
                % Electron density solved (or prescribed directly)
                if ~prime
                    [ne,dne_dx] = eval_ne(x,g,v,model,gauss);
                else
                    [ne,dne_dx] = eval_nep(x,g,v,model,gauss);
                end
                zx = zis(:,1); % charge of unknown ion
                % Unknown density from quasineutrality
                nx     = (ne - njZj) ./ zx; % ne = njZj + nx*Zx
                dnx_dx = bsxfun(@rdivide,dne_dx - dnjZj_dx,zx);
                % Zeff
                zx2 = zx.^2;
                ze = (njZj2+nx.*zx2)./ne; % zeff from definition % neZe = njZj2+nx*zx^2
                dze_dx =  bsxfun(@times,dnjZj2_dx,1./ne) ...
                    + bsxfun(@times,dnx_dx,zx2./ne) ...
                    - bsxfun(@times,ze./ne,dne_dx);
                
                % Output names
                namestr(2,:) = 'ze';
                % Output values
                profiles(:,1) = nx;
                profiles(:,2) = ze;
                
                % Jacobians
                if der
                    varargout{5} = dnx_dx;
                    varargout{6} = dze_dx;
                end
        end
end

varargout{1} = namestr(1,:);
varargout{2} = namestr(2,:);
varargout{3} = profiles(:,1);
varargout{4} = profiles(:,2);

return

function [na,nb,dna_dx,dnb_dx] = solve_twoions(ze,zis,njZj,njZj2,...
    dnjZj_dx,dnjZj2_dx,x,g,v,model,gauss,prime)
% two ion species are unknown
% update RHS with electron terms
% and solve linear system of equations
% AX=B

for jj=1:size(zis,1);
  
A = [zis(jj,:);zis(jj,:).^2];
   % split ne and other contribution
   % split done for easier calc of Jacobian
   B1 = -[njZj, njZj2]';

   if nargout==2
       if ~prime
          ne = eval_ne(x,g,v,model,gauss);
       else
          ne = eval_nep(x,g,v,model,gauss);
       end
       B2 = [ne,ze.*ne]';
       X = A\(B1+B2);
    elseif nargout==4
       dB1a_dx = -dnjZj_dx'; dB1b_dx = -dnjZj2_dx';
       if ~prime
           [ne,dne_dx] = eval_ne(x,g,v,model,gauss);
       else
           [ne,dne_dx] = eval_nep(x,g,v,model,gauss);
       end
       B2 = [ne,ze.*ne]'; % [B2a,B2b];
       dB2a_dx = dne_dx';
       dB2b_dx = bsxfun(@times,ze,dne_dx)';
       % Solve the linear system
       iA = A\eye(2); % A inverse (2x2 so easy)
       X = iA*(B1+B2);
       dBa_dx = dB1a_dx + dB2a_dx;
       dBb_dx = dB1b_dx + dB2b_dx;
    
       dna_dx = (iA(1,1)*dBa_dx + iA(1,2)*dBb_dx)';
       dnb_dx = (iA(2,1)*dBa_dx + iA(2,2)*dBb_dx)';
   end
   na(jj) = X(1,jj)'; nb(jj) = X(2,jj)';
end
return

function  [njZj,njZj2] = known_ion(njZj,njZj2,ni,zi)
njZj  = njZj  + ni.*zi;
njZj2 = njZj2 + ni.*zi.^2;
return
