clear;close all;
addpath('data')
addpath('functions')
addpath('figures')
addpath('results')

% Parameters
Ts = [50,100,200,500];          % Time points
d = 3;            % Number of features
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));
sigma2 = 4;
sigma  = sqrt(sigma2);
coeff = 1;
indep = 1;

MCs = 2;
infid = zeros(6,length(Ts),MCs);
err = zeros(6,length(Ts),MCs);
TR  = zeros(6,length(Ts),MCs);
Rs  = zeros(6,d,length(Ts),MCs);

% define loss function parameters
loss_params.lower = 0;
loss_params.upper = 0;
loss_params.alpha = 2;
loss_params.beta  = 2;
loss_params.w     = 0.5;

fprintf('Setup %.5f, MC replication:     %.5f\n', Ts(1), 1);
for tt = 1:length(Ts)
    T = Ts(tt);

    infids        = zeros(6,MCs);
    errs          = zeros(6,MCs);
    hit_rates     = zeros(6,MCs);
    pearson_corrs = zeros(6,MCs);

    parfor mc = 1:MCs

        infid        = zeros(6,1);
        err          = zeros(6,1);
        hit_rate     = zeros(6,1);
        pearson_corr = zeros(6,1);

        if mod(mc,50) == 0
            fprintf('Setup %.5f, MC replication:     %.5f\n', T, mc);
        end
        y = zeros(T+2,d);
        
        f = @f_loss_TVPSV;
        % Simulate data
        X = randn(T, d);
    
        if indep == 1
            %beta = [0.85 0 0;0 0.8 0; 0 0 -0.6];
   
           beta  = [0.85 0.1 0.4;0.1 0.8 -0.2; 0.3 0.5 -0.6];
           betaI = diag(diag(beta));
           c    = [0.2 0.3 0.1; 
                  -0.4 0.5 0.2;
                   0.3 -.2 0.6];
           Sigma = c'*c;
           gamm =  diag(reshape((eye(d^2)-kron(beta,beta))\Sigma(:),d,d));
        
        
           Sigma    = diag(gamm.*(1-diag(betaI).^2));
           beta = betaI;
        else
            beta = [0.85 0.1 0.4;0.1 0.8 -0.2; 0.3 0.5 -0.6];
            c    = [0.2 0.3 0.1; 
                   -0.4 0.5 0.2;
                    0.3 -.2 0.6];
            Sigma = c'*c;
        end
        
     
        for t = 2:T+2
            y(t,:) = beta* y(t-1,:)' + chol(Sigma)'*randn(d,1);
        end
        y = y(3:end,:);
        X = y(2:end-1,:);
        
        cc = chol(Sigma)';
        sig = cc(1,1);

        % Model output: bin probabilities of forecast density
        bin_probs = @(mu) normcdf((bin_edges(2:end) - mu) / sigma) - ...
                          normcdf((bin_edges(1:end-1) - mu) / sigma);
        
        n_bg = 100;
        X_bg =y(1:end-1,:);
        x_target = y(end,:);
        % X_bg = X(1:end-1,:);

        [phi_shap]    = shap_loss_comp(x_target,X_bg,f,[beta(1,:),sig],loss_params,[],bin_edges,"exact",1,0,1);
        [phi_shapR]   = shap_loss_comp(x_target,X_bg,f,[beta(1,:),sig],loss_params,[],bin_edges,"random",1,0,1);
        [phi_shapK]   = shap_loss_comp(x_target,X_bg,f,[beta(1,:),sig],loss_params,[],bin_edges,"kernel",1,0,1);
        [phi_shapknn] = shap_loss_comp(x_target,X_bg,f,[beta(1,:),sig],loss_params,[],bin_edges,"knn",1,0,1);
        
        [phi_grad]    = f_compute_loss_grad(x_target,X_bg,f,[beta(1,:),sig],loss_params,bin_edges);
        [phi_Igrad]   = f_compute_loss_Igrad(x_target,X_bg,f,[beta(1,:),sig],loss_params,bin_edges);
        
        dp = phi_shap.p_full - phi_shap.phi0;  % total change to be decomposed
        
        delta_std = 0.01;
        samp      = T-1;
        deltas    = delta_std * randn(samp, d);
        diffs     = zeros(samp, 6);
        x_perturbed = x_target - deltas;
        
        [infid(1)] = sqrt(eval_loss_decomp(f,x_perturbed,phi_shap.p_full,deltas*phi_shap.phi,[beta(1,:),sig],loss_params,bin_edges));
        [infid(2)] = sqrt(eval_loss_decomp(f,x_perturbed,phi_shap.p_full,deltas*phi_shapR.phi,[beta(1,:),sig],loss_params,bin_edges));
        [infid(3)] = sqrt(eval_loss_decomp(f,x_perturbed,phi_shap.p_full,deltas*phi_shapK.phi,[beta(1,:),sig],loss_params,bin_edges));
        [infid(4)] = sqrt(eval_loss_decomp(f,x_perturbed,phi_shap.p_full,deltas*phi_shapknn.phi,[beta(1,:),sig],loss_params,bin_edges));
        [infid(5)] = sqrt(eval_loss_decomp(f,x_perturbed,phi_shap.p_full,deltas*phi_grad,[beta(1,:),sig],loss_params,bin_edges));
        [infid(6)] = sqrt(eval_loss_decomp(f,x_perturbed,phi_shap.p_full,deltas*phi_Igrad,[beta(1,:),sig],loss_params,bin_edges));
        
        
        err(1) = norm((sum(phi_shap.phi,1) - dp));
        err(2) = norm((sum(phi_shapR.phi,1) - dp));
        err(3) = norm((sum(phi_shapK.phi,1) - dp));
        err(4) = norm((sum(phi_shapknn.phi,1) - (phi_shapknn.p_full-phi_shapknn.phi0cond)));
        err(5)  = norm((sum(phi_grad,1) - dp));
        err(6)  = norm((sum(phi_Igrad,1) - dp));


        % Top-k precision
        k = 1; 
        bench = abs(beta(1,:).*X_bg(end,:));
        [~, true_rank] = sort(bench', 'descend');

        hr = zeros(3,1); 
        [~, est_rank]       = sort(abs(phi_shap.phi), 'descend');
        for kk = 1:size(est_rank(1:k,:),2)
            hr(kk)=numel(intersect(true_rank(1:k), est_rank(1:k,kk))) / k;
        end
        hit_rate(1)         = mean(hr);
        [~, est_rank]       = sort(sum(abs(phi_shapR.phi),2), 'descend');
        for kk = 1:size(est_rank(1:k,:),2)
            hr(kk)=numel(intersect(true_rank(1:k), est_rank(1:k,kk))) / k;
        end
        hit_rate(2)         = mean(hr);
        [~, est_rank]       = sort(sum(abs(phi_shapK.phi),2), 'descend');
        for kk = 1:size(est_rank(1:k,:),2)
            hr(kk)=numel(intersect(true_rank(1:k), est_rank(1:k,kk))) / k;
        end
        hit_rate(3)         = mean(hr);
        [~, est_rank]       = sort(sum(abs(phi_shapknn.phi),2), 'descend');
        for kk = 1:size(est_rank(1:k,:),2)
            hr(kk)=numel(intersect(true_rank(1:k), est_rank(1:k,kk))) / k;
        end
        hit_rate(4)         = mean(hr);
        [~, est_rank]       = sort(sum(abs(phi_grad),2), 'descend');
        for kk = 1:size(est_rank(1:k,:),2)
            hr(kk)=numel(intersect(true_rank(1:k), est_rank(1:k,kk))) / k;
        end
        hit_rate(5)         = mean(hr);
        [~, est_rank]       = sort(sum(abs(phi_Igrad),2), 'descend');
        for kk = 1:size(est_rank(1:k,:),2)
            hr(kk)=numel(intersect(true_rank(1:k), est_rank(1:k,kk))) / k;
        end
        hit_rate(6)         = mean(hr);
        
    
        %Pearson correlation
        pearson_corr(1) = mean(corr(bench', abs(phi_shap.phi), 'Type', 'Pearson'));
        pearson_corr(2) = mean(corr(bench', abs(phi_shapR.phi), 'Type', 'Pearson'));
        pearson_corr(3) = mean(corr(bench', abs(phi_shapK.phi), 'Type', 'Pearson'));
        pearson_corr(4) = mean(corr(bench', abs(phi_shapknn.phi), 'Type', 'Pearson'));
        pearson_corr(5) = mean(corr(bench', abs(phi_grad), 'Type', 'Pearson'));
        pearson_corr(6) = mean(corr(bench', abs(phi_Igrad), 'Type', 'Pearson'));

        errs(:,mc)                = err;
        infids(:,mc)              = infid;
        hit_rates(:,mc)           = hit_rate;
        pearson_corrs(:,mc)       = pearson_corr;
        
    end
    results(tt).err      = errs;
    results(tt).infid    = infids;
    results(tt).hit_rate = hit_rates;
    results(tt).pearson  = pearson_corrs;
end

save("SIM3_shap_loss_indep_" + string(indep) + "_.mat");

%% gradient based approximation


for tt = 1:length(Ts)
    err(:,tt)     = mean(results(tt).err,2);
    infid(:,tt)   = mean(results(tt).infid,2);
    pearson(:,tt) = mean(results(tt).pearson,2);
    hitrate(:,tt) = mean(results(tt).hit_rate,2);
 end

models = ["SHAP_{Full}", "SHAP_{Sampling}" , "SHAP_{Kernel}" , "SHAP_{KNN}" , "Gradients", "I-Gradients"];
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)]



