clear;
clc;
close all;

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


%% simulate simple AR(1)

MCreps = 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]; 

bin_edges = [-Inf -20:0.5:20 Inf];%[-Inf -2 -1 0 1 2 Inf];
B=length(bin_edges)-1;
bins_center = 0.5 * (bin_edges(1:end-1) + bin_edges(2:end));


d = 5;
B = 6;            % Number of bins
bin_edges = [-Inf -20:0.5:20 Inf];%[-Inf -2 -1 0 1 2 Inf];
B=length(bin_edges)-1;
bins_center = 0.5 * (bin_edges(1:end-1) + bin_edges(2:end));


hit_rates = zeros(MCreps,2);
pearsons  = zeros(MCreps,2);
infids    = zeros(MCreps,2);
errs      = zeros(MCreps,2);

tic

Ts             = [50,100,200,500];
lags = 1;
std = 2;


% 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       = 2; 

for tt = 1:length(Ts)
    T = Ts(tt);
    parfor mcreps = 1:MCreps
   
        %%
        hit_rate     = zeros(1,2);
        pearson_corr = zeros(1,2);
        infid        = zeros(1,2);
        err          = zeros(1,2);


        y = zeros(T+1,1);
        
        b  = [.6 .5 -.2]';
        xx = randn(T+1,3);
        for t = 2:T+1
            y(t) = 0.5 + 0.8*y(t-1) + xx(t-1,:)*b +std*randn; 
        end
        
        y = y(lags+1:end);
        xx = xx(lags+1:end,:);
        
        
        %% estimate simple AR(1)
        
        X = [ones(T-1,1) y(1:end-1) xx(1:end-1,:)];
        beta = (X'*X)\(X'*y(2:end));
        
        yhat = X*beta;
        sig2 = (1/(T-1+1-length(beta)))*sum((y(2:end)-yhat).^2);
        
        % Model definitions
        forecast_mu = @(x) x * beta;
        forecast_sigma = sqrt(sig2);
        bin_probs = @(mu, sigma) normcdf((bin_edges(2:end) - mu) ./ sigma) - ...
                                 normcdf((bin_edges(1:end-1) - mu) ./ sigma);
        
        % Background and target
        x_target = X(end, :);
        X_bg = X(1:end-1, :);
        n_bg = size(X_bg, 1);
        y    = y(2:end);
    
    
        %% shaps for DR
        [out] = VB_BDR_pipe(y,X,X,H_max,init_method,prior_mixing,prior_kernel,10000,tol,prior);
        
        f = @f_DR;
        [shaps] = shap_dens_comp4(x_target,X_bg,f,out,1,bin_edges,"random",1,0,[],[],500);

        % Error
        err(1) = shaps.err;
        
        % Top-k precision
        [~, true_rank] = sort([0; 0.9; abs(b)], 'descend');
        [~, est_rank] = sort(sum(abs(shaps.phi),2), 'descend');
        k = 5; 
        hit_rate(1) = numel(intersect(true_rank(1:k), est_rank(1:k))) / k;
        
        % Pearson correlation
        pearson_corr(1) = corr([0; 0.8; abs(b)], sum(abs(shaps.phi),2), 'Type', 'Pearson');

        delta_std = 1;
        samp      = T-1;
        deltas    = delta_std * randn(samp, d);
        deltas(:,1) = 0;
        diffs     = zeros(samp, 6);
        x_perturbed = x_target - deltas;
        
        % Infidelity score
        [infid(1)] = sqrt(eval_decomp2(f,x_perturbed,shaps.p_full,deltas*shaps.phi,out,bin_edges));
        
        %% shaps for Mitchell QR
        [yforeq,beta_mat] = VBTVPQR(y,X,X,max_iter,0.05:0.05:0.95,0);
        
        f = @f_QR2;
        [shaps] = shap_dens_comp4(x_target,X_bg,f,squeeze(beta_mat(end,:,:)),1,bin_edges,"random",1,0,[],[],500);

        % Error
        err(2) = shaps.err;

        % Top-k precision
        [~, true_rank] = sort([0; 0.9; abs(b)], 'descend');
        [~, est_rank] = sort(sum(abs(shaps.phi),2), 'descend');
        k = 5; 
        hit_rate(2) = numel(intersect(true_rank(1:k), est_rank(1:k))) / k;
            
        %Pearson correlation
        pearson_corr(2) = corr([0; 0.8; abs(b)], sum(abs(shaps.phi),2), 'Type', 'Pearson');

        % Infidelity score 
        [infid(2)] = sqrt(eval_decomp2(f,x_perturbed,shaps.p_full,deltas*shaps.phi,squeeze(beta_mat(end,:,:)),bin_edges));
       
        infids(mcreps,:)         = infid;
        hit_rates(mcreps,:)      = hit_rate;
        pearsons(mcreps,:)       = pearson_corr;
        errs(mcreps,:)           = err;
    end
    results(tt).hit_rates = hit_rates;
    results(tt).pearsons  = pearsons;
    results(tt).infids    = infids;
    results(tt).errs      = errs;
end
toc

save(["SIM4_shaps_model_dens"]);

%% 

for tt = 1:length(Ts)
    pearson(:,tt) = mean(results(tt).pearsons,1)';
    infid(:,tt)   = mean(results(tt).infids,1)';
    hitrate(:,tt) = mean(results(tt).hit_rates,1)';
    err(:,tt)     = mean(results(tt).errs,1)';
 end

models = {"DR","QR"};

j = 1;
for i = 1:4
    metrics(:,:,j) = [err(:,i)  infid(:,i) pearson(:,i) hitrate(:,i)];
    j = j+1;
end

Metrics = [repmat(cellstr(models)',1,1,4) 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)]

