%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%                                                                           
%       cake3.m: Cake eating problem with habit formation
%
%       The program solves the standard cake eating model using the method
%       of undetermined coefficients.
%
%       Carlos Garriga, September 2004
%
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
clear all;close all; clc; tic; format short
disp(' ');disp('  CAKE EATING PROBLEM (Undetermined coefficients)');disp(' ');

% PARAMETERS
beta  = 0.5;                             % Discount rate
gamma = 1.0;                             % Relative weight of habit

%----------------------------------------------------------------
% INITIALIZE VARIABLES AND FUNCTIONS
%----------------------------------------------------------------
% Define grid points for cake
wlow  = 0.1;                             % Lowest point in the grid  
whigh = 2;                               % Highest point in the grid
nump  = 50;                              % Number of grid points
ws    = linspace(wlow,whigh,nump)';      % State space
tol   = 1e-5;                            % Tolerance criterium

% Guess for the initial value function and policy function
v1 = zeros(nump,1);                      % Initialize value function
wprime  = zeros(nump,1);                 % Initialize savings policy function
const   = zeros(nump,1);                 % Initialize present consumption policy function
consy   = zeros(nump,1);                 % Initialize past consumption policy function

% Guess a functional form for the value function
%  v(c-1,ws)= a0 + a1 log(c-1) + a2 log(ws)
a0  = -5.0;
a1  = 1.0;
a2  = 1.0;

%----------------------------------------------------------------
% INITIALIZE VALUE FUNCTION ITERATION
%----------------------------------------------------------------
difa0 = 1;
difa1 = 1;
difa2 = 1;
count = 1;

%while( difa0>tol & difa0>tol & difa2>tol )

% For every point in the grid solve the FOC and compute consumption
%  -U'(w-w')+beta*(V1(w-w',w')+V(w-w',w')=0

% Optimal choices today
    wprime = (a2*beta/(gamma+beta*(a1+a2)))*ws;
    const  = ws-wprime;
    const1 = ((gamma+beta*a1)/(gamma+beta*(a1+a2)))*ws;
    
    dif = max(abs(const -const1));
    if dif>1e-6
        disp('   Policy function has not converged')
    break;end
    
% Optimal decision rules implied from habit (last period consumption)    
    wsy    = ((gamma +beta*(a1+a2))/beta*a2)*ws;
    consy  = ((gamma+beta*a1)/(gamma +beta*(a1+a2)))*wsy;    

% Compute the value function associated to the optimal policy

    v1 = log(consy)+gamma*log(const) + beta* (a0 + a1*log(const)+a2*log(wprime));

% Update the coefficients of the value function using OLS
% beta = (x'x)^-1(x'y)

    v0 = a0 + a1*log(consy) + a2*log(ws);
    
    x = [ones(size(ws)) log(consy) log(ws)];
    y = v1;
break
    anew = (x'*x)^-1*(x'*y);
    %anew  = x\y;
    a0new = anew(1);
    a1new = anew(2);
    a2new = anew(3);

% Check convergence on the coeeficients and update
    difa0 = (abs(a0new-a0));
    difa1 = (abs(a1new-a1));
    difa2 = (abs(a2new-a2));
    count = count +1;
    
    disp('   Iter.     a0new      a1new     Dif.a0    Dif.a1    Dif.a2')
    disp([count a0new a1new difa0 difa1 difa2])%
    a0 = a0new;
    a1 = a1new;
    a2 = a2new;
    %end        

break    
    disp(' Convergence has been achieved ')
    a1true = 1/(1-beta);
        
    if abs(a1true-a1)<tol
       disp(' The coefficient has converged to the true value')
    end
break
%----------------------------------------------------------------
% DISPLAY RESULTS
%----------------------------------------------------------------
disp('Solve the model took');
time = toc;
if time<60
    disp([time]);disp('Seconds')
else
    disp([time/60]);disp('Minutes')
end    

num = 4;
plot(ws,wprime,'r:','LineWidth',num); hold
plot(ws,beta*ws*Rate,'*','LineWidth',num)
plot(ws,cons,'r','LineWidth',num); 
plot(ws,ws,'b:')
legend('Savings','True Policy','Consumption',2)
xlabel('Cake Size','FontWeight','bold','FontSize',12)
ylabel('Policy function w`=g(w)','FontWeight','bold','FontSize',12)
title('Optimal Decision Rules','FontWeight','bold','FontSize',14)