% Main file to generate intial guess for solution to model A6
% (asymmetric CRRA utility function over cons & 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 eta b
global kmin kmax amin amax
global xgrid
global k1_mat k2_mat k3_mat k4_mat K_mat
global a1_mat a2_mat a3_mat a4_mat
global ap1_mat ap2_mat ap3_mat ap4_mat 
global y1_mat y2_mat y3_mat y4_mat Y_mat

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

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

% 1) Calibration
N=4;
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];

etalow=0.1;
etahigh=1;
eta=-[etalow:(etahigh-etalow)/(N-1):etahigh];


% 2) Compute the steady state, the Negishi weights and the labor weight in
% the utility function

% a)
kss=1;
lss=1;
Kss=N*kss;
Lss=N*lss;
css=A*kss^alpha;
Css=N*css;
lnass=0;
ass=1;
thetass=1;
Yss=A*Kss^alpha*Lss^(1-alpha);

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

% c) Labor weights
b=zeros(1,N);
for j=1:N
    b(j)=(1-alpha)*A^(1-1/gamma(j));
end


% 3) Compute the steady state in logs
lnkss=log(kss);
lnKss=log(Kss);
lncss=log(css);
lnCss=log(Css);
lnthetass=log(thetass);
lnYss=log(Yss);
lnLss=log(Lss);
lnlss=log(lss);
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 l1 lp1
syms c2 lna2 k2 theta2 cp2 lnap2 kp2 thetap2 l2 lp2
syms c3 lna3 k3 theta3 cp3 lnap3 kp3 thetap3 l3 lp3
syms c4 lna4 k4 theta4 cp4 lnap4 kp4 thetap4 l4 lp4
syms C Cp K Kp Y Yp CAC CACp mu mup L Lp

% 1) Define symbolic variables
syms lnc1 lna1 lnk1 lntheta1 lncp1 lnap1 lnkp1 lnthetap1 lnl1 lnlp1
syms lnc2 lna2 lnk2 lntheta2 lncp2 lnap2 lnkp2 lnthetap2 lnl2 lnlp2
syms lnc3 lna3 lnk3 lntheta3 lncp3 lnap3 lnkp3 lnthetap3 lnl3 lnlp3
syms lnc4 lna4 lnk4 lntheta4 lncp4 lnap4 lnkp4 lnthetap4 lnl4 lnlp4
syms lnC lnCp lnK lnKp lnY lnYp CAC CACp lnmu lnmup lnL lnLp

% 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)*exp(lnlp1)^(1-alpha)];
f1b=b(1)*[exp(lnl1)^(alpha-1/eta(1))]-[exp(lnc1)^(-1/gamma(1))]*exp(lna1)*(1-alpha)*A*exp(lnk1)^alpha;
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)*exp(lnlp2)^(1-alpha)];
f5b=b(2)*[exp(lnl2)^(alpha-1/eta(2))]-[exp(lnc2)^(-1/gamma(2))]*exp(lna2)*(1-alpha)*A*exp(lnk2)^alpha;
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;

f9=exp(lnc3)^(-1/gamma(3))*[1+phi*(exp(lnkp3)/exp(lnk3)-1)]-beta*exp(lncp3)^(-1/gamma(3))*[exp(lnthetap3)+exp(lnap3)*A*alpha*exp(lnkp3)^(alpha-1)*exp(lnlp3)^(1-alpha)];
f9b=b(3)*[exp(lnl3)^(alpha-1/eta(3))]-[exp(lnc3)^(-1/gamma(3))]*exp(lna3)*(1-alpha)*A*exp(lnk3)^alpha;
f10=exp(lntheta3)-[1+phi*[1+0.5*(exp(lnkp3)/exp(lnk3)-1)]*(exp(lnkp3)/exp(lnk3)-1)];
f11=tau(3)*exp(lnc3)^(-1/gamma(3))-exp(lnmu);
f12=lnap3-rho*lna3;

f13=exp(lnc4)^(-1/gamma(4))*[1+phi*(exp(lnkp4)/exp(lnk4)-1)]-beta*exp(lncp4)^(-1/gamma(4))*[exp(lnthetap4)+exp(lnap4)*A*alpha*exp(lnkp4)^(alpha-1)*exp(lnlp4)^(1-alpha)];
f13b=b(4)*[exp(lnl4)^(alpha-1/eta(4))]-[exp(lnc4)^(-1/gamma(4))]*exp(lna4)*(1-alpha)*A*exp(lnk4)^alpha;
f14=exp(lntheta4)-[1+phi*[1+0.5*(exp(lnkp4)/exp(lnk4)-1)]*(exp(lnkp4)/exp(lnk4)-1)];
f15=tau(4)*exp(lnc4)^(-1/gamma(4))-exp(lnmu);
f16=lnap4-rho*lna4;

f33=exp(lnc1)+exp(lnc2)+exp(lnc3)+exp(lnc4)-exp(lnC);
f34=exp(lnk1)+exp(lnk2)+exp(lnk3)+exp(lnk4)-exp(lnK);
f35=exp(lna1)*A*exp(lnk1)^alpha*exp(lnl1)^(1-alpha)+exp(lna2)*A*exp(lnk2)^alpha*exp(lnl2)^(1-alpha)+...
    exp(lna3)*A*exp(lnk3)^alpha*exp(lnl3)^(1-alpha)+exp(lna4)*A*exp(lnk4)^alpha*exp(lnl4)^(1-alpha)-exp(lnY);
f36=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+...
    exp(lnk3)*(exp(lnkp3)/exp(lnk3)-1)^2+exp(lnk4)*(exp(lnkp4)/exp(lnk4)-1)^2];

f=[f1;f1b;f2;f3;f4;f5;f5b;f6;f7;f8;f9;f9b;f10;f11;f12;f13;f13b;f14;f15;f16;f33;f34;f35;f36];

% 3) Define vectors of states and controls
xx=[lnk1 lnk2 lnk3 lnk4 lna1 lna2 lna3 lna4];
xxp=[lnkp1 lnkp2 lnkp3 lnkp4 lnap1 lnap2 lnap3 lnap4];

yy=[lnc1 lnc2 lnc3 lnc4 lntheta1 lntheta2 lntheta3 lntheta4 lnmu lnC lnK lnY lnl1 lnl2 lnl3 lnl4];
yyp=[lncp1 lncp2 lncp3 lncp4 lnthetap1 lnthetap2 lnthetap3 lnthetap4 lnmup lnCp lnKp lnYp lnlp1 lnlp2 lnlp3 lnlp4];

% 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_steadystate4;
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);
k3_vec=scalup(xgrid(:,3),kmin,kmax);
k4_vec=scalup(xgrid(:,4),kmin,kmax);

lna1_vec=scalup(xgrid(:,5),log(amin),log(amax));
lna2_vec=scalup(xgrid(:,6),log(amin),log(amax));
lna3_vec=scalup(xgrid(:,7),log(amin),log(amax));
lna4_vec=scalup(xgrid(:,8),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) log(k3_vec)-log(kss) log(k4_vec)-log(kss) ...
    lna1_vec-log(ass) lna2_vec-log(ass) lna3_vec-log(ass) lna4_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,:)';
lnkp3_vec=log(kss)+INIT_REG*hx(3,:)';
lnkp4_vec=log(kss)+INIT_REG*hx(4,:)';
kp1_vec = exp(lnkp1_vec);
kp2_vec = exp(lnkp2_vec);
kp3_vec = exp(lnkp3_vec);
kp4_vec = exp(lnkp4_vec);

lnc1_vec=lncss+INIT_REG*gx(1,:)';
lnc2_vec=lncss+INIT_REG*gx(2,:)';
lnc3_vec=lncss+INIT_REG*gx(3,:)';
lnc4_vec=lncss+INIT_REG*gx(4,:)';
c1_vec = exp(lnc1_vec);
c2_vec = exp(lnc2_vec);
c3_vec = exp(lnc3_vec);
c4_vec = exp(lnc4_vec);

OUT_VEC=[kp1_vec c1_vec kp2_vec c2_vec kp3_vec c3_vec kp4_vec c4_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('initguess4.csv',OUT_VEC1,'delimiter',',','precision',12)
