clear;
clc;
close all;

addpath('data')
addpath('functions')
addpath('figures')
addpath('results')

%% |---------------- Simulations Settings

T           = [50,100,200,500];
DGPs        = 1:8;
MCs         = 500;

%% |---------------- Estimation settings

lags         = 4;
nhor         = 12;
H_max        = 5;            % Number of mixing components 
init_method  = 'cluster';   % cluster initialization method random|cluster|standard
iter         = 20;           % Print every "iter^th" iteration
nthin        = 100;
prior_mixing = 2;           % prior for mixing coefficients: 0: Normal Prior | 1: Tipping's ML prior | 2: Horseshoe Prior
prior_kernel = 0;           % prior for mixing coefficients: 0: Normal Prior | 1: Tipping's ML prior | 2: Horseshoe Prior
basis_exp    = 2;           % 0: no basis function expansion | 1: basis function expansion | 2: panel
max_iter     = 10000;
tol          = 1e-6;        % 1e-10;
R            = 1000;
quants       = [0.05,0.1,0.2,0.5,0.8,0.9,0.95]; 


% prior settings
prior.a_tau = 1;
prior.b_tau = 0.1;

prior.Ab_mix  = 1;
prior.tau_mix = 0.1;     
prior.nu_mix  = 0.1;  

sampler       = 1; 

for dgp = 1:length(DGPs)
    for tt = 1:length(T)
        kl_divs   = zeros(MCs,2);
        mean_divs = zeros(MCs,2);
        var_divs  = zeros(MCs,2);
        skew_divs = zeros(MCs,2);
        kurt_divs = zeros(MCs,2);
        parfor iter = 1:MCs
            if DGPs(dgp) == 1
                y = randn(T(tt),1);
            elseif DGPs(dgp) == 2
                mu    = 1;
                sig   = 1;
                alpha = -.5;
                nu    = 10;
                y     = rskt(T(tt),mu,sig,alpha,nu)';  
            elseif DGPs(dgp) == 3
                mu    = 1;
                sig   = 1;
                alpha = 1;
                nu    = 5;
                y     = rskt(T(tt),mu,sig,alpha,nu)';  
            elseif DGPs(dgp) == 4
                probs = rand(T(tt),1);
                y = (probs<=1/3).*(0.04*randn(T(tt),1)) + (probs>1/3).*(1+0.04*randn(T(tt),1));
            elseif DGPs(dgp) == 5
                probs = rand(T(tt),1);
                y = (probs<=1/6).*(0.2*randn(T(tt),1)) + (probs>1/6 & probs<=1/6+1/3).*(1+0.2*randn(T(tt),1))+ (probs>1/2).*(2+0.2*randn(T(tt),1));
            elseif DGPs(dgp) == 6
                probs = rand(T(tt),1);
                y = (probs<=3/4).*(-43/100+randn(T(tt),1)) + (probs>3/4).*(107/100+1/3*randn(T(tt),1));
            elseif DGPs(dgp) == 7
                probs = rand(T(tt),1);
                y = (probs<=1/10).*(randn(T(tt),1)) + (probs>1/10).*(1/10*randn(T(tt),1));
            elseif DGPs(dgp) == 8
                probs = rand(T(tt),1);
                y = (probs<=2/3).*(randn(T(tt),1)) + (probs>2/3).*(1/10*randn(T(tt),1));    
            end
    
            x = ones(T(tt),1);
            [out] = VB_BDR_pipe(y,x,x,H_max,init_method,prior_mixing,prior_kernel,10000,tol,prior);
    
            %[draws] = f_sample(1,1,out,20000,0);
            [draws] = f_sample_sel(1,1,out,20000,0,sampler);
            %draws   = draws(draws<10 & draws >-10);
            % Estimate densities using ksdensity
            [p_pdf] = ksdensity(draws, draws);
            [q_pdf] = ksdensity(y, draws);  % evaluate Q at points from P
    
            % Avoid log(0)
            epsilon = 1e-10;
            p_pdf = max(p_pdf, epsilon);
            q_pdf = max(q_pdf, epsilon);
    
            % KL divergence approximation
            meandiv1 = mean(draws)-mean(y);
            vardiv1  = var(draws)-var(y);
            skewdiv1 = skewness(draws)-skewness(y);
            kurtdiv1 = kurtosis(draws)-kurtosis(y);
            kl_div1  = mean(log(p_pdf ./ q_pdf));
            %disp(['Estimated KL(P || Q): ', num2str(kl_div)])
    
            [yforeq,beta_mat] = VBTVPQR(y,x,1,max_iter,0.05:0.05:0.95,0);
            post_mean_quants  = yforeq;   
            samples           = QR1_adjusted(0.05:0.05:0.95,post_mean_quants' );
            
            % Estimate densities using ksdensity
            [p_pdf] = ksdensity(samples, samples);
            [q_pdf] = ksdensity(y, samples);  % evaluate Q at points from P
    
            % Avoid log(0)
            epsilon = 1e-10;
            p_pdf = max(p_pdf, epsilon);
            q_pdf = max(q_pdf, epsilon);
    
            % KL divergence approximation
            meandiv2 = mean(samples)-mean(y);
            vardiv2  = var(samples)-var(y);
            skewdiv2 = skewness(samples)-skewness(y);
            kurtdiv2 = kurtosis(samples)-kurtosis(y);
            kl_div2 = mean(log(p_pdf ./ q_pdf));
            %disp(['Estimated KL(P || Q): ', num2str(kl_div)])
            
            
            mean_divs(iter,:)  = [meandiv1,meandiv2];
            var_divs(iter,:)   = [vardiv1,vardiv2];
            skew_divs(iter,:)  = [skewdiv1,skewdiv2];
            kurt_divs(iter,:)  = [kurtdiv1,kurtdiv2];
            kl_divs(iter,:)    = [kl_div1,kl_div2];
        end
        results.dgp(dgp).T(tt).mean     = mean_divs;
        results.dgp(dgp).T(tt).var      = var_divs;
        results.dgp(dgp).T(tt).skewness = skew_divs;
        results.dgp(dgp).T(tt).kurtosis = kurt_divs;
        results.dgp(dgp).T(tt).KL       = kl_divs;
    end
end

%recursive forecasts
save("SIM1_DGPs_"+ string(H_max) + "_priorM_"+ string(prior_mixing) + "_priorK_"+ string(prior_kernel) + "_sampler_" + string(sampler) + "_.mat");


for dgp = 1:length(DGPs)
    for tt = 1:length(T)
        means(tt,:,dgp) = mean(results.dgp(dgp).T(tt).mean.^2);
        vars(tt,:,dgp)  = mean(results.dgp(dgp).T(tt).var.^2);
        skews(tt,:,dgp) = mean(results.dgp(dgp).T(tt).skewness.^2);
        kurts(tt,:,dgp) = mean(results.dgp(dgp).T(tt).kurtosis.^2);
        KLs(tt,:,dgp)   = mean(results.dgp(dgp).T(tt).KL);
    end
end

models = ["DR", "QR"];
j = 1;
for i = 1:8
    metrics(:,:,j) = [reshape(means(:,:,i)',8,1)  reshape(vars(:,:,i)',8,1)  reshape(skews(:,:,i)',8,1)  reshape(kurts(:,:,i)',8,1)  reshape(KLs(:,:,i)',8,1)];
    j = j+1;
end
Metrics = [repmat(cellstr(models)',4,1,8) cellfun(@(x) sprintf('%.2f', x), num2cell(round(metrics,2)), 'UniformOutput', false)];

% Size info
[nRows, nCols, nLayer] = size(Metrics);

% New array: double the number of columns
B = cell(nRows, 2 * nCols, nLayer);

% Fill original columns into odd positions
B(:, 1:2:end, :) = Metrics;

% Fill symbol into even positions
B(:, 2:2:end,:) = repmat({'&'}, nRows, nCols, nLayer);

Metrics = B;
Metrics(:,end,:) = repmat({'"\\"'},nRows,nLayer);

[Metrics(:,1:end-1,1) Metrics(:,2:end,2)]
[Metrics(:,1:end-1,3) Metrics(:,2:end,4)]
[Metrics(:,1:end-1,5) Metrics(:,2:end,6)]
[Metrics(:,1:end-1,7) Metrics(:,2:end,8)]

%% charts for distributions

T = 10000;
y = randn(T,1);
figure()
histogram(y,'Normalization','pdf',NumBins=100);
export_uniform_plot(gcf, 'figures/OS_Figure1_DGP1_hist', ...
    'Width', 700, 'Height', 500, 'FontSize', 12, 'Format', 'pdf');
export_uniform_plot(gcf, 'figures/OS_Figure1_DGP1_hist', ...
    'Width', 700, 'Height', 500, 'FontSize', 12, 'Format', 'eps');

y     = rskt(T,1,1,-.5,10)'; 
figure()
histogram(y,'Normalization','pdf',NumBins=100);
export_uniform_plot(gcf, 'figures/OS_Figure1_DGP2_hist', ...
'Width', 700, 'Height', 500, 'FontSize', 12, 'Format', 'pdf');
export_uniform_plot(gcf, 'figures/OS_Figure1_DGP2_hist', ...
'Width', 700, 'Height', 500, 'FontSize', 12, 'Format', 'eps');

y     = rskt(T,1,1,1,5)';  
figure()
histogram(y,'Normalization','pdf',NumBins=100);
export_uniform_plot(gcf, 'figures/OS_Figure1_DGP3_hist', ...
'Width', 700, 'Height', 500, 'FontSize', 12, 'Format', 'pdf');
export_uniform_plot(gcf, 'figures/OS_Figure1_DGP3_hist', ...
'Width', 700, 'Height', 500, 'FontSize', 12, 'Format', 'eps');
 
probs = rand(T,1);
y = (probs<=1/3).*(0.04*randn(T,1)) + (probs>1/3).*(1+0.04*randn(T,1));
figure()
histogram(y,'Normalization','pdf',NumBins=100);
export_uniform_plot(gcf, 'figures/OS_Figure1_DGP4_hist', ...
'Width', 700, 'Height', 500, 'FontSize', 12, 'Format', 'pdf');
export_uniform_plot(gcf, 'figures/OS_Figure1_DGP4_hist', ...
'Width', 700, 'Height', 500, 'FontSize', 12, 'Format', 'eps');

probs = rand(T,1);
y = (probs<=1/6).*(0.2*randn(T,1)) + (probs>1/6 & probs<=1/6+1/3).*(1+0.2*randn(T,1))+ (probs>1/2).*(2+0.2*randn(T,1));
figure()
histogram(y,'Normalization','pdf',NumBins=100);
export_uniform_plot(gcf, 'figures/OS_Figure1_DGP5_hist', ...
'Width', 700, 'Height', 500, 'FontSize', 12, 'Format', 'pdf');
export_uniform_plot(gcf, 'figures/OS_Figure1_DGP5_hist', ...
'Width', 700, 'Height', 500, 'FontSize', 12, 'Format', 'eps');

probs = rand(T,1);
y = (probs<=3/4).*(-43/100+randn(T,1)) + (probs>3/4).*(107/100+1/3*randn(T,1)); 
figure()
histogram(y,'Normalization','pdf',NumBins=100);
export_uniform_plot(gcf, 'figures/OS_Figure1_DGP6_hist', ...
'Width', 700, 'Height', 500, 'FontSize', 12, 'Format', 'pdf');      
export_uniform_plot(gcf, 'figures/OS_Figure1_DGP6_hist', ...
'Width', 700, 'Height', 500, 'FontSize', 12, 'Format', 'eps');      
            
probs = rand(T,1);
y = (probs<=1/10).*(randn(T,1)) + (probs>1/10).*(1/10*randn(T,1));
figure()
histogram(y,'Normalization','pdf',NumBins=100);
export_uniform_plot(gcf, 'figures/OS_Figure1_DGP7_hist', ...
'Width', 700, 'Height', 500, 'FontSize', 12, 'Format', 'pdf'); 
export_uniform_plot(gcf, 'figures/OS_Figure1_DGP7_hist', ...
'Width', 700, 'Height', 500, 'FontSize', 12, 'Format', 'eps'); 

probs = rand(T,1);
y = (probs<=2/3).*(randn(T,1)) + (probs>2/3).*(1/10*randn(T,1));   
figure()
histogram(y,'Normalization','pdf',NumBins=100);
export_uniform_plot(gcf, 'figures/OS_Figure1_DGP8_hist', ...
'Width', 700, 'Height', 500, 'FontSize', 12, 'Format', 'pdf'); 
export_uniform_plot(gcf, 'figures/OS_Figure1_DGP8_hist', ...
'Width', 700, 'Height', 500, 'FontSize', 12, 'Format', 'eps'); 



