% compute eqm or Nash eqm or coordinated solution

function P = compute(P)


% intermediary profits
P.D  = create_demand(P.Gamma,P.Gamma_H,P.mass,P.mass_H);
temp = 0; % not important if demand fn only dependent on r-rstar
profits = @(tau) tau * P.D.fn(tau+temp,temp,0,0,0);


%%% additional functions
% cost of consumption function h
% P.h     = @(z) h0 + h1*z;
P.yT_h2m  = P.gamma_T*P.yT/P.mu; % per capita T income of h2m's
P.yNT_h2m = P.gamma_NT*P.yNT/P.mu; % per capita NT income of h2m's
a0 = 1-P.mu; b0 = P.mu; c0 = 1-P.alpha;
b1 = P.mu*P.alpha/(P.gamma_NT*(1-P.mu)*(1-P.alpha))*(P.c_lbar-P.yT_h2m);
c1 = b1*(1-P.alpha)*P.gamma_NT;
h0 = -P.mu*P.alpha*(P.c_lbar-P.yT_h2m)/(1-(1-P.alpha)*P.gamma_NT) - P.yT + P.c_lbar;
h1 = 1-P.mu + P.gamma_NT*(1-P.mu)*(1-P.alpha)/(1-(1-P.alpha)*P.gamma_NT);

if b1 < 0
    warning 'Need c_lbar > yT_h2m to have the right sign of intervention.'
    P.yT_h2m
end


% AE period 0 consumption function
Euler = @(rstar,tau) ( P.beta_AE * (1+rstar) )^(1/P.Sigma);
C0 = @(rstar,tau) 1/(1+Euler(rstar,tau)/(1+rstar)) * (P.Y0 + P.Y1/(1+rstar) + P.mass_EME/P.mass_AE*1/2*profits(tau) - P.mass_EME/P.mass_AE*P.nfa0);
C1 = @(rstar,tau) Euler(rstar,tau) * C0(rstar,tau);

% EME period 0 consumption function
Euler = @(rstar,tau) ( P.beta_EME * (1+rstar) * (1+tau) )^(1/P.sigma);
z0 = @(rstar,tau) 1/(1+Euler(rstar,tau)/(1+rstar)) * 1/h1 * (-h0*(1+P.yT_decline/(1+rstar)) - profits(tau) + P.nfa0);
z1 = @(rstar,tau) Euler(rstar,tau) * z0(rstar,tau);

% Objective functions
if P.Sigma == 1
    V_AE  = @(C0,C1) log(C0) + P.beta_AE * log(C1);
else
    V_AE  = @(C0,C1) (C0^(1-P.Sigma)-1)/(1-P.Sigma) + P.beta_AE * (C1^(1-P.Sigma)-1)/(1-P.Sigma);
end
if P.sigma == 1
    v = @(x) a0*log(max(x,0)) + b0*log(max(x-b1,0)) - c0*log(max(x-c1,0));
else
    v = @(x) (max(x-c1,0))^(c0*(P.sigma-1)) * ( ...
                    a0*max(x,0)^(1-P.sigma) + b0 * (a0*c1/b0/b1/(1-c1/b1))^(1-P.sigma) * (max(x-b1,0))^(1-P.sigma) ...
                    )/(1-P.sigma) - 1/(1-P.sigma);
end
V_EME = @(z0,z1) v(z0) + P.beta_EME * v(z1);

% Market clearing
mkt0 = @(z0,C0,profits) P.mass_EME*h1*z0 + P.mass_EME*h0 + P.mass_AE*C0 + P.mass_EME*1/2*profits - P.mass_AE*P.Y0;
mkt1 = @(z1,C1) P.mass_EME*h1*z1 + P.mass_EME*h0*P.yT_decline + P.mass_AE*C1 - P.mass_AE*P.Y1;

% budget 

% nfa and intermediaries
nfa_EME = @(z0) -h1*z0 - h0 + P.nfa0;
bI = @(tau) P.D.fn(tau+temp,temp,0,0,0);
bIH = @(tau) P.D.Hfn(tau+temp,temp,0,0,0);
reserves = @(z0,tau) nfa_EME(z0) + bI(tau) + bIH(tau);

% Compute rstar and tau
switch P.mode
    case 'exog_tau'
        % compute rstar for a given tau
        tau = P.tau;
        fun = @(rstar) mkt0( z0(rstar,tau), C0(rstar,tau), profits(tau) );
        rstar_guess = 0; fun(rstar_guess);
        P.rstar = fzero(fun,rstar_guess);
        
    case 'zero_res'
        % find tau not to optimize but to set reserves to zero
        tau_guess = 0; mode = 1;
        fun = @(tau) zero_res(tau,mkt0,reserves,P.res_target,z0,C0,profits,mode);
        fun(tau_guess);
        P.tau = fzero(fun,tau_guess);
        
        mode = 0;
        P.rstar = zero_res(P.tau,mkt0,reserves,P.res_target,z0,C0,profits,mode);
        
    case 'Nash'
        % first find optimum given rstar
        rstar_guess = 0; mode = 1; % mode = optimization mode
        fun = @(rstar) Nash(rstar,mkt0,V_EME,z0,z1,C0,profits,mode);
        fun(rstar_guess);
        
        if 1==0
            % diagnostics
            rstar_grid = linspace(-0.05,0.05,10); mkt0_grid = rstar_grid*0;
            for i=1:length(rstar_grid)
                mkt0_grid(i) = fun(rstar_grid(i));
            end
            plot(rstar_grid,mkt0_grid)
        end
        
        P.rstar = fzero(fun,rstar_guess);
        
        mode = 0;
        P.tau = Nash(P.rstar,mkt0,V_EME,z0,z1,C0,profits,mode);
        
    case 'Coord'
        % bla
        tau_guess = 0; mode = 1;
        fun = @(tau) Coord(tau,mkt0,V_EME,z0,z1,C0,profits,mode);
        fun(tau_guess);
        P.tau = fminsearch(fun,tau_guess);
        
        mode = 0;
        P.rstar = Coord(P.tau,mkt0,V_EME,z0,z1,C0,profits,mode);
        
end

% compute objects
P.C0 = C0(P.rstar,P.tau); P.C1 = C1(P.rstar,P.tau);
P.z0 = z0(P.rstar,P.tau); P.z1 = z1(P.rstar,P.tau);
P.V_AE    = V_AE(P.C0,P.C1);
P.V_EME   = V_EME(P.z0,P.z1);
P.profits = profits(P.tau);
P.nfa_EME = nfa_EME(P.z0);
P.bI      = bI(P.tau);
P.bIH     = bIH(P.tau);
P.reserves = reserves(P.z0, P.tau);
P.GDP0     = (P.mass_AE*P.Y0 + P.mass_EME*P.yT)/P.alpha;
P.reserves_GDP = P.reserves/P.GDP0;
P.priv_flows   = P.bI + P.bIH;
P.priv_flows_GDP = P.priv_flows/P.GDP0;

% check both market clearing conditions
% display(P.mode)
P.mkt0 = mkt0(P.z0,P.C0,P.profits);
P.mkt1 = mkt1(P.z1,P.C1);


assert(abs(P.mkt0) < 1e-9);
assert(abs(P.mkt1) < 1e-9);

end



%%%%% ZERO RESERVES %%%%%%
function  out = zero_res(tau,mkt0,reserves,res_target,z0,C0,profits,mode)

rstar_guess = 0;
fun = @(rstar) mkt0( z0(rstar,tau), C0(rstar,tau), profits(tau) );
fun(rstar_guess);
rstar = fzero(fun,rstar_guess);

% compute objective
if mode == 1
    out = res_target - reserves(z0(rstar,tau),tau);
else
    out = rstar;
end

end



%%%%% NASH %%%%%%
function  out = Nash(rstar,mkt0,V_EME,z0,z1,C0,profits,mode)

% maximize objective using tau
fun = @(tau) -V_EME(z0(rstar,tau),z1(rstar,tau));
tau_guess = 0.05; fun(tau_guess);
if 1==0
    % diagnostics
    tau_grid = linspace(0,0.8,100); V_EME_grid = tau_grid*0; z0_grid = tau_grid*0; z1_grid = tau_grid*0;
    for i=1:length(tau_grid)
        z0_grid(i) = z0(rstar,tau_grid(i));
        z1_grid(i) = z1(rstar,tau_grid(i));
        V_EME_grid(i) = -fun(tau_grid(i));
    end
    plot(tau_grid,V_EME_grid)
end
tau = fminsearch(fun, tau_guess);

% compute mkt clearing
mkt0_value = mkt0(z0(rstar,tau),C0(rstar,tau),profits(tau));

if mode == 1
    % optim mode: spit out mkt clearing
    out = mkt0_value;
else
    out = tau;
end

end




%%%%% Coord %%%%%%
function  out = Coord(tau,mkt0,V_EME,z0,z1,C0,profits,mode)

rstar_guess = 0;
fun = @(rstar) mkt0( z0(rstar,tau), C0(rstar,tau), profits(tau) );
fun(rstar_guess);
rstar = fzero(fun,rstar_guess);

% compute objective
if mode == 1
    out = -V_EME(z0(rstar,tau),z1(rstar,tau));
else
    out = rstar;
end

end


