clear;
clc;

global beta r ro se gamma

%declare parameters

beta=0.90;
r=0.92/beta-1;
ro=0.70;
se=0.1;
gamma=2;


% Discretize the distribution of e shocks

k=5;                            % nodes for the distribution of income shocks
[e,w]=qnwnorm(k,0,se^2);        % Gaussian quadrature nodes and weights

% Bounds for state space

ymin=e(1)/(1-ro);
ymax=e(end)/(1-ro);

amin =  0;                        % impose for now natural borrowing limit
amax = 10*exp(ymax);             % guess an upper bound on a, check later that do not exceed it


% Declare function space to approximate a'(a,y)

n=[25,21];

smin=[amin,ymin];
smax=[amax,ymax];


fspace=fundefn('cheb',n,smin,smax);
fspace=fundefn('spli',n,smin,smax);
fspace=fundefn('lin',n,smin,smax);          

 scale=1/2;
 fspace=fundef({'spli',  nodeunif(n(1),(smin(1)-amin+.01).^scale,(smax(1)-amin+.01).^scale).^(1/scale)+amin-.01,0,1},...
               {'spli',  nodeunif(n(2),smin(2),smax(2)),0,1});           
           
grid=funnode(fspace);                          %4d grid of collocation nodes 
s=gridmake(grid);                              %collection of  states
ns=length(s);

c=funfitxy(fspace,s,r/(1+r)*s(:,1)+exp(s(:,2)));                   %guess that keep constant assets


for it=1:101
cnew=c;
solve;      
c=funfitxy(fspace,s,x);

fprintf('%4i %6.2e\n',[it,norm(c-cnew)]);
if norm(c-cnew)<1e-4, break, end
end

sfine=gridmake(nodeunif(n(1)*2,smin(1),smax(1)),nodeunif(n(1)*2,smin(2),smax(2)));
xfine=funeval(c,fspace,sfine);
disp('mean and max errors')
disp(mean(abs(euler(xfine,c,fspace,sfine,e,w).*((1+r)*sfine(:,1)+exp(sfine(:,2))-xfine>amin+.01))))
disp(max(abs(euler(xfine,c,fspace,sfine,e,w).*((1+r)*sfine(:,1)+exp(sfine(:,2))-xfine>amin+.01))))

figure(1)
subplot(2,2,1)
sfine=gridmake(nodeunif(n(1)*4,smin(1),smax(1)),0);
xfine=funeval(c,fspace,sfine);
plot(sfine(:,1),xfine)
xlabel('a')
ylabel('c')
title('c(a)')

subplot(2,2,2)
sfine=gridmake(0,nodeunif(n(2)*4,smin(2),smax(2)));
xfine=funeval(c,fspace,sfine);
plot(sfine(:,2),xfine)
xlabel('y')
ylabel('c')
title('c(y)')

subplot(2,2,3)
sfine=gridmake(nodeunif(n(1)*4,smin(1),smax(1)),0);
xfine=funeval(c,fspace,sfine);
plot(sfine(:,1),(1+r)*sfine(:,1)+exp(sfine(:,2))-xfine)
xlabel('a')
ylabel('aprime')
title('aprime(a)')

subplot(2,2,4)
sfine=gridmake(0,nodeunif(n(2)*4,smin(2),smax(2)));
xfine=funeval(c,fspace,sfine);
plot(sfine(:,2),(1+r)*sfine(:,1)+exp(sfine(:,2))-xfine)
xlabel('y')
ylabel('aprime')
title('aprime(y)')