function [zk_k,zk_km1,xdotk,Sk_k,Sk_km1,Omk,Lk,Fk,Gk,stapk,geopk,trapk,res,icrash,exitflag] = ...
  RAPTOR_EKF_step(zkm1_km1,xdotkm1,gk,gdotk,vk,vdotk,uk,hk,hk_km1,Hk,Skm1_km1,Qk,Rk,dtk,it,RAPTORmodel,RAPTORparams,KFmodel,KFparams)
%#codegen

% zk_km1: state with only knowledge from from previous step (x_{k|k-1})
% gk: geom profiles
% vk: kinetic profiles
% zk: measurements at current step

persistent stapkm1 geopkm1 trapkm1

if isempty(stapkm1)
  stapkm1 = init_stap(RAPTORmodel);
end
if isempty(geopkm1)
  geopkm1=init_geop(RAPTORmodel);
end
if isempty(trapkm1)
  trapkm1=init_trap(RAPTORmodel);
end

coder.extrinsic('KFplot');
coder.extrinsic('KFdisp');
coder.extrinsic('RAPTOR_plot_step');
coder.extrinsic('RAPTOR_out');
coder.extrinsic('warning');

% checks for some sizes
assert(KFmodel.nz<200);
assert(KFmodel.nx<200);
assert(KFmodel.nu<50);
assert(KFmodel.nd<200);

icrash = 0;
ny = numel(hk_km1);
nz = KFmodel.nz;
assert(ny<200);


%% KALMAN FILTER STEP
% measurement yk
% augmented state zk
% kalman gain Lk
%
% Roughly follows notation of Anderson (1979) Ch. 8

%% state update
xkm1_km1 = zkm1_km1(KFmodel.ix); % get RAPTOR model state
dkm1 = zkm1_km1(KFmodel.id); % disturbance

% init
if it==1
  stapkm1 = state_profiles(xkm1_km1,xdotkm1,gk,gdotk,vk,vdotk,RAPTORmodel,stapkm1);
  geopkm1 = geometry_profiles(gk,gdotk,RAPTORmodel,geopkm1);
  trapkm1 = transport_profiles(stapkm1,stapkm1,geopkm1,dkm1,uk,it,RAPTORmodel,RAPTORparams,0,trapkm1);
end

% nonlinear predictive step - one step simulation of physics
[xk_km1,xdotk,~,~,dftilde_dxk,dftilde_dxkm1,~,dftilde_dd,dftilde_dw,~,...
  stapk,geopk,trapk,res,inewt,exitflag] = ...
  RAPTOR_predictive_step(it,dtk,xkm1_km1, xdotkm1,gk,gdotk,vk,vdotk,uk,dkm1,stapkm1,geopkm1,trapkm1,RAPTORmodel,RAPTORparams);

dk_km1 = dkm1; % trivial disturbance prediction step

[Fk,Gk] = calc_linearization(dftilde_dxk,dftilde_dxkm1,...
  dftilde_dd,dftilde_dw,KFmodel);

% Covariance matrix time step update
% with extra alpha parameter for lower weighting of past measurements
% (cf. Anderson(1979), Ch6.2)
Sk_km1 = KFparams.alpha.^2*Fk(1:nz,1:nz)*Skm1_km1(1:nz,1:nz)*Fk(1:nz,1:nz)' ...
  + Gk*Qk*Gk';

if state_ok(exitflag) && strcmp(RAPTORmodel.saw,'porcelli');
    % possible MHD
    % only if porcelli ST model selected, else no MHD
    [xk_mhd,icrash] = RAPTOR_mhd(xk_km1,gk,vk,uk,it,RAPTORmodel,RAPTORparams);
    mhd_exitflag = check_state_validity(xk_mhd,gk,vk,RAPTORmodel);
    % Recover if things go wrong
    if mhd_exitflag ~= -1
      exitflag = mhd_exitflag; % update exitflag
    else
      % update xk+1 with post-MHD state
      xk_km1 = xk_mhd;
    end
end

if state_ok(exitflag)
  zk_km1 = [xk_km1;dk_km1];
else
  % Backup option in case things go wrong.
  % Restart everything from previous step, deflating disturbance
  
  zk_km1 = zkm1_km1;
  zk_k = zkm1_km1; % previous state
  zk_k(KFmodel.id) = 0.5 * zk_k(KFmodel.id); % decrease disturbance
  
  Sk_km1 = Skm1_km1;  Sk_k = Skm1_km1;
  
  [Omk,Lk,Fk,Gk,xdotk,res] = default_fallback(KFmodel,ny);
  
  stapk = stapkm1;
  geopk = geopkm1;
  trapk = trapkm1;
  
  return
end


%% Measurement updates

% Kalman gain
Omk = Hk*Skm1_km1*Hk' + Rk; % innovation covariance
Lk = (Skm1_km1*Hk')/Omk; % right inverse

% state measurement update
upd = Lk*(hk-hk_km1);
zk_k = zk_km1 + KFparams.Lgain*upd; % extra scalar to control update (usually =1)

% Check whether this state is unphysical
if check_state_validity(zk_k(KFmodel.ix),gk,vk,RAPTORmodel) ~= -1
  % ignore measurement update
  zk_k = zk_km1;
end

% Covariance matrix measurement update
Sk_k = (eye(KFmodel.nz) - Lk*Hk)*Sk_km1;

%% Next step
stapkm1 = stapk;
geopkm1 = geopk;
trapkm1 = trapk;

%% Plots etc
if KFparams.iterdisp && (rem(it,double(KFparams.iterdisp))==0)
  KFdisp(it,inewt,res,exitflag,norm(upd)/norm(zkm1_km1),KFparams);
end

end

function KFdisp(it,inewt,res,exitflag,updn,KFparams)

if ~mod(it-1,KFparams.iterdisp*20)
  fprintf('\n%5s%5s%9s%9s%10s\n','it','inewt','res','|upd|/|zk|','exitflag')
end
fprintf('%5d%5d%9.2e%9.2e%10d\n',[it,inewt,res,updn,exitflag])

end

function [Omk,Lk,Fk,Gk,xdotk,res] = default_fallback(KFmodel,ny)
Omk    = zeros(ny);
Lk     = zeros(KFmodel.nz,ny);
Fk = eye(KFmodel.nz);
Gk = zeros(KFmodel.nz,KFmodel.nw);
xdotk = zeros(KFmodel.nx,1); % set to 0
res = 0;
end

function [Fk,Gk] = calc_linearization(dftilde_dxk,dftilde_dxkm1,...
  dftilde_dd,dftilde_dw,KFmodel)
%% Linearized state equation atrices
dfxkinv  =   dftilde_dxk\eye(KFmodel.nx);
AA       = - dfxkinv * dftilde_dxkm1;
Bxd      = - dfxkinv * dftilde_dd;
Bxw      = - dfxkinv * dftilde_dw; % noise model

Odx = zeros(KFmodel.nd,KFmodel.nx); % zeros matrix
Id = eye(KFmodel.nd,KFmodel.nd); % identity

Fk = [AA, Bxd; Odx, Id];
Gk = blkdiag(Bxw,Id); % Gkw

% xk = A*xkm1 + Bxd*dkm1 + Bxu * uk + Bxw*wxk
% dk = 0*xkm1 +  I *dkm1 +   0 * uk +  Id*wdk
%
% zk = [xk;dk];
% zk = Fk*zkm1 + Gku*uk + Gkw*wk + s
end

function statok = state_ok(exitflag)
statok = any(exitflag == [0 1]);
end