clear all;
close all;

% Definovani strukturalnich parametru modelu
% tj. parametry preferenci (domacnosti) a technologii (firmy)
% strukturalni = policy invariant = nezavisle na hopo

% alpha : kapitalovy podil na vystupu
% beta  : diskontni faktor
% delta : mira depreciace
% sigma : parametr rizikove averze, take (inverzni) mezicasova elast.  substituce 
% gamma : uroven technologie

alpha = .35;
beta  = .98;
delta = .025;
sigma = 2;
kappa = .6;

% pocet exogennich stavu
gg = 2;

% hodnoty gamma
gamma = [4.95
         5.05];
     
% pravdepodobnosti pro gamma

Pi = [kappa 1-kappa
      1-kappa kappa];

  
% ocekavana hodnota gamma
gamma_ss = 5;

%  Najdete steady state uroven kapitalu jako funkci strukturalnich  parametru
kstar = ((1/beta - 1 + delta)/(alpha*gamma_ss))^(1/(alpha-1));


%  Definovani poctu diskretnich hodnot k 
gk = 101;
k  = linspace(0.95*kstar,1.05*kstar,gk);


%  vypocitejte matici spotreby c o rozmerech (gk x gk x gg) 
%  pro vsechny mozne hodnoty  k_t, k_t+1 a gamma_t

for h = 1 : gk
   for i = 1 : gk
       for j = 1 : gg
            c(h,i,j) = gamma(j)*k(h)^alpha + (1-delta)*k(h) - k(i);
            if c(h,i,j) < 0
                c(h,i,j) = 0;
            end
             % h je citac pro stavovou promenou  k_t
             % i je citac pro ridici promenou k_t+1
             % j je citac pro exogenni stavovou promenou gamma_t
       end
   end
end

% vypocitejte matici uzitku (gk x gk x gg) 
%  pro vsechny mozne hodnoty  k_t, k_t+1 a gamma_t

for j = 1 : gg
            if sigma == 1
                u(:,:,j) = log(c(:,:,j));
            else
                u(:,:,j) = (c(:,:,j).^(1-sigma) - 1)/(1-sigma);    
            end
            % j je citac pro exogenni promenou gamma_t
end


   
    
% definujte pocatecni hodnotovou funkci v (vektro nul)
v = zeros(gk,gg);
one = ones(gk,1); 

% nastaveni parametru pro cyklus
convcrit = 1E-9;  % konvergencni kriterium (male cislo)
diff = 1;         % arbitrarne zvolena cislo (vetsi jak convcrit)
iter = 0;         % citac iteraci


while diff > convcrit
    %   pro kazdou kombinaci  k_t  a gamma_t najdi hodnotu k_t+1, ktera maximalizuje sumu okamziteho
    %   uzitku a diskonovaneho budouciho uzitku

    for j = 1 : gg
            Tv(:,j) = max( u(:,:,j) + beta*one*(Pi(j,1)*v(:,1)'+Pi(j,2)*v(:,2)'),[],2);
      
    end
    iter = iter + 1;
    diff = abs(Tv - v);
    v = Tv;
end

% Najdi rozhodovaci pravidlo pro  k_t+1 jako funkce stavu k_t a gamm_t
for j = 1 : gg
        % najdi, ktery prvek vektoru  k_t+1 dava optimalni hodnotu 
        [Tv,gridpoint] = max(u(:,:,j) + beta*one*(Pi(j,1)*v(:,1)'+Pi(j,2)*v(:,2)'),[],2);
        % najdi rozhodovaci pravidlo (jako prvky vektoru) - pouzito pozdeji
        % pro simulaci
        kgridrule(:,j) = gridpoint;
        % najdi  rozhodovaci pravidlo (hodnotu pro k_t+1 ktera dava optimum)
        k_dr(:,j) = k(gridpoint);
end


figure
plot(k,k_dr,k,k);
xlabel('k_t')
ylabel('k_{t+1}')
title('rozhodovaci pravidlo pro k')

% vypocitejte optimalni rozhodovaci pravidlo jako funkce stavu k 

for j = 1 : gg   
        c_dr(:,j) = gamma(j)*k.^alpha + (1-delta)*k - k_dr(:,j)';
end

figure
plot(k,c_dr)
xlabel('k_t')
ylabel('c_t')
title('rozhodovaci pravidlo pro c')

%% simulace

% arbitrarni startovaci bod
endostate = round(gk/2); % prvek vektoru
ksim(1) = k(endostate);  % hodnota kaptialu (prvni simulovana hodnota)

t = (1:100)';  % vektor casu
exostate = 1;  % urceni prvniho soku 

% vyber sekvenci 100 nahodnych promenych a pouzij rozhodovaci pravidlo

for i = 1 : 100
    draw = rand;
    
    if exostate == 1
            if draw < Pi(1,1)
                exostate = 1;
            else
                exostate = 2;
            end
        else
            if draw < Pi(2,2)
                exostate = 2;
            else
                exostate = 1;
            end
    end
     
     % urceni rozhodovaciho pravidla (prvky vekotru) podle stavu k a gamma
     kprimegrid = kgridrule(endostate,exostate);
     % prevod na hodnotu kapitalu k_t+1
     kprime = k(kprimegrid);
     
      
    gammasim(i) = gamma(exostate);
    ksim(i+1) = kprime;
    ysim(i) = gamma(exostate)*ksim(i)^alpha;
    isim(i) = ksim(i+1) - (1-delta)*ksim(i);
    csim(i) = ysim(i) - isim(i);
end

ksim(end) = [];

figure
plot(t,ysim/mean(ysim),t,csim/mean(csim),t,isim/mean(isim),t,ksim/mean(ksim))
legend('y','c','i','k')
title('simulace 100 soku, odchylka od prumeru')