function [x, xdot, u, out_data, res_vec] = stationary_state_solver(model, params, init, g, v, u, ...
                                                                   parameter, inputs, verbosity, varargin)    
% Interface for use of the state equation solver for stationary state parametrizations.
% Wrapper function for calling the general state equation solver for
% the specific case of a stationary state.
uindex = 2;
x0_provided = false;
if ~isempty(varargin)
  for i= 1:numel(varargin) 
    if numel(varargin{i})==1
      uindex = varargin{i};
    elseif numel(varargin{i})>1
      x0_provided = true;
      x0 = varargin{i};
    end
  end
end

% function handles passing the case specific parametrization
switch parameter
  case 'Upl'
    x_z    = @(z) x_z_statUpl(z, model);
    xdot_z = @(z) xdot_z_statUpl(z, model, inputs);
    u_z    = @(z) u_z_statUpl(z, model, u);
  case 'Ip'
    x_z    = @(z) x_z_statIp(z, model);
    xdot_z = @(z) xdot_z_statIp(z, model, inputs);
    u_z    = @(z) u_z_statIp(z, model, init, inputs, u);
  case 'Act'
    x_z    = @(z) x_z_statPec(z, model);
    xdot_z = @(z) xdot_z_statPec(z, model, inputs);
    u_z    = @(z) u_z_statPec(z, model, inputs, u, uindex);
end
xxdotu_z = @(z) compact_args(z, x_z, xdot_z, u_z);
if x0_provided
  [x, xdot, u, out_data, res_vec] = solve_parametrized(model, params, init, g, v, xxdotu_z, verbosity, x0);
else
  [x, xdot, u, out_data, res_vec] = solve_parametrized(model, params, init, g, v, xxdotu_z, verbosity);
end
end

function [x, dx_dz, xdot, dxdot_dz, u, du_dz] = compact_args(z, x_z, xdot_z, u_z)
[x, dx_dz] = x_z(z);
[xdot, dxdot_dz] = xdot_z(z);
[u, du_dz]= u_z(z);
end

%% CASE SPECIFIC PARAMETRIZATIONS
% x(z), xdot(z), u(z)
% case 1: stationary state with loop voltage unknown

function [x, dx_dz] = x_z_statUpl(z, model)
psisize = numel(model.psi.xind);
zsize = numel(z);
xsize = model.dims.nx;
kinetic_xind = [model.te.xind, model.ti.xind, model.ne.xind, model.ni.xind];

x = zeros(xsize, 1);
x(model.psi.xind(1:end-1)) = z(1:psisize-1);
x(kinetic_xind) = z(psisize:zsize-1);

dx_dz = zeros(xsize, zsize);
dx_dz(model.psi.xind(1:end-1), 1:psisize-1) = eye(psisize-1);
dx_dz(kinetic_xind, psisize:zsize-1) = eye(zsize-psisize);
end

function [xdot, dxdot_dz] = xdot_z_statUpl(z, model, inputs)
if isfield(inputs, 'Upl')
  scaling = inputs.Upl;
else
  scaling = 1;
end
zsize = numel(z);
xsize = model.dims.nx;

xdot = zeros(xsize, 1);
xdot(model.psi.xind) = scaling*z(zsize)*ones(numel(model.psi.xind), 1);

dxdot_dz = zeros(xsize, zsize);
dxdot_dz(model.psi.xind, zsize) = scaling;
end

function [u, du_dz] = u_z_statUpl(z, model, u)
du_dz = zeros(model.dims.nu, numel(z));
end

%% case 2: stationary state with plasma current unknown
% include scaling, elements p should have same order of magnitude

function [x, dx_dz] = x_z_statIp(z, model)
psisize = numel(model.psi.xind);
zsize = numel(z);
xsize = model.dims.nx;
kinetic_xind = [model.te.xind, model.ti.xind, model.ne.xind, model.ni.xind];

x = zeros(xsize, 1);
x(model.psi.xind(1:end-1)) = z(1:psisize-1);
x(kinetic_xind) = z(psisize:zsize-1);

dx_dz = zeros(xsize, zsize);
dx_dz(model.psi.xind(1:end-1), 1:psisize-1) = eye(psisize-1);
dx_dz(kinetic_xind, psisize:zsize-1) = eye(zsize-psisize);
end

function [xdot, dxdot_dz] = xdot_z_statIp(z, model, inputs)
Upl = inputs.Upl;
zsize = numel(z);
xsize = model.dims.nx;

xdot =  zeros(xsize, 1);
xdot(model.psi.xind) = Upl*ones(numel(model.psi.xind), 1);

dxdot_dz = zeros(xsize, zsize);
end

function [u, du_dz] = u_z_statIp(z, model, init, inputs, u)
if isfield(inputs, 'Ip')
  scaling = inputs.Ip;
else
  scaling = init.Ip0;
end
zsize = numel(z);
u(1) = scaling*z(zsize);

du_dz = zeros(model.dims.nu, zsize);
du_dz(1, zsize) = scaling;
end

%% case 3: stationary state with Pec unknown

function [x, dx_dz] = x_z_statPec(z, model)
psisize = numel(model.psi.xind);
zsize = numel(z);
xsize = model.dims.nx;
kinetic_xind = [model.te.xind, model.ti.xind, model.ne.xind, model.ni.xind];

x = zeros(xsize, 1);
x(model.psi.xind(1:end-1)) = z(1:psisize-1);
x(kinetic_xind) = z(psisize:zsize-1);

dx_dz = zeros(xsize, zsize);
dx_dz(model.psi.xind(1:end-1), 1:psisize-1) = eye(psisize-1);
dx_dz(kinetic_xind, psisize:zsize-1) = eye(zsize-psisize);
end

function [xdot, dxdot_dz] = xdot_z_statPec(z, model, inputs)
Upl = inputs.Upl;
zsize = numel(z);
xsize = model.dims.nx;

xdot =  zeros(xsize, 1);
xdot(model.psi.xind) = Upl*ones(numel(model.psi.xind), 1);

dxdot_dz = zeros(xsize, zsize);
end

function [u, du_dz] = u_z_statPec(z, model, inputs, u, uindex)
if isfield(inputs, 'Act')
  scaling = inputs.Act;
else
  scaling = 1e6;
end
zsize = numel(z);
u(uindex) = scaling*z(zsize);

du_dz = zeros(model.dims.nu, zsize);
du_dz(uindex, zsize) = scaling;
end
