% Main file to generate initial guess for solution to model A5 
% (asymmetric CRRA utility function labor, symmetric CD production function)
% Approximation is done in levels (k) and logs (a)

clear all

cpu0=cputime;

global N beta alpha delta sigma rho phi A gamma tau
global kmin kmax amin amax
global xgrid
global k1_mat k2_mat K_mat
global a1_mat a2_mat
global ap1_mat ap2_mat
global y1_mat y2_mat Y_mat

%%%%%%%%%%%%%%%%%%%%%%
% 1. Getting started %
%%%%%%%%%%%%%%%%%%%%%%

approx=1;   % set order of approximation for perturbation solution (initilization)


% 1) Calibration
N=2;
beta=0.99;
alpha=0.36;
delta=0.025;
sigma=0.01;
rho=0.95;
phi=0.5;
A=(1-beta)/(alpha*beta);
% gammalow=0.25;
% gammahigh=1;
% gamma=[gammalow:(gammahigh-gammalow)/(N-1):gammahigh];
gamma=ones(1,N);


% 2) Compute the steady state and the Negishi weights
kss=1;
Kss=N*kss;
css=A*kss^alpha;
Css=N*css;
lnass=0;
ass=1;
thetass=1;
Yss=N*A*kss^alpha;


% 3) Compute Negishi weights
tau=zeros(1,N);
for j=1:N
    tau(j)=css^(1/gamma(j));
end
muss=1;


% 4) Compute the steady state in logs
lnkss=log(kss);
lnKss=log(Kss);
lncss=log(css);
lnCss=log(Css);
lnthetass=log(thetass);
lnYss=log(Yss);
lnmuss=log(muss);



%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% 2. Construct first-order approximation using method by Klein (2000)  %
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

% 1) Define symbolic variables
syms c1 lna1 k1 theta1 cp1 lnap1 kp1 thetap1
syms c2 lna2 k2 theta2 cp2 lnap2 kp2 thetap2
syms C Cp K Kp Y Yp CAC CACp mu mup
syms lnc1 lna1 lnk1 lntheta1 lncp1 lnap1 lnkp1 lnthetap1
syms lnc2 lna2 lnk2 lntheta2 lncp2 lnap2 lnkp2 lnthetap2
syms lnC lnCp lnK lnKp lnY lnYp CAC CACp lnmu lnmup


% 2) Construct the system of equilibrium conditions

f1=exp(lnc1)^(-1/gamma(1))*[1+phi*(exp(lnkp1)/exp(lnk1)-1)]-beta*exp(lncp1)^(-1/gamma(1))*[exp(lnthetap1)+exp(lnap1)*A*alpha*exp(lnkp1)^(alpha-1)];
f2=exp(lntheta1)-[1+phi*[1+0.5*(exp(lnkp1)/exp(lnk1)-1)]*(exp(lnkp1)/exp(lnk1)-1)];
f3=tau(1)*exp(lnc1)^(-1/gamma(1))-exp(lnmu);
f4=lnap1-rho*lna1;

f5=exp(lnc2)^(-1/gamma(2))*[1+phi*(exp(lnkp2)/exp(lnk2)-1)]-beta*exp(lncp2)^(-1/gamma(2))*[exp(lnthetap2)+exp(lnap2)*A*alpha*exp(lnkp2)^(alpha-1)];
f6=exp(lntheta2)-[1+phi*[1+0.5*(exp(lnkp2)/exp(lnk2)-1)]*(exp(lnkp2)/exp(lnk2)-1)];
f7=tau(2)*exp(lnc2)^(-1/gamma(2))-exp(lnmu);
f8=lnap2-rho*lna2;

f25=exp(lnc1)+exp(lnc2)-exp(lnC);
f26=exp(lnk1)+exp(lnk2)-exp(lnK);
f27=exp(lna1)*A*exp(lnk1)^alpha+exp(lna2)*A*exp(lnk2)^alpha-exp(lnY);
f28=exp(lnC)+exp(lnKp)-exp(lnK)-exp(lnY)+...
    phi/2*[exp(lnk1)*(exp(lnkp1)/exp(lnk1)-1)^2+exp(lnk2)*(exp(lnkp2)/exp(lnk2)-1)^2];

f=[f1;f2;f3;f4;f5;f6;f7;f8;f25;f26;f27;f28];


% 3) Define vectors of states and controls
xx=[lnk1 lnk2 lna1 lna2];
xxp=[lnkp1 lnkp2 lnap1 lnap2];

yy=[lnc1 lnc2 lntheta1 lntheta2 lnmu lnC lnK lnY];
yyp=[lncp1 lncp2 lnthetap1 lnthetap2 lnmup lnCp lnKp lnYp];


% 4) Compute analytical derivatives and evaluate them at the steady state
nx = length(xx);
ny = length(yy);

[fx,fxp,fy,fyp,fypyp,fypy,fypxp,fypx,fyyp,fyy,fyxp,fyx,fxpyp,fxpy,...
    fxpxp,fxpx,fxyp,fxy,fxxp,fxx]=anal_deriv(f,xx,yy,xxp,yyp);

assign_steadystate2;
num_eval;


% 5) Construct log-linear approximation using the Klein (2000) algorithm
[gx,hx] = gx_hx(nfy,nfx,nfyp,nfxp);


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% 3. Construct initial guess for decision rule coefficients %
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

% 1) Construct bounds for the state space
tr=1;  % scale parameter
kmin=1/1.05;
kmax=1.05;

amin=exp(-tr*2*sigma/sqrt(1-rho^2));
amax=exp(+tr*2*sigma/sqrt(1-rho^2));

 tr=0.8;  % scale parameter
 kmin=0.8;
 kmax=1.2;
 amin=exp(-tr*sigma/(1-rho));
 amax=exp(+tr*sigma/(1-rho));


% 2) Construct Smolyak Grid
nx = 2*N;
x_Smol = zeros(1,5);
for j = 1:5
  x_Smol(j) = -cos((j-1)*pi/4);
end

%Build Grid
xgrid = zeros(1+4*nx+2*nx*(nx-1),nx);
xgrid = x_Smol(3) * xgrid;
for j = 1:nx
  xgrid(1+4*(j-1)+1,j) = x_Smol(1);
  xgrid(1+4*(j-1)+2,j) = x_Smol(2);
  xgrid(1+4*(j-1)+3,j) = x_Smol(4);
  xgrid(1+4*(j-1)+4,j) = x_Smol(5);
end
ind = 1+4*nx+1;
points1 = [x_Smol(1);x_Smol(1);x_Smol(5);x_Smol(5)];
points2 = [x_Smol(1);x_Smol(5);x_Smol(1);x_Smol(5)];
for j = 1:(nx-1)
  for i = j+1: nx
    xgrid(ind:ind+3,j) = points1;
    xgrid(ind:ind+3,i) = points2;
    ind = ind+4;
  end
end

% Construct state vectors from grid
k1_vec=scalup(xgrid(:,1),kmin,kmax);
k2_vec=scalup(xgrid(:,2),kmin,kmax);

lna1_vec=scalup(xgrid(:,3),log(amin),log(amax));
lna2_vec=scalup(xgrid(:,4),log(amin),log(amax));

% Construct regressor matrix to be used for perturbation initialization
INIT_REG=[log(k1_vec)-log(kss) log(k2_vec)-log(kss) ...
    lna1_vec-log(ass) lna2_vec-log(ass)];

% Construct next period's capital stocks and current consumption
lnkp1_vec=log(kss)+INIT_REG*hx(1,:)';
lnkp2_vec=log(kss)+INIT_REG*hx(2,:)';
kp1_vec = exp(lnkp1_vec);
kp2_vec = exp(lnkp2_vec);

lnc1_vec=lncss+INIT_REG*gx(1,:)';
lnc2_vec=lncss+INIT_REG*gx(2,:)';
c1_vec = exp(lnc1_vec);
c2_vec = exp(lnc2_vec);

OUT_VEC=[kp1_vec c1_vec kp2_vec c2_vec];
for j = 1:(1+4*nx+2*nx*(nx-1))
  for i = 1:N
  OUT_VEC1(N*(j-1)+i,:) = OUT_VEC(j,2*i-1:2*i);
  end
end

% Output linearization to .csv file
dlmwrite('initguess2.csv',OUT_VEC1,'delimiter',',','precision',12)
%save initguess2.csv OUT_VEC1 -ascii
