function [varargout] = RAPTOR_predictive(x0,g0,v0,U_ff,model,params,varargin)
% function [X,U,df_dxk1,df_dx,df_du,df_dp,dx_dp] = RAPTOR_predictive(x0,g0,v0,U_ff,model,params)
% function [simres] = RAPTOR_predictive(x0,g0,v0,U_ff,model,params) % returns structure

[verbosity] = parse_inputs(varargin{:});
params = set_verbosity(params,verbosity);

input_checks(model,params,U_ff)

nt = numel(params.tgrid);
% call mode
if (nargout==1 || nargout>2) && (params.numerics.calc_sens) % flag to decide whether forward sensitivities are to be calculated
    calcfwsens = 1;
else
    calcfwsens = 0;
end

% check nans in inputs
assert(~any(any(isnan(x0))),'NaNs in x0');
%assert(~any(any(isnan(g0))),'NaNs in g0');
assert(~any(any(isnan(v0))),'NaNs in v0');
assert(~any(any(isnan(U_ff))),'NaNs in U_ff');
assert(all(U_ff(1,:))>=0,'current must be positive');
assert(~isfield(model,'tnz'),'tnz field no longer in use for RAPTOR, please remove');

% check what is fixed and what is not

        if size(g0,2)==1
            G = repmat(g0,[1,nt]);
        else
            G = g0; % assign g from outside (e.g. time-varying externally specified profiles)
        end
        assert(size(G,2)==nt,'g0 has invalid size')

if (size(v0,2)>=nt)
    V = v0; % assign v from outside (e.g. time-varying density)
elseif size(v0,2) == 1;
    V = repmat(v0,[1,nt]);
else
    error('v0 has invalid size')
end

% number of inputs
if ~isstruct(model.controllers)
    % standard case, no controllers
    nu = size(U_ff,1);
else
    % with controllers
    nu=0; % init
    for ictr = 1:numel(model.controllers.mapping)
        % count actuators as defined in mapping
        nu = nu+numel(model.controllers.mapping{ictr}{1});
    end
end

assert(size(U_ff,2) == nt,'size(2) of U_ff must match numel(params.tgrid)')

% allocate
U    = NaN*ones(nu,nt);
Xdot = NaN*ones(size(x0,1),nt);
X    = NaN*ones(size(x0,1),nt);
niter = NaN*ones(nt,1);

traps = cell(1,nt);
staps = cell(1,nt);
geops = cell(1,nt);

if calcfwsens
    df_dxk  = cell(1,nt);
    df_dxk1 = cell(1,nt);
    df_dx   = cell(1,nt);
    df_dxdot= cell(1,nt);
    df_du   = cell(1,nt);
    df_dd   = cell(1,nt);
    df_du_du_dp = cell(1,nt);
    df_dp   = cell(1,nt);
    dxk_dp  = cell(1,nt+1);
end

newtres = NaN*zeros(params.numerics.nmax+1,nt); % preallocate, assign NaNs here
exitflags = NaN*zeros(1,nt); % preallocate, assign NaNs here

% initialize
xk = x0(:,1); X(:,1) = x0(:,1);

% flying start, useful to get ramp-ups to converge
xdot = initialize_xdot(model,params);

errorflag = 0; % flag to signal we have to exit from loop prematurely (initialize)

tic;

for it = 1:nt % TIME ITERATIONS
    % default input sensitivity
    if calcfwsens && model.np==0
        % civ - control input vector
        dciv_dp = RAPTOR_cvp(it);        
        if ~isempty(params.cvp.transp_matrix)
            % Get dciv_dp for reduced p which doesn't contain fixed values.
            dciv_dpfull = dciv_dp*params.cvp.transp_matrix;
        else
            dciv_dpfull = dciv_dp;
        end
        nu = model.dims.nu;
        ij=1;
        duk_dp = zeros(nu,size(dciv_dpfull,2));
        for ii=1:nu
            % check which actuators are optimized
            if params.cvp.uind(ii)==1
                duk_dp(ii,:) = dciv_dpfull(ij,:);
                ij = ij+1;
            end
        end
    else
        duk_dp = 0.0;
    end
    
    % geometry and kinetic profiles
    gk = G(:,it);
    vk = V(:,it);
    
    %%% B.C. INPUT & FEEDBACK CONTROL ISSUES
    % sometimes input is not defined directly but only via a control law
    % in this case the dU_dp from the input parametrization is wrong
    % and needs to be fixed here
    
    u = RAPTOR_controllers(xk,gk,vk,U_ff,it,model,params); % calculate controller input based on x=xk
    
    % determine time step
    if it ~= nt
        dt = params.tgrid(it+1)-params.tgrid(it);
    else % last time step, need derivative only
        %dt = dt; % reuse previous dt (the one for last step)
    end
    
    % time derivative of kinetic profiles, enters in some equation
    if it==1
        vdot = zeros(size(vk));
    else
        vdot = (V(:,it)-V(:,it-1))/dt;
    end
    
    % time derivative of geometry profiles, enters in some equations
    if it==1
        gdot = zeros(size(gk));
    else
        gdot = (G(:,it)-G(:,it-1))/dt;
    end
    
    dk = zeros(model.dims.nd,1); % disturbance is zero when called from RAPTOR_predictive.
    
    % store first 
    
    %%% RAPTOR time step
    if params.numerics.usemex
        [xk1,xdotk1,dftilde_dx,dftilde_dxdot,dftilde_dxk1,dftilde_dxk,dftilde_du,dftilde_dd,dftilde_dw,dftilde_dp,stap,geop,trap,residues,inewt,exitflag] = ...
        RAPTOR_predictive_step_mex(it,dt,xk,xdot,gk,gdot,vk,vdot,u,dk,model,params);
    else
        [xk1,xdotk1,dftilde_dx,dftilde_dxdot,dftilde_dxk1,dftilde_dxk,dftilde_du,dftilde_dd,dftilde_dw,dftilde_dp,stap,geop,trap,residues,inewt,exitflag] = ...
        RAPTOR_predictive_step(it,dt,xk,xdot,gk,gdot,vk,vdot,u,dk,model,params);
    end
    
    %% forward sensitivity analysis
    % store jacobians
    df_dxk{it} = dftilde_dxk;
    df_dx{it} = dftilde_dx;
    df_dxdot{it} = dftilde_dxdot;
    df_dxk1{it} = dftilde_dxk1;
    df_du{it} = dftilde_du;
    df_dd{it} = dftilde_dd;
    
    % forward sens equation if necessary
    if calcfwsens
        % Check if we need df_dp or whether no model parameters are to be
        % estimated.
        if model.np==0 % No parameter estimation
            % init
            if it==1
                dxk_dp{it} = sparse(numel(xk),size(duk_dp,2));
            end
            
            % make more efficient later by computing LU decomposition of dftilde_dxk1 only
            % once (see "doc LU" for details)
            
            df_dxk_dxk_dp = df_dxk{it}*dxk_dp{it}; % (df/dpsi_k * dpsi_k/dp)
            df_du_du_dp{it} = df_du{it}*duk_dp;
            
            % only df_du*du_dp
            dxk_dp{it+1} = - dftilde_dxk1 \ (df_dxk_dxk_dp + df_du_du_dp{it});
        else
            if it==1
                dxk_dp{it} = sparse(numel(xk),size(dftilde_dp,2));
            end
            df_dxk_dxk_dp = df_dxk{it}*dxk_dp{it}; % (df/dpsi_k * dpsi_k/dp)
            df_dp{it} = dftilde_dp;
            
            % use df_dp
            dxk_dp{it+1} = - dftilde_dxk1 \ (df_dxk_dxk_dp + df_dp{it} );
        end
    end
    
    % sawtooth module: simply substitute resulting state
    %xk1_preMHD = xk1; 
    xdot1_preMHD = xdotk1; % store old values

    if ~isempty(params.saw) && isfield(params.saw, 'active') && params.saw.active
        xk1 = saw(xk1,gk,vk,u,it,model,params.saw);
        xdotk1 = (xk1-xk)/dt;
        xdot1_preMHD = 0.5.*xdot1_preMHD;
    end
    
    %%%% store in matrices
    U(:,it) = u;
    newtres(:,it) = residues;  % newton iteration residual
    exitflags(it) = exitflag;
    niter(it) = inewt;
    X(:,it) = xk;
    Xdot(:,it) = xdotk1;
    G(:,it) = gk;
    V(:,it) = vk;
    
    staps{it} = stap;
    geops{it} = geop;
    traps{it} = trap;
    
    %%% CHECKS that things have not gone wrong
    % check exitflags
    if exitflag % nonzero
        % special case, exitflag=1 with RT=true is ok
        if model.realtime && (exitflag==1)
            errorflag=0;
        else
        errorflag = RAPTOR_errorhandler(exitflag);
        end
        % handle errors in RAPTOR_predictive
    end
    
    % other external checks
    if it~=nt
        if any(isnan(xk1));
            disp('NaNs in xnew - stopping');
        %elseif (model.mout.Cjtor(1,:)*xk1(iPsi) ) < 0;
        %    warning('j0 < 0'); errorflag = 5;
        elseif any(stap.psip < 0);
            % check for negative or zero iota
            warning('dpsi/drho < 0'); errorflag = 6;
        elseif any(imag(xk1)~=0)
            warning('imag(x) ~=0');errorflag = 7;
        end
    end
    
    %%% PLOTS & DISPLAYS
    if params.debug.iterplot && ((it<nt) || it == nt)
        if ~mod(it-1,params.debug.iterplot) || errorflag || it == nt
            timeslice = RAPTOR_out(xk1,gk,vk,xdotk1,u,it,stap,geop,trap,model,params);
            RAPTOR_plot_step(timeslice)
            drawnow;
        end
    end
    
    % optional summary of macroscopic quantities
    %%
    if params.debug.iterdisp
        if ~mod(it-1,params.debug.iterdisp) || errorflag
            iterdisp(xk1,xk,gk,vk,xdotk1,residues,dt,it,inewt,trap,model,params)
        end
    end
    
    % optional break at given iteration
    if isfield(params.debug,'keyboard_at') && ~isempty(params.debug.keyboard_at)
        if any(it == params.debug.keyboard_at)
            disp(['keyboard at it= ',num2str(it),' as requested']);
            dbstack;
            keyboard;
        end
    end
    
    % check exit flag and exit if necessary
    if errorflag
        % fill remaining result with NaNs
        X(:,it+1:end) = NaN;
        Xdot(:,it+1:end) = NaN;
        U(:,it+1:end) = NaN;
        break
    end
    
    xk = xk1;
    xdot = xdot1_preMHD; 
    % use the pre-MHD xdot1 as initialization for the next step
    % the one modified after MHD can be too large
    % so that the extrapolation in the first iteration of the 
    % next step may be too large...
    
    
    %update the equilibrium from chease 
    if model.equi.updateequi > 0
        gk_new = update_g_from_CHEASE(xk,gk,vk,xdot,u,stap,geop,trap,it,dt,model,params);
        % update G for the rest of the foreseeable future
        G(:,it:end) = repmat(gk_new,1,size(G,2)-it+1); %use this
    end
end % of time loop

% outputs, depending on call varargout
%% If one out argument, return all as structure
if nargout == 1;
    % assign outputs
    
    simres.X = X;
    simres.G = G;
    simres.V = V;
    simres.Xdot = Xdot;
    simres.U = U;
    simres.staps = staps;
    simres.geops = geops;
    simres.traps = traps;

    simres.exitflags = exitflags;
    simres.newtres = newtres;
    simres.niter = niter;
    
    simres.df_dx = df_dx;
    simres.df_dxdot = df_dxdot;
    simres.df_dxk = df_dxk;
    simres.df_dxk1 = df_dxk1;
    simres.df_du = df_du;
    simres.df_dd = df_dd;
    if calcfwsens % assign sensitivities if calculated
        simres.df_dp = df_dp;
        simres.dxk_dp = dxk_dp;
    end
    
    varargout{1} = simres;
elseif nargout == 2;
    varargout = {X,U};
elseif nargout == 7;
    varargout = {X,U,df_dxk1,df_dx,df_du,df_dp,dx_du};
else
    varargout = [];
end

return


function [verbosity] = parse_inputs(varargin)
p = inputParser;
addParameter(p,'verbosity',-1,@(x) isnumeric(x) && isscalar(x) && (x>=0 || x==-1));

parse(p,varargin{:});
verbosity = p.Results.verbosity;

return
