function [x, xdot, u, out_data, res_vec, varargout] = solve_parametrized(model, params, init, g, v, ...
                                                                         xxdotu_z, verbosity, varargin)
% Solve the state equation for a general parametrization of x, xdot and u
x0_provided = false;
if ~isempty(varargin)
  x0_provided = true;
  x0 = varargin{1};
end
f_and_jac = @(z, stap, geop, trap, itn) ... 
            general_state(z, xxdotu_z, model, params, g, v, stap, geop, trap, itn); 
% function handle to be passed to the solver containing the equation to be
% solved f(x, xdot, u) and the parametrizations (x(z), xdot(z), u(z)) 

if x0_provided
  [z_init, stap, geop, trap] = get_init_z(xxdotu_z, model, params, init, g, v, x0);
else
  [z_init, stap, geop, trap] = get_init_z(xxdotu_z, model, params, init, g, v);
end

% calling the Newton solver
[z, out_data, res_vec, inner_jacs] = Newton_Raphson_solver(z_init, f_and_jac, xxdotu_z, model, params, g, v, stap, geop, trap, verbosity);

[x, dx_dz, xdot, dxdot_dz, u, du_dz]       = xxdotu_z(z);

% varargout contains jacbian information for optimization routines
if nargout == 6
    jacs.dx_dz = dx_dz;
    jacs.dxdot_dz = dxdot_dz;
    jacs.du_dz = du_dz;
    jacs.df_dz = inner_jacs.df_dz;
    % jacs.df_dx = inner_jacs.df_dx;
    % jacs.df_dxdot = inner_jacs.df_dxdot;
    jacs.df_du = inner_jacs.df_du;
    varargout{1} = jacs;
end
end

function[z_init, stap, geop, trap] = get_init_z(xxdotu_z, model, params, init, g, v, varargin)
% calculate initial conditions and setup state profiles, geometric profiles
% and transport profiles
z = zeros(model.dims.nx, 1); 
[~, dx_dz, ~, dxdot_dz, ~, du_dz] = xxdotu_z(z);

% determine x_0
if ~isempty(varargin)
  x = varargin{1};
else
  x = RAPTOR_initial_conditions(model, init, g, v);
end
x(model.psi.xind) = x(model.psi.xind) - x(model.psi.xind(end));

% 2. translate to p
% linear paramtrizations x    = dx_dp    * p
%                        xdot = dxdot_dp * p
%                        u    = du_dp    * p
% jacs can be obtained as an output of the parametrization functions
% all_vars = all_jacs * p
% p        = all_jacs\all_vars
all_vars = [x; zeros(model.dims.nx + model.dims.nu, 1)];
all_jacs = [dx_dz; dxdot_dz; du_dz];
z_init = all_jacs\all_vars;
z_init(end) = 1;
[~, ~, xdot, ~, u, ~]    = xxdotu_z(z_init);

d = zeros(model.dims.nd, 1); 
vdot = zeros(size(v));
gdot = zeros(size(g));
stapi = init_stap(model); geopi = init_geop(model); trapi = init_trap(model);
stap = state_profiles(x, xdot, g, gdot, v, vdot, model, stapi);
geop = geometry_profiles(g, gdot, model, geopi);
trap = transport_profiles(stap, stap, geop, d, u, 1, model, params, 0, trapi);
end

%% Anonymous function for evaluation of the transport equations f(x(p), xdot(p), u(p)) and its jacobian
function [f, df_dz, stap, geop, trap, jacs] = general_state(z, xxdotu_z, model, params, g, v, stap, geop, trap, itn)
% case specific: parametrisations x, xdot and u
time_step = 1e-3;   
d = zeros(model.dims.nd, 1);

[x, dx_dz, xdot, dxdot_dz, u, du_dz] = xxdotu_z(z);

[f, df_dxdot, df_dx, df_du, ~, ~, ~, stap, geop, trap] = ...
state_equation(x, x-xdot*time_step, g, v, u, d, xdot, zeros(size(g)), zeros(size(v)), 1, model, params, itn, stap, geop, trap); 

df_dz = df_dx * dx_dz + df_dxdot * dxdot_dz + df_du * du_dz;

jacs.df_dx = df_dx;
jacs.df_dxdot = df_dxdot;
jacs.df_du = df_du;

end
