function [h,dh_dx,dh_du,Rk,T,iT] = diag_ITER(x,g,v,u,it,model,diaglist)
% RAPTOR diagnostic definition file

% initialize
h = zeros(0,1); dh_dx = zeros(0,model.dims.nx); dh_du = zeros(0,model.dims.nu);
Rk = []; T = []; iT = [];

for ii=1:numel(diaglist)
    
    switch diaglist{ii}
        case 'Psib'
            [idh,iddh_dx,iddh_du,idRk,idT,idiT] = diag_ITER_Psib(x,g,v,u,it,model);
        case 'Ipl'
            [idh,iddh_dx,iddh_du,idRk,idT,idiT] = diag_ITER_Ipl(x,g,v,u,it,model);
        case 'iota'
            [idh,iddh_dx,iddh_du,idRk,idT,idiT] = diag_ITER_iota(x,g,v,u,it,model);
        case 'Te'
            [idh,iddh_dx,iddh_du,idRk,idT,idiT] = diag_ITER_Te(x,g,v,u,it,model);
        otherwise 
            error('invalid diagonstic name')
    end
            h = [h;idh]; dh_dx = [dh_dx;iddh_dx]; dh_du = [dh_du;iddh_du];
            Rk = blkdiag(Rk,idRk); T = blkdiag(T,idT); iT = blkdiag(iT,idiT);
end

return

function [h,dh_dx,dh_du,Rk,T,iT] = diag_ITER_Psib(x,g,v,u,it,model)

scal = 1; % typical scale
serr = 0.01; % typical error std

dh_dx = [model.psi.Lam(end,:),zeros(1,numel(model.te.xind))];

h = dh_dx*x;

dh_du = zeros(numel(h),numel(u));

T = 1/scal * eye(numel(h));
iT = scal * eye(numel(h));

Rk = T*(serr.^2 * eye(numel(h)))*T';

return

function [h,dh_dx,dh_du,Rk,T,iT] = diag_ITER_Ipl(x,g,v,u,it,model)

scal = 1e6; % typical scale
serr = 10e3; % error standard deviation

[~,dIpl_dx] = eval_Ipl(x,g,[],model,false);
dIp_dx = dIpl_dx(end,:);

dh_dx = [1/scal*dIp_dx];
h = dh_dx*x;
dh_du = zeros(numel(h),numel(u));

T = 1/scal * eye(numel(h));
iT = scal * eye(numel(h));

Rk = T*(serr.^2 * eye(numel(h)))*T';

return

function [h,dh_dx,dh_du,Rk,T,iT] = diag_ITER_Te(x,g,v,u,it,model)

scal = 1e3; % typical scale
serr = 100; % error standard deviation

%irhomeas = 1:3:model.rgrid.nrho;
irhomeas = 1:1:model.rgrid.nrho;
dh_dx = [zeros(numel(irhomeas),model.psi.nsp),... 
    1/scal*model.te.Lam(irhomeas,:)];

h = dh_dx*x;

dh_du = zeros(numel(h),numel(u));

T = 1/scal * eye(numel(h));
iT = scal * eye(numel(h));

%herr = linspace(100,10,numel(h));
%Rk = T*diag(herr.^2)*T';
Rk = T*(serr.^2 * eye(numel(h)))*T';

return


function [h,dh_dx,dh_du,Rk,T,iT] = diag_ITER_iota(x,g,v,u,it,model)

scal = 10; % typical scale
serr = 0.001; % error standard deviation

%irhomeas = 1:3:model.rgrid.nrho;
irhomeas = 1:1:ceil(model.rgrid.nrho*0.7);

%irhomeas = 7;
[iota,diota_dx] = eval_iota(x,g,[],model,false);
dh_dx = 1/scal*diota_dx;

h = dh_dx*x;

dh_du = zeros(numel(h),numel(u));

T = 1/scal * eye(numel(h));
iT = scal * eye(numel(h));

Rk = T*(serr.^2 * eye(numel(h)))*T';

return