%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% simulate.m executes the AMPL model, simulates random paths over the     %
% life cycle and draws respective figures.                                %
% NOTE: you can directly change the data input file in                    %
%       projectwithrisk.mod                                               %
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

clear;

% call AMPL program
dos 'AMPL projectwithrisk.mod';

% call policy function coefficients, ...
datapolicy;

% number of simulated random paths
numPoints = 100;

% draw random numbers
pathsUnif = random('unif', 0, 1, [numPoints, T]);

% calculate cumulated density function
cdfProb = zeros(ns);
for s = 1:ns
    cdfProb(1, s) = prob(1, s);
    for t = 2:ns
        cdfProb(t, s) = cdfProb(t-1, s) + prob(t, s);
    end
end

% calculate invariant distribution for periode 1
[V, D] = eig(prob);
for s = 1:ns
    if(D(s,s) >= 0.99999)
        j = s;
    end
end
a_inv = V(:, j);
a_inv = a_inv./sum(a_inv);

% calculate cumulated density function
for s = 2:ns
    a_inv(s) = a_inv(s) + a_inv(s-1);
end

% set up first draw for (y, r)
path = zeros(numPoints, T);
for p = 1:numPoints
    for s = 1:ns
        if(pathsUnif(p, 1) <= a_inv(s))
            path(p, 1) = s;
            break;
        end
    end
end

% calculate (y, r) path for any other period according to transition
% matrices
for t = 2:T
    for p = 1:numPoints
        for s = 1:ns
            if(pathsUnif(p, t) <= cdfProb(s, path(p, t-1)))
                path(p, t) = s;
                break;
            end
        end
    end
end

% initialize policy variables
asafe = zeros(numPoints, T);
arisk = zeros(numPoints, T);
cons  = zeros(numPoints, T);
wealth = zeros(numPoints, T);

% calculate investment, consumption and wealth paths
for p = 1:numPoints
    wealth(p, 1) = 0;
    for t = 1:T-1
        asafe(p, t) = chebbas(M, lb, ub, wealth(p,t))*coeffsafe(path(p, t), :, t)';
        arisk(p, t) = chebbas(M, lb, ub, wealth(p,t))*coeffrisk(path(p, t), :, t)';
        cons(p, t) = wealth(p, t) + y(path(p, t), t) - asafe(p, t) - arisk(p, t);
        wealth(p, t+1) = asafe(p, t)*(1+r) + arisk(p, t)*R(path(p, min(t+1)), t);
    end
    cons(p, T) = wealth(p, T) + y(path(p, T), T) - asafe(p, T) - arisk(p, T);
end

% in order to ignore divide by zero messages
warning off all;

% draw figures
subplot(2,2,1);
title('Share of risky asset in total invested amount'); 
hold on;
box on;
grid on;
for j = 1:100
    plot([1:T], arisk(j,:)./(asafe(j,:) + arisk(j,:)));
end
hold off

subplot(2,2,2);
title('Consumption'); 
hold on;
box on;
grid on;
for j = 1:100
    plot([1:T], cons(j,:));
end
hold off

income(T+1) = 0;
income(T) = prob(:,1)'*y(:,T);
for t = T-1:-1:1
    income(t) = prob(:,1)'*y(:,t) + income(t+1)/(1+r);
end

subplot(2,2,3);
title('Share of risky asset in total PDV of income'); 
hold on;
box on;
grid on;
for j = 1:100
    plot([1:T], arisk(j,:)./(asafe(j,:) + arisk(j,:)+ income(2:T+1)/(1+r)'));
end
hold off

subplot(2,2,4);
title('Financial Wealth'); 
hold on;
box on;
grid on;
for j = 1:100
    plot([1:T], wealth(j, :));
end
hold off

warning on all;