%This function calculates and returns the first order Sobol indices S and
%ST for each parameter of the ODE. y0 is the vector of initial values for 
%the ODE and M is the number of iterations
%
%Author: Natalie M. Clark, nmclark2@ncsu.edu
%Date: May 19, 2016
function [S,ST] = shr_sensitivity_analysis(y0,M) 

%build the matrices A and B, which consist of random draws of the
%parameters from their distributions. We assume uninformative priors for
%all parameters.
k_1= 0.1+(700-0.1)*rand(2*M,1); %range from Cruz-Ramirez
a_1 = 0.01+(20-0.01)*rand(2*M,1); %arbitrary range
a_2 = 0.01+(20-0.01)*rand(2*M,1); %arbitrary range
d_1 = 0.01+(1000-0.01)*rand(2*M,1); %range from Cruz-Ramirez
d_2 = 0.01+(1000-0.01)*rand(2*M,1); %based on d1
k_2 = 0.1+(700-0.01)*rand(2*M,1); %based on k1
d_3 = 0.01+(1000-0.01)*rand(2*M,1); %based on d1
k_3 = 0.01+(3000-0.01)*rand(2*M,1); %range from Cruz-Ramirez
K_1D = 1.5+(1000-1.5)*rand(2*M,1); %range from Cruz-Ramirez
K_2D = 1.5+(1000-1.5)*rand(2*M,1); %based on K_1D
d_4 = 0.01+(100-0.01)*rand(2*M,1); %range from Cruz-Ramirez
k_4 = 0.05+(1000-0.05)*rand(2*M,1); %range from Cruz-Ramirez
d_5 = 0.01+(200-0.01)*rand(2*M,1); %range from Cruz-Ramirez
k_5 = 0.05+(1000-0.05)*rand(2*M,1); %based on k4
d_6 = 0.01+(200-0.01)*rand(2*M,1); %based on d5

all = [k_1,a_1,a_2,d_1,d_2,k_2,d_3,k_3,K_1D,K_2D,d_4,k_4,d_5,k_5,d_6];
A= all(1:M,:);
B= all(M+1:2*M,:);

%preallocate
ya=zeros(M,6);
yb=zeros(M,6);
yc=zeros(M,6);
S=zeros(size(A,2),6);
ST=zeros(size(A,2),6);

%use to record the computational time
tic;

%calculate Sobol indices for each parameter
for i = 1:size(A,2)
    i %prints the iteration we are on
    %create C, which is equal to A except that the ith column of B is
    %replaced by the ith column of B
    C = A;
    C(:,i)=B(:,i);
    %for each model simulation
    for j = 1:M
        %save the parameters for this simulation and feed them into the ODE
        %model
        aparams = A(j,:);
        bparams = B(j,:);
        cparams = C(j,:);
        myfuna = @(t,x) shr_model_fixed(t,x,aparams);
        myfunb = @(t,x) shr_model_fixed(t,x,bparams);
        myfunc = @(t,x) shr_model_fixed(t,x,cparams);
        %use ODE15s to solve the ODE
        [T,YA]=ode15s(myfuna,0:0.1:7,y0,[]);
        [~,YB]=ode15s(myfunb,0:0.1:7,y0,[]);
        [~,YC]=ode15s(myfunc,0:0.1:7,y0,[]);
        %use trapz to numerically integrate the solution over time. This is
        %used to calcualte the Sobol index. trapz integrates each column
        %separately
        %sometimes, the solver might not make it all the way through the
        %time series (this is rare). in this case, we will just use the
        %values from a random previous iteration
        if size(YA,1) ~= length(T) || size(YB,1) ~= length(T) || size(YC,1) ~= length(T)
            iter = floor(1 + rand*(j-2));
            ya(j,:) = ya(iter,:);
            yb(j,:) = yb(iter,:);
            yc(j,:) = yc(iter,:);
        else
            ya(j,:) = trapz(T,YA);
            yb(j,:) = trapz(T,YB);
            yc(j,:) = trapz(T,YC);
        end
    end
        %calculate sobol indices for each of the responses
        %estimators from Saltelli et al, 2010
        for k = 1:6
            %f0 = mean(ya(:,k))*mean(yb(:,k));
            amc = (ya(:,k)-yc(:,k))'*(ya(:,k)-yc(:,k));
            S(i,k) = (1/M*yb(:,k)'*(yc(:,k)-ya(:,k)))/var(yb(:,k));
            ST(i,k) = amc/(2*M)/var(yb(:,k));
        end
end

%report computational time
toc;


end
        