close all;

%----------------------------------------------------------------
% 1. Defining variables
%----------------------------------------------------------------
var Ck1 Ck2 Ch1 Ch2 C1 C2 Nk1 Nk2 Nh1 Nh2 Lambda1 Lambda2 Xk1 Xk2 Xh1 Xh2 bk1 bk2 D1 D2 W1 W2 r1 r2 L1 L2 K1 K2 Y1 Y2 Dp1 Dp2 MC1 MC2 Dk1 Dk2 I1 I2 V1 V2 q1 q2 Pi1_c Pi2_c PC1 PC2 Pi1 Pi2 P2 Pstar1 Pstar2 s1 s2 Gamma11 Gamma12 Gamma21 Gamma22 B1 B2 R Pi Th1 Th2 chi1 chi2 z1 z2 a1 a2 ui G1 G2 g1 g2;
varexo e_chi1 e_chi2 e_z1 e_z2 e_a1 e_a2 e_g1 e_g2 e_i; 

parameters beta upsilon alpha delta theta nu eta xi gamma zeta w_bar tau_w psi tau_c tau_k tau_n tau_p v phi_b rho_z sigma_z rho_chi sigma_chi rho_a sigma_a rho_g sigma_g phi_pi rho_i sigma_i sigma k Ck1_ss Ck2_ss Ch1_ss Ch2_ss C1_ss C2_ss Nk1_ss Nk2_ss Nh1_ss Nh2_ss Lambda1_ss Lambda2_ss Xk1_ss Xk2_ss Xh1_ss Xh2_ss bk1_ss bk2_ss D1_ss D2_ss W1_ss W2_ss r1_ss r2_ss L1_ss L2_ss K1_ss K2_ss Y1_ss Y2_ss Dp1_ss Dp2_ss MC1_ss MC2_ss Dk1_ss Dk2_ss I1_ss I2_ss V1_ss V2_ss q1_ss q2_ss Pi1_c_ss Pi2_c_ss PC1_ss PC2_ss Pi1_ss Pi2_ss P2_ss Pstar1_ss Pstar2_ss s1_ss s2_ss Gamma11_ss Gamma12_ss Gamma21_ss Gamma22_ss G1_ss G2_ss B1_ss B2_ss R_ss Pi_ss Th1_ss Th2_ss g1_ss g2_ss;

%----------------------------------------------------------------
% 2. Calibration
%----------------------------------------------------------------

load parameters

% Load SS values

set_param_value('Ck1_ss',Ck1_ss);
set_param_value('Ck2_ss', Ck2_ss);
set_param_value('Ch1_ss', Ch1_ss);
set_param_value('Ch2_ss', Ch2_ss);
set_param_value('C1_ss', C1_ss);
set_param_value('C2_ss', C2_ss);
set_param_value('Nk1_ss', Nk1_ss);
set_param_value('Nk2_ss', Nk2_ss);
set_param_value('Nh1_ss', Nh1_ss);
set_param_value('Nh2_ss', Nh2_ss);
set_param_value('Lambda1_ss', Lambda1_ss);
set_param_value('Lambda2_ss', Lambda2_ss);
set_param_value('Xk1_ss', Xk1_ss);
set_param_value('Xk2_ss', Xk2_ss);
set_param_value('Xh1_ss', Xh1_ss);
set_param_value('Xh2_ss', Xh2_ss);
set_param_value('bk1_ss', bk1_ss);
set_param_value('bk2_ss', bk2_ss);
set_param_value('D1_ss', D1_ss);
set_param_value('D2_ss', D2_ss);
set_param_value('W1_ss', W1_ss);
set_param_value('W2_ss', W2_ss);
set_param_value('r1_ss', r1_ss);
set_param_value('r2_ss', r2_ss);
set_param_value('L1_ss', L1_ss);
set_param_value('L2_ss', L2_ss);
set_param_value('K1_ss', K1_ss);
set_param_value('K2_ss', K2_ss);
set_param_value('Y1_ss', Y1_ss);
set_param_value('Y2_ss', Y2_ss);
set_param_value('Dp1_ss', Dp1_ss);
set_param_value('Dp2_ss', Dp2_ss);
set_param_value('MC1_ss', MC1_ss);
set_param_value('MC2_ss', MC2_ss);
set_param_value('Dk1_ss', Dk1_ss);
set_param_value('Dk2_ss', Dk2_ss);
set_param_value('I1_ss', I1_ss);
set_param_value('I2_ss', I2_ss);
set_param_value('V1_ss', V1_ss);
set_param_value('V2_ss', V2_ss);
set_param_value('q1_ss', q1_ss);
set_param_value('q2_ss', q2_ss);
set_param_value('Pi1_c_ss', Pi1_c_ss);
set_param_value('Pi2_c_ss', Pi2_c_ss);
set_param_value('PC1_ss', PC1_ss);
set_param_value('PC2_ss', PC2_ss);
set_param_value('Pi1_ss', Pi1_ss);
set_param_value('Pi2_ss', Pi2_ss);
set_param_value('P2_ss', P2_ss);
set_param_value('Pstar1_ss', Pstar1_ss);
set_param_value('Pstar2_ss', Pstar2_ss);
set_param_value('s1_ss', s1_ss);
set_param_value('s2_ss', s2_ss);
set_param_value('Gamma11_ss', Gamma11_ss);
set_param_value('Gamma12_ss', Gamma12_ss);
set_param_value('Gamma21_ss', Gamma21_ss);
set_param_value('Gamma22_ss', Gamma22_ss);
set_param_value('G1_ss', G1_ss);
set_param_value('G2_ss', G2_ss);
set_param_value('B1_ss', B1_ss);
set_param_value('B2_ss', B2_ss);
set_param_value('R_ss', R_ss);
set_param_value('Pi_ss', Pi_ss);
set_param_value('Th1_ss', Th1_ss);
set_param_value('Th2_ss', Th2_ss);
set_param_value('g1_ss', g1_ss);
set_param_value('g2_ss', g2_ss);

% Parameter values

set_param_value('sigma',sigma);     %       Sinze of country 1
set_param_value('w_bar',w_bar);     %       Upper threshold on wage	
set_param_value('beta',beta);       %		Discount factor	
set_param_value('upsilon',upsilon); %		Labour supply elasticity	
set_param_value('alpha',alpha);     %		Capital Share in production						
set_param_value('delta',delta);     %		Capital depreciation rate						
set_param_value('theta',theta);     %		Calvo price stickiness							
set_param_value('nu',nu);           %		Share of non-Ricardian households 				
set_param_value('eta',eta);     	%		Elasticity of substitution between C^1 and C^2	
set_param_value('xi',xi);           %		Elasticity of substitution between differentiated goods 	
set_param_value('gamma',gamma);     %		Home bias 									
set_param_value('zeta',zeta);     	%		Investment adjustment cost 						
set_param_value('tau_w',tau_w);     %		Personal income tax rate							
set_param_value('psi',psi);     	%		Measure of transfer to the poor						
set_param_value('tau_c',tau_c);     %		Consumption tax rate							
set_param_value('tau_k',tau_k);     %		Corporate income tax rate 						
set_param_value('tau_n',tau_n);     %		Payroll tax rate								
set_param_value('tau_p',tau_p);     %		Property tax rate 									
set_param_value('v',v);             %		Deduction of capital costs 						
set_param_value('phi_b',phi_b);     %		Coefficient of government debt in expenditure dynamic	
set_param_value('rho_z',rho_z);     %		Autocorrelation shock to labour disutility				
set_param_value('sigma_z',sigma_z); %		Sd. of labour disutility shock						
set_param_value('rho_chi',rho_chi); %		Autocorrelation consumption demand shock			
set_param_value('sigma_chi',sigma_chi);  %		Sd. of consumption demand shock					
set_param_value('rho_a',rho_a);     %		Autocorrelation productivity shock				
set_param_value('sigma_a',sigma_a); %		Sd. productivity shock							
set_param_value('rho_g',rho_g);     %		Autocorrelation fiscal expenditure 				
set_param_value('sigma_g',sigma_g); %		Sd. fiscal expenditure							
set_param_value('phi_pi',phi_pi);   %		Coefficient of inflation in interest rate rule		
set_param_value('rho_i',rho_i);     %		Autocorrelation monetary policy shock 				
set_param_value('sigma_i',sigma_i); %		Sd. of monetary policy shock
set_param_value('k',k);             %       Constant of the approximation to the indicator function


%----------------------------------------------------------------
% 3. Model
%----------------------------------------------------------------
model;

% HH - Ricardian (12)
%------------------
exp(chi1 - z1)*Ck1*Nk1^upsilon = (1-tau_w)*W1/(1+tau_c);
exp(chi2 - z2)*Ck2*Nk2^upsilon = (1-tau_w)*W2/(1+tau_c);

Ck1(+1)/Ck1 = beta*exp(z1(+1)-z1)*(1+(1-tau_w)*R(+1))/Pi1_c(+1);    %>>>PROBLEM!!!!
Ck2(+1)/Ck2 = beta*exp(z2(+1)-z2)*(1+(1-tau_w)*R(+1))/Pi2_c(+1);    %>>>PROBLEM!!!!

Lambda1 = beta*Ck1/Ck1(+1);                                         %>>>PROBLEM!!!!   
Lambda2 = beta*Ck2/Ck2(+1);                                         %>>>PROBLEM!!!!        

Xk1 = R*(bk1(-1)/PC1)*(1/Pi1) + D1 + W1*Nk1;                        %>>>PROBLEM!!!!
Xk2 = R*(bk2(-1)/PC2)*(1/Pi1) + D2 + W2*Nk2;                        %>>>PROBLEM!!!!

(1+tau_c)*PC1*Ck1 + bk1 = (1+(1-tau_w)*R)*bk1(-1)/Pi1 + PC1*(1-tau_w)*(D1 + W1*Nk1);             %>>>PROBLEM!!!!
(1+tau_c)*PC2*Ck2 + bk2 = (1+(1-tau_w)*R)*bk2(-1)/Pi1 + PC2*(1-tau_w)*(D2 + W2*Nk2);             %>>>PROBLEM!!!!

D1 = (1/(sigma*(1-nu)))*((1-tau_k)*Dp1 + Dk1);
D2 = (1/((1-sigma)*(1-nu)))*((1-tau_k)*Dp2 + Dk2);


% HH - Non Ricardian (8)
%----------------------

Th1 = (psi/(1+2^(-2*k*(w_bar - W1*Nh1))))*(1-Nh1*W1/w_bar);
Th2 = (psi/(1+2^(-2*k*(w_bar - W2*Nh2))))*(1-Nh2*W2/w_bar);

exp(chi1 - z1)*Ch1*Nh1^upsilon = (1 - (psi/w_bar)*(1/(1+2^(-2*k*(w_bar - W1*Nh1)))))*(1-tau_w)*W1/(1+tau_c);
exp(chi2 - z2)*Ch2*Nh2^upsilon = (1 - (psi/w_bar)*(1/(1+2^(-2*k*(w_bar - W2*Nh2)))))*(1-tau_w)*W2/(1+tau_c);

Xh1 = W1*Nh1 + Th1;
Xh2 = W2*Nh2 + Th2;

(1+tau_c)*Ch1 = (1-tau_w)*Xh1;
(1+tau_c)*Ch2 = (1-tau_w)*Xh2;


% Firm - Final Good (6)
%-------------------

Dp1 = Y1 - (1+tau_n)*W1*L1 - r1*K1(-1);
Dp2 = Y2 - (1+tau_n)*W2*L2 - r2*K2(-1);
    
K1(-1)/L1 = (1+tau_n)*(W1/r1)*(alpha/(1-alpha));                                 %>>>PROBLEM!!!!
K2(-1)/L2 = (1+tau_n)*(W2/r2)*(alpha/(1-alpha));                                 %>>>PROBLEM!!!!

MC1 = exp(-a1)*(1/alpha^alpha)*(1/(1-alpha)^(1-alpha))*(r1^alpha)*((1+tau_n)*W1)^(1-alpha);
MC2 = exp(-a2)*(1/alpha^alpha)*(1/(1-alpha)^(1-alpha))*(r2^alpha)*((1+tau_n)*W2)^(1-alpha);


% Firm - Capital Good (10)
%--------------------

Dk1 = r1*K1(-1) - I1 - tau_p*V1;
Dk2 = r2*K2(-1) - I2 - tau_p*V2;

V1 = q1*K1(-1);
V2 = q2*K2(-1);

q1*(1 + zeta*I1) = 1;                        %>>>PROBLEM!!!!
q2*(1 + zeta*I2) = 1;                        %>>>PROBLEM!!!!

I1 = K1 - (1-delta)*K1(-1) + (zeta/2)*I1^2;
I2 = K2 - (1-delta)*K2(-1) + (zeta/2)*I2^2;

q1 = Lambda1*(r1(+1) + (1-delta)*q1(+1));
q2 = Lambda2*(r2(+1) + (1-delta)*q2(+1));


% Prices (15)
%---------
s1 = theta*(1/Pi1^(-xi)) + (1-theta)*(Pstar1^(-xi));                                 %>>>PROBLEM!!!!
s2 = theta*((P2(-1)/P2)^(-xi))*(1/Pi1^(-xi)) + (1-theta)*((Pstar2/P2)^(-xi));        %>>>PROBLEM!!!!

PC1 = (gamma + (1-gamma)*P2^(1-eta))^(1/(1-eta));
PC2 = ((1-gamma) + gamma*P2^(1-eta))^(1/(1-eta));

Gamma11 = MC1*Y1 + theta*Lambda1*(Pi1(+1)^(xi))*Gamma11(+1);
Gamma12 = MC2*Y2 + theta*Lambda2*(Pi2(+1)^(xi))*Gamma12(+1);

Gamma21 = Y1 + theta*Lambda1*(Pi1(+1)^(xi-1))*Gamma21(+1);
Gamma22 = Y2 + theta*Lambda2*(Pi2(+1)^(xi-1))*Gamma22(+1);

Pstar1*Gamma21 = (xi/(xi-1))*Gamma11;
(Pstar2/P2)*Gamma22 = (xi/(xi-1))*Gamma12;                       %>>>PROBLEM!!!!

P2 = (theta*(P2(-1)^(1-xi))*(1/Pi1^(1-xi)) + (1-theta)*(Pstar2^(1-xi)))^(1/(1-xi));

Pi1_c = (PC1/PC1(-1))*Pi1;
Pi2_c = (PC2/PC2(-1))*Pi1;

1 = theta*Pi1^(xi-1) + (1-theta)*(Pstar1^(1-xi));                        %>>>PROBLEM!!!!
1 = theta*Pi2^(xi-1) + (1-theta)*((Pstar2/P2)^(1-xi));                   %>>>PROBLEM!!!!


% Government (2)
%------------

tau_c*C1 + tau_p*V1 + tau_w*sigma*(nu*Xh1 + (1-nu)*Xk1) + tau_n*W1*L1 + tau_k*Dp1 - sigma*nu*Th1 = G1/PC1 + (1+R)*(B1(-1)/PC1)*(1/Pi1) - B1/PC1;                      %>>>PROBLEM!!!!
tau_c*C2 + tau_p*V2 + tau_w*(1-sigma)*(nu*Xh2 + (1-nu)*Xk2) + tau_n*W2*L2 + tau_k*Dp2  - (1-sigma)*nu*Th2 = P2*G2/PC2 + (1+R)*(B2(-1)/PC2)*(1/Pi1) - B2/PC2;          %>>>PROBLEM!!!!

%53

% CB (2)
%---
Pi = (Pi1^sigma)*(Pi2^(1-sigma));

R/R_ss = (Pi^phi_pi)*exp(ui);


% Aggregate demand, feasibility, and market clearing (8)
%---------------------------------------------------
C1 = sigma*(nu*Ch1 + (1-nu)*Ck1);
C2 = (1-sigma)*(nu*Ch2 + (1-nu)*Ck2);

Y1 = gamma*(1/PC1^(-eta))*C1 + (1-gamma)*(1/PC2^(-eta))*C2 + G1/PC1 + I1;                 %>>>PROBLEM!!!!
Y2 = (1-gamma)*((P2/PC1)^(-eta))*C1 + gamma*((P2/PC2)^(-eta))*C2 + G2*P2/PC2 + I2;        %>>>PROBLEM!!!!

L1 = sigma*(nu*Nh1 + (1-nu)*Nk1);
L2 = (1-sigma)*(nu*Nh2 + (1-nu)*Nk2);

(K1(-1)^alpha)*(L1^(1-alpha)) = s1*Y1;
(K2(-1)^alpha)*(L2^(1-alpha)) = s2*Y2;


% Shocks
%-------
g1 = G1/Y1;             %>>>PROBLEM!!!!
g2 = G2/Y2;             %>>>PROBLEM!!!!

log(g1) = (1-rho_g)*log(g1_ss) + rho_g*log(g1(-1)) + e_g1;          %>>>PROBLEM!!!!
log(g2) = (1-rho_g)*log(g2_ss) + rho_g*log(g2(-1)) + e_g2;          %>>>PROBLEM!!!!

a1 = rho_a*a1(-1) + e_a1;
a2 = rho_a*a2(-1) + e_a2;

z1 = rho_z*z1(-1) + e_z1;
z2 = rho_z*z2(-1) + e_z2;

chi1 = rho_chi*chi1(-1) + e_chi1;
chi2 = rho_chi*chi2(-1) + e_chi2;

ui = rho_i*ui(-1) + e_i;

end;

steady(solve_algo=4,maxit=10000);
check;

%----------------------------------------------------------------
% 4. Steady State
%----------------------------------------------------------------

initval;

Ck1 = Ck1_ss;
Ck2 = Ck2_ss;
Ch1 = Ch1_ss;
Ch2 = Ch2_ss;
C1 = C1_ss;
C2 = C2_ss;
Nk1 = Nk1_ss; 
Nk2 = Nk2_ss;
Nh1 = Nh1_ss;
Nh2 = Nh2_ss;
Lambda1 = Lambda1_ss; 
Lambda2 = Lambda2_ss;
Xk1 = Xk1_ss;
Xk2 = Xk2_ss;
Xh1 = Xh1_ss;
Xh2 = Xh2_ss;
bk1 = bk1_ss;
bk2 = bk2_ss;
D1 = D1_ss;
D2 = D2_ss;
W1 = W1_ss;
W2 = W2_ss;
r1 = r1_ss;
r2 = r2_ss;
L1 = L1_ss;
L2 = L2_ss;
K1 = K1_ss;
K2 = K2_ss;     
Y1 = Y1_ss;
Y2 = Y2_ss;
Dp1 = Dp1_ss;
Dp2 = Dp2_ss;
MC1 = MC1_ss;
MC2 = MC2_ss;
Dk1 = Dk1_ss;
Dk2 = Dk2_ss;
I1 = I1_ss;
I2 = I2_ss;
V1 = V1_ss;
V2 = V2_ss;
q1 = q1_ss;
q2 = q2_ss;
Pi1_c = Pi1_c_ss;
Pi2_c = Pi2_c_ss;
PC1 = PC1_ss;
PC2 = PC2_ss;
Pi1 = Pi1_ss;
Pi2 = Pi2_ss;
P2 = P2_ss;
Pstar1 = Pstar1_ss;
Pstar2 = Pstar2_ss;
s1 = s1_ss;
s2 = s2_ss;
Gamma11 = Gamma11_ss;
Gamma12 = Gamma12_ss;
Gamma21 = Gamma21_ss;
Gamma22 = Gamma22_ss;
%G1 = G1_ss;
%G2 = G2_ss;
B1 = B1_ss;
B2 = B2_ss;
R = R_ss;
Pi = Pi_ss;
Th1 = Th1_ss;
Th2 = Th2_ss;
chi1 = 0; 
chi2 = 0;
z1 = 0;
z2 = 0;
a1 = 0;
a2 = 0;
%g1 = g1_ss;
%g2 = g2_ss;
ui = 0;
end;

%----------------------------------------------------------------
% 5. Shocks
%----------------------------------------------------------------

shocks;
var e_chi1 = sigma_chi^2;
var e_chi2 = sigma_chi^2;
var e_z1 = sigma_z^2;
var e_z2 = sigma_z^2;
var e_a1 = sigm_a^2;
var e_a2 = sigm_a^2;
%var e_g1 = sigma_g^2;
%var e_g2 = sigma_g^2;
var e_i = sigma_1^2;
end;



%----------------------------------------------------------------
% 6. Solve
%----------------------------------------------------------------


stoch_simul(hp_filter = 1600, order=1, irf=80);