%% This function determines the approximated welfare under a given policy and realization of shocks

function [W, Fst_AW, Fst_SS, Snd_OG, Snd_WS, Snd_CS,  Snd_VA, welfare_highorder]= welfarefunction (y_tilde, dlogz)

global N mu theta sigma_c sigma varepsilon omega_tilde alpha_tilde beta_tilde gamma varphi,

b_tilde                  = [beta_tilde;0];

% Cost-based IO variables

Omega_tilde           = [omega_tilde, alpha_tilde; zeros(1,N+1)];
Psi_tilde             = (eye(N+1)-Omega_tilde)^(-1);
Domar_tilde           = (b_tilde' * Psi_tilde)';
Domar_tilde           = Domar_tilde(1:N);

% Revenue-based IO variables

mu_trans              = [mu;1];
Omega                 = diag(mu_trans)^(-1) * Omega_tilde;
Psi                   = (eye(N+1)-Omega)^(-1);
Domar                 = (b_tilde' * Psi)';
Lambda                = Domar(N+1);
Domar                 = Domar(1:N);

% Rigidity-adjusted IO variables

theta_trans           = [theta;1];
Psi_hat               = (eye(N+1)-diag(theta_trans) * Omega_tilde)^(-1);
Domar_hat             = (b_tilde' * Psi_hat)';
Lambda_hat            = Domar_hat(N+1);
Domar_hat             = Domar_hat(1:N);


M        = zeros(N,1);


for k = 1:N
    M(k) = sigma_c * weightedcov(Psi_tilde(:,k), Psi(:,N+1)/Lambda, b_tilde);
    for j = 1:N
        M(k) = M(k) + sigma(j) * Domar(j)/mu_trans(j) * weightedcov(Psi_tilde(:,k), Psi(:,N+1)/Lambda, Omega_tilde(j,:));
    end
end

J        = M .* (1-theta.^(-1));
xi       = 1 - Lambda_hat - J' * Psi_hat(1:N,N+1);
H1       = diag(varepsilon) * diag(Domar.* Psi(1:N,N+1)/Lambda - M) * (diag(theta)^(-1)-eye(N));
B        = zeros(N,N);

if sigma_c == 1 && all(sigma == 1)

    D = diag(Domar) * Psi(1:N,1:N) * diag(Psi(1:N,N+1)/Lambda);
    B = D + D' - diag(Domar.*Psi(1:N,N+1)/Lambda) - (Domar.*Psi(1:N,N+1)/Lambda) * (Domar.*Psi(1:N,N+1)./Lambda)';
else
    for k = 1:N
        for l = 1:N
            B(k,l) = sigma_c * weightedcov(Psi(:,k), Psi(:,l), b_tilde);
            for j = 1:N
                B(k,l) = B(k,l) + sigma(j) * Domar(j) * weightedcov(Psi(:,k), Psi(:,l), Omega(j,:));
            end
        end
    end
end

H2 = (diag(theta)^(-1)-eye(N)) * B * (diag(theta)^(-1)-eye(N));

K        = (gamma + varphi)/(1-Lambda_hat + varphi * (1-Lambda_hat-xi)) * Psi_hat(1:N,N+1);
V        = 1/(1-Lambda_hat + varphi * (1-Lambda_hat-xi)) * Psi_hat(1:N,N+1) * (Domar_tilde' - Domar_hat' * diag(theta) + varphi*J'*Psi_hat(1:N,1:N)*diag(theta)) - Psi_hat(1:N,1:N)*diag(theta);
pi       = K * y_tilde + V * dlogz;

Fst_AW = (1-Lambda) * y_tilde; % first-order bias from aggregate wedge
Fst_SS = Lambda * J' * pi; % first-order bias from allocative efficiency
Snd_OG = - 1/2 * (gamma + varphi) * (y_tilde - (1+varphi)/(gamma+varphi)*J'*pi)^2; % output gap volatility
Snd_WS = - 1/2 * pi' * H1 * pi; % within-sector price dispersion
Snd_CS = - 1/2 * pi' * H2 * pi; % cross-sector price dispersion
Snd_VA = - 1/2 * (gamma-1)*(1+varphi)/(gamma+varphi) * pi'*J*J'*pi; % variation in allocative efficiency
y_hat = y_tilde - 1/2 * pi' * diag(varepsilon) * diag(1/(gamma+varphi) * Domar_tilde + (varphi/(gamma+varphi)) * (Domar.* Psi(1:N,N+1)/Lambda - M)) * (diag(theta)^(-1)-eye(N)) * pi - 1/2 * (varphi/(gamma+varphi)) * pi' * H2 * pi;
welfare_highorder = (1-Lambda) * ((y_hat - y_tilde) + 1/2 * pi' * (H1 + H2) * pi + (1+varphi)/2 * (y_hat-J'*pi)^2);
W = Fst_AW + Fst_SS + Snd_OG + Snd_WS + Snd_CS + Snd_VA;

end