function [z_draws, q_draws, n_draws, alpha_q, alpha_n, R0d_draws, R0c_draws] = MHsampler_Epi(m, sigma_z, sigma_n, n_0, z_0, q_0, I, NN, Data, alfa, betta, recov_d, recov_c)
%% MHsampler_Epi(n_0, q_0, I, N, alpha, beta)
%  Function to generate m MH draws for the epi model
%  Input:
%  m is number of MH draws
%  sigma_z: tuning parameter for z draws
%  sigma_n: tuning parameter for n draws
%  n_0: initial draw for nature n (1 scalar)
%  z_0 is initial draws for z (Nx1 vector)
%  q_0 initial draws for q (NxN matrix)
%  I is NxN matrix with number of successful contaminations
%  NN is NxN matrix with number of failed contaminations
%  data is TxN matrix with 0,1,2's to indicate states of sectors
%  alpha is parameter matrix NxN from Beta distribution for prior of q
%  beta is parameter matrix NxN from Beta distribution for prior of q
%  recov_d is mxN matrix with draws from the probability of recovering from a downturn
%  recov_c is mxN matrix with draws from the probability of recovering from a crisis

g_z = @(n_star, n_m, trunc_up, trunc_dw, sigma_n) normpdf(n_star, n_m, sigma_n)/(normcdf(trunc_up, n_m, sigma_n) - normcdf(trunc_dw, n_m, sigma_n));

N = length(z_0); % number of sectors
q_draws = zeros(m, N, N);
q_draws(1, :, :) = q_0;
alpha_q = zeros(m,N,N);
z_draws = zeros(m, N);
z_draws(1, :) = z_0;
n_draws = zeros(m,1);
n_draws(1) = n_0;
alpha_n = zeros(m,1);
alpha_z = zeros(m,N);

R0d_draws = zeros(m,N);
R0c_draws = zeros(m,N);

for i=1:N
    R0d_draws(1,i) = sum(z_draws(1,i).*q_draws(1,i,:)) + (1-recov_d(1,i))*(1/recov_c(1,i))*sum(z_draws(1,i).*q_draws(1,i,:));
    R0c_draws(1,i) = (1/recov_c(1,i))*sum(z_draws(1,i).*q_draws(1,i,:));
end

for draw = 1:(m-1)
    draw        
    n = n_draws(draw);
    qnew = reshape(q_draws(draw,:,:),N,N);
    for i = 1:N
        % first generate proposal draw q_ij
        z_i = z_draws(draw,i);
        for j=1:N
            if i~=j
                % proposal draw:
                q_ij_prev = q_draws(draw,i,j);
                q_ij_prop = betarnd(alfa(i,j)+I(i,j), betta(i,j)+NN(i,j));
                qkjs = qnew(:,j);
                qkjs_prop = qkjs;
                tries = 1;
                while z_i*q_ij_prop>1 && tries<1000
                    q_ij_prop = betarnd(alfa(i,j)+I(i,j), betta(i,j)+NN(i,j));
                    tries = tries +1;
                end
                if tries>=1000
                    q_draws(draw+1,i,j) = q_draws(draw,i,j);
                    z_draws(draw, i) = z_draws(draw-1,i);
                    alpha_q(draw,i,j) = 0;
                else
                    qkjs_prop(i)=q_ij_prop;
                    qkjs_prev = qkjs;
                    qkjs_prev(i) = q_ij_prev;
                    u = rand;
                    f_ster = (q_ij_prop^alfa(i,j))*((1-q_ij_prop)^betta(i,j))*kernel_pijs(Data, j, i, z_i, qkjs_prop, n);
                    f_prev = (q_ij_prev^alfa(i,j))*((1-q_ij_prev)^betta(i,j))*kernel_pijs(Data, j, i, z_i, qkjs_prev, n);
                    g_prev = betapdf(q_ij_prev, alfa(i,j)+I(i,j), betta(i,j)+NN(i,j));
                    g_ster = betapdf(q_ij_prop, alfa(i,j)+I(i,j), betta(i,j)+NN(i,j));
                    alpha_q(draw,i,j) = min((f_ster*g_prev)/(f_prev*g_ster), 1);
                    if u<alpha_q(draw,i,j)
                        q_draws(draw+1,i,j) = q_ij_prop;
                    else
                        q_draws(draw+1,i,j) = q_draws(draw,i,j);
                    end
                    qnew(i,j) = q_draws(draw+1,i,j);
                end
            end
        end
    end
    zi_new = z_draws(draw,:);
    
    for i=1:N
        %next, generate proposal draw z_i
        trunc_up = 1/max(q_draws(draw+1,i,:));
        trunc_dw = 0;
        reject   = 1;
        while reject==1
            z_i_prop = sigma_z*randn+z_i;
            if z_i_prop > trunc_dw && z_i_prop < trunc_up
                reject=0;
            end
        end
        g_prev_cond_prop = g_z(z_i_prop, z_i, trunc_up, trunc_dw, sigma_z);
        g_prop_cond_prev = g_z(z_i, z_i_prop, trunc_up, trunc_dw, sigma_z);
        zi_s_prop = zi_new;
        zi_s_prop(i) = z_i_prop;
        zi_s_prev = zi_new;
        f_z_prop = kernel_zis(Data, i, zi_s_prop, reshape(q_draws(draw+1,:,:),[N,N]), n);
        f_z_prev = kernel_zis(Data, i, zi_s_prev, reshape(q_draws(draw+1,:,:),[N,N]), n);
        alpha_z(draw,i) = min((f_z_prop*g_prev_cond_prop)/(f_z_prev*g_prop_cond_prev), 1);
        u = rand;
        if u<alpha_z(draw,i)
            z_draws(draw+1, i) = z_i_prop;
        else
            z_draws(draw+1, i) = z_i;
        end
        zi_new(i) = z_draws(draw+1,i);
        
    end
    
    % now I have to simulate the nature parameter
    nprev = n_draws(draw);
    trunc_up = 1;
    trunc_dw = 0;
    reject   = 1;
    while reject==1
        n_prop = sigma_n*randn+nprev;
        if n_prop > trunc_dw && n_prop < trunc_up
            reject=0;
        end
    end
    g_prev_cond_prop = g_z(n_prop, nprev, trunc_up, trunc_dw, sigma_n);
    g_prop_cond_prev = g_z(nprev, n_prop, trunc_up, trunc_dw, sigma_n);
    
    log_f_n_prop = log_kernel_n(Data, z_draws(draw+1,:), reshape(q_draws(draw+1,:,:),[N,N]), n_prop)  ;
    log_f_n_prev = log_kernel_n(Data, z_draws(draw+1,:), reshape(q_draws(draw+1,:,:),[N,N]), nprev);
    alpha_n(draw) = min(exp((log_f_n_prop+log(g_prev_cond_prop))-(log_f_n_prev+log(g_prop_cond_prev))), 1);
    u = rand;
    if u<alpha_n(draw)
        n_draws(draw+1) = n_prop;
    else
        n_draws(draw+1) = nprev;
    end
    
    
    for i=1:N
        R0d_draws(draw+1,i) = sum(z_draws(draw+1,i).*q_draws(draw+1,i,:)) + (1-recov_d(draw+1,i))*(1/recov_c(draw+1,i))*sum(z_draws(draw+1,i).*q_draws(draw+1,i,:));
        R0c_draws(draw+1,i) = (1/recov_c(draw+1,i))*sum(z_draws(draw+1,i).*q_draws(draw+1,i,:));
    end
    
end







end


function [ptilde_i] = ptilde_it(Data_tmin1, cursec, pjis, natur)
prodje = 1-natur;
for j=1:length(Data_tmin1)
    if j~=cursec
        if Data_tmin1(j)>0
            prodje = prodje*(1-pjis(j));
        end
    end
end
ptilde_i = 1 - prodje;
end

function [prodje] = kernel_pijs(Data, cursec_s, cursec_inf, zi_s, qkjs, natur)
[T,N]=size(Data);
prodje = 1;
for t=2:T
    if Data(t-1,cursec_s) == 0 && Data(t-1,cursec_inf)>= 1
        ptilde_j = ptilde_it(Data(t-1,:), cursec_s, (zi_s.*qkjs), natur);
        ISD = Data(t,cursec_s)==1;
        ISS = Data(t,cursec_s)==0;
        prodje = prodje*(ptilde_j^(ISD))*((1-ptilde_j)^(ISS));
    end
end
end

function [prodje] = kernel_zis(Data, cursec_inf, zi_s, qijs, natur)
[T,N]=size(Data);
prodje = 1;
for t=2:T
    for cursec_s=1:N
        if cursec_s~=cursec_inf
            if Data(t-1,cursec_s)==0 && Data(t-1,cursec_inf)>0
                qkcur = qijs(:,cursec_s);
                ptilde_j = ptilde_it(Data(t-1,:), cursec_s, (zi_s.*qkcur), natur);
                ISD = Data(t,cursec_s)==1;
                ISS = Data(t,cursec_s)==0;
                prodje = prodje*(ptilde_j^(ISD))*((1-ptilde_j)^(ISS));
            end
        end
    end
end
end


function [prodje] = kernel_n(Data, zi_s, qijs, natur)
[T,N]=size(Data);
prodje = 1;
for t=2:T
    for cursec_i=1:N
        for cursec_s=1:N
            if cursec_s~=cursec_i
                if Data(t-1,cursec_s)==0 && Data(t-1,cursec_i)>0
                    qkcur = qijs(:,cursec_s);
                    ptilde_j = ptilde_it(Data(t-1,:), cursec_s, (zi_s.*qkcur), natur);
                    ISD = Data(t,cursec_s)==1;
                    ISS = Data(t,cursec_s)==0;
                    prodje = prodje*(ptilde_j^(ISD))*((1-ptilde_j)^(ISS));
                end
            end
        end
    end
end
end


function [prodje] = log_kernel_n(Data, zi_s, qijs, natur)
[T,N]=size(Data);
prodje = 1;
for t=2:T
    for cursec_i=1:N
        for cursec_s=1:N
            if cursec_s~=cursec_i
                if Data(t-1,cursec_s)==0 && Data(t-1,cursec_i)>0
                    qkcur = qijs(:,cursec_s);
                    ptilde_j = ptilde_it(Data(t-1,:), cursec_s, (zi_s.*qkcur), natur);
                    ISD = Data(t,cursec_s)==1;
                    ISS = Data(t,cursec_s)==0;
                    prodje = prodje + ISD*log(ptilde_j) + ISS*log(1-ptilde_j);
                end
            end
        end
    end
end
end