clear all
clc

% parameters
global epsi w beta gam tau rstar gstar Zstar phi_r phi_g phi_Z sig_r sig_g sig_Z Pi

epsi = [-2.02018;-0.958572;0;0.958572;2.02018]*sqrt(2);
w = [0.0199532;0.393619;0.945309;0.393619;0.0199532]/sqrt(pi);

beta = 0.97;
gam = 2;
tau = 0.2;

rstar = 0.0723;
gstar = 0.0427;
Zstar = 0.2640*0.15;

phi_r = 0.0536;
phi_g = -0.1027;
phi_Z = 0.8105;
sig_r = 0.0075;
sig_g = 0.0229;
sig_Z = 0.0350*0.15;

Pi = 0.5;

function f = foc0(ahat,r,g,Z,n)
  global epsi w beta gam tau rstar gstar Zstar phi_r phi_g phi_Z sig_r sig_g sig_Z Pi

  r1 = (1-phi_r)*rstar+phi_r*r;
  g1 = (1-phi_g)*gstar+phi_g*g;
  Z1 = (1-phi_Z)*Zstar+phi_Z*Z;

  h = 0;
  for i1 = 1:5
    for i2 = 1:5
      for i3 = 1:5
        h = h+1;
        F(h) = Pi*beta*(1+r1+sig_r*epsi(i1)).^(1-gam).*(ahat+tau+Z).^(-gam)+...
        (1-Pi)*Pi*beta*(1+r1+sig_r*epsi(i1)).^(1-gam).*(ahat).^(-gam)+...
        (1-Pi)^2*beta*(1+r1+sig_r*epsi(i1)).^(1-gam).*(ahat+((1+n)*(1+g1+sig_g*epsi(i2)))./(1+r1+sig_r*epsi(i1)).*(tau+Z1+sig_Z*epsi(i3))).^(-gam)-...
        (1-tau-ahat)^(-gam);
      end
    end
  end
  w3 = kron(w,w,w);
  f = w3'*F';
end

function f = foc1(ahat,r,g,Z,n)
  global epsi w beta gam tau rstar gstar Zstar phi_r phi_g phi_Z sig_r sig_g sig_Z

  r1 = (1-phi_r)*rstar+phi_r*r;
  g1 = (1-phi_g)*gstar+phi_g*g;
  Z1 = (1-phi_Z)*Zstar+phi_Z*Z;

  h = 0;
  for i1 = 1:5
    for i2 = 1:5
      for i3 = 1:5
        h = h+1;
        F(h) = beta*(1+r1+sig_r*epsi(i1)).^(1-gam).*(ahat+tau+Z).^(-gam)-(1-tau-ahat)^(-gam);
      end
    end
  end
  w3 = kron(w,w,w);
  f = w3'*F';
end

tchange = 2
tt = 8; % even number
rr(1) = rstar-2*sig_r;
gg(1) = gstar;
ZZ(1) = Zstar+2*sig_Z;
AA(1) = 1;
NN(1) = 1;

tic
for iter = 1:1000
  disp(iter)
  for t = 1:tt

    rr(t+1) = (1-phi_r)*rstar+phi_r*rr(t)+sig_r*randn;
    gg(t+1) = (1-phi_g)*gstar+phi_g*gg(t)+sig_g*randn;
    ZZ(t+1) = (1-phi_Z)*Zstar+phi_Z*ZZ(t)+sig_Z*randn;

    AA(t+1) = AA(t)*(1+gg(t+1));
    if t <= tt/2
      n = 0.0217;
    else
      n = 0.011;
    end
    NN(t+1) = (1+n)*NN(t);
    if t <= tchange
      func0 = @(ahat) foc0(ahat,rr(t),gg(t),ZZ(t),n);
      aahat(t) = fzero(func0,1-tau);
    else
      func1 = @(ahat) foc1(ahat,rr(t),gg(t),ZZ(t),n);
      aahat(t) = fzero(func1,1-tau);
    end
  end

  aa = aahat.*AA(1:tt);
  ccy = (1-tau)*AA(1:tt)-aa;
  QP = ZZ.*NN.*AA;
  T0 = (1+n)*(1+gg(2:tt+1)).*(tau+ZZ(2:tt+1)).*AA(1:tt);
  T1 = (1+rr(2:tt+1)).*(tau+ZZ(1:tt)).*AA(1:tt);
  TT = [T0(1:tchange) T1(tchange+1:tt)];
  cco = (1+rr(2:tt+1)).*aa+TT;

  V(iter) = (((ccy).^(1-gam)-1)/(1-gam)+beta*((cco).^(1-gam)-1)/(1-gam))*NN(1:tt)';
end
toc
EV = mean(V)

%figure
%plot([ccy;cco;aa]','LineWidth',2);legend('cy','co','a');grid;