%% Preliminary Output Analysis for Bayesian ANOVA with missing values. 
% DESIGNED FOR TESTING, WHERE WE KNOW THE TRUE VALUES. 
clear all; close all;
%% get the 3 pers of the imputed values:
Imputation_Pers_Extractor
clear all; close all;
%% LOAD the true pars, factor levels, and complete data array:
load Data_Array_Complete
load TRUE_FLS
load TRUE_PARS

%% load the original, incomplete data array:
load Data_INPUT
% Calculate relevant quantities:
N_levels=size(Data_INPUT);
N_factors=length(N_levels);

% load the priors:
load PRIORS
%% Load the samples of the parameters and factor levels:
load PARS_ARRAY
load FLS_ARRAY

% get the number of samples
N_Its_ALL=length([PARS_ARRAY.GM]);
burnin=min(600, 0.5*N_Its_ALL);
N_Its=N_Its_ALL-burnin;
%% PLOT the time series of scalar parameters, along with the true values:
figure(1), clf
N_SPs=length(fields(PARS_ARRAY));
%GRAND MEAN
subplot(N_SPs, 1, 1)
plot(1:1:N_Its_ALL, [PARS_ARRAY.GM], 'k.-', 'linewidth', 2, 'markersize', 4), hold on
plot(1:1:N_Its_ALL, ones(1, N_Its_ALL)*TRUE_PARS.GM, 'r-', 'linewidth', 1), hold on
plot(1:1:N_Its_ALL, ones(1, N_Its_ALL)*mean(Data_Array_Complete(:)), 'r--', 'linewidth', 1), hold on
ylabel('Grand Mean')

%ERROR VARIANCE
subplot(N_SPs, 1, 2)
plot(1:1:N_Its_ALL, [PARS_ARRAY.sigma2], 'k.-', 'linewidth', 2, 'markersize', 4), hold on
plot(1:1:N_Its_ALL, ones(1, N_Its_ALL)*TRUE_PARS.sigma2, 'r-', 'linewidth', 1), hold on
ylabel('Error Variance')

%REDCUDED SPACE FACTOR LEVEL VARIANCES
for KK=1:1:N_factors
    subplot(N_SPs, 1, 2+KK)
    plot(1:1:N_Its_ALL, eval(['[PARS_ARRAY.Red_space_var_', num2str(KK), ']']), 'k.-', 'linewidth', 2, 'markersize', 4), hold on
    plot(1:1:N_Its_ALL, ones(1, N_Its_ALL)*[eval(['TRUE_PARS.Red_space_var_', num2str(KK)])], 'r-', 'linewidth', 1), hold on
    ylabel(['Fact. Var. ', num2str(KK)])
end


%% THROWOUT the burnin for the parameter
PARS_ARRAY_Good=PARS_ARRAY(burnin+1:1:end);
%and the factor levels:
FLS_ARRAY_Good=FLS_ARRAY(burnin+1:1:end);

%% PLOT histograms of the parameters:
fonter=15;
N_prior_plot=200;

figure(2), clf
%GRAND MEAN
subplot(N_SPs,1,1)
set(gca, 'FontSize', fonter)
[cs_R1, vs_R1]=hist([PARS_ARRAY_Good(:).GM], 15);
%plot the prior: Normal
ax_vals=[vs_R1(1)-diff(vs_R1(1:1:2))  vs_R1(end)+diff(vs_R1(1:1:2)) 0 1.1*max(cs_R1)/(sum(cs_R1)*diff(vs_R1(1:1:2)))];
X1=linspace(ax_vals(1), ax_vals(2), N_prior_plot);
prior_1=normpdf(X1, PRIORS.GM(1), sqrt(PRIORS.GM(2)));
plot(X1, prior_1, 'k--', 'LineWidth', 2.5, 'color', [0.850,0.0,0.850]), hold on
barplot_tops(cs_R1/(sum(cs_R1)*diff(vs_R1(1:1:2))), vs_R1, 3, 1, [0.00,0.00,0.00]), hold on
axis(ax_vals)
ylabel('Grand Mean', 'FontSize', fonter);
%plot truth:
plot(ones(1,2)*TRUE_PARS.GM, ax_vals(3:1:4), 'r-', 'LineWidth', 2.5), hold on

%ERROR VARIANCE:
subplot(N_SPs,1,2)
set(gca, 'FontSize', fonter)
%plot the prior
[cs_R1, vs_R1]=hist([PARS_ARRAY_Good(:).sigma2], 15);
ax_vals=[vs_R1(1)-diff(vs_R1(1:1:2))  vs_R1(end)+diff(vs_R1(1:1:2)) 0 1.1*max(cs_R1)/(sum(cs_R1)*diff(vs_R1(1:1:2)))];
%plot the prior: inverse gamma with parameters, using the notation of the pink book, sigma2_prior=[0.5, 0.5];
X2_pre=linspace(ax_vals(1), ax_vals(2), N_prior_plot);
X2_good_ins=find(X2_pre>0);
X2=X2_pre(X2_good_ins);
prior_2=InverseGamma_pdf(PRIORS.sigma2(1), PRIORS.sigma2(2), X2);
plot(X2, prior_2, 'k--', 'LineWidth', 2.5, 'color', [0.850,0.0,0.850]), hold on
barplot_tops(cs_R1/(sum(cs_R1)*diff(vs_R1(1:1:2))), vs_R1, 3, 1, [0.00,0.00,0.00]), hold on
axis(ax_vals)
ylabel('Err Var', 'FontSize', fonter);
%plot truth:
plot(ones(1,2)*TRUE_PARS.sigma2, ax_vals(3:1:4), 'r-', 'LineWidth', 2.5), hold on

%FACTOR VARIANCES:
for KK=1:1:N_factors
    subplot(N_SPs,1,2+KK)
    set(gca, 'FontSize', fonter)
    %plot the prior
    [cs_R1, vs_R1]=hist(eval(['[PARS_ARRAY_Good.Red_space_var_', num2str(KK), ']']), 15);
    ax_vals=[vs_R1(1)-diff(vs_R1(1:1:2))  vs_R1(end)+diff(vs_R1(1:1:2)) 0 1.1*max(cs_R1)/(sum(cs_R1)*diff(vs_R1(1:1:2)))];
    ax_vals(1)=0.05; %TO HELP TROUBLE SHOOT AN APPARENT PROBLEM
    XK_pre=linspace(ax_vals(1), ax_vals(2), N_prior_plot);
    XK_good_ins=find(XK_pre>0);
    XK=XK_pre(XK_good_ins);
    prior_K=InverseGamma_pdf(eval(['[PRIORS.Red_space_var_', num2str(KK), '(1)]']), eval(['[PRIORS.Red_space_var_', num2str(KK), '(2)]']), XK);
    plot(XK, prior_K, 'k--', 'LineWidth', 2.5, 'color', [0.850,0.0,0.850]), hold on
    barplot_tops(cs_R1/(sum(cs_R1)*diff(vs_R1(1:1:2))), vs_R1, 3, 1, [0.00,0.00,0.00]), hold on
    axis(ax_vals)
    ylabel(['Red. Space Var ', num2str(KK)], 'FontSize', fonter);
    %plot truth:
    plot(ones(1,2)*eval(['[TRUE_PARS.Red_space_var_', num2str(KK), ']']), ax_vals(3:1:4), 'r-', 'LineWidth', 2.5), hold on
end


%% GET percentiles for the factor levels:
pers3=round([0.05, 0.5, 0.95]*N_Its);
FLS_ARRAY_3Pers=struct('fact1', []);
for KK=1:1:N_factors
    tmp=eval(['[FLS_ARRAY_Good.Levels_', num2str(KK), ']']);
    tmp2=sort(tmp, 2);
    tmp_pers=tmp2(:, pers3);
    FLS_ARRAY_3Pers=setfield(FLS_ARRAY_3Pers, ['fact', num2str(KK)], tmp_pers);
    clear tmp tmp2 tmp_pers
end

%% PLOT THE FACTOR LEVELS, UNCERTAINTIES
fonter=15;
fill_trip=[0.7, 0.7, 0.95];

figure(3), clf
for KK=1:1:N_factors
    subplot(N_factors,1,KK)
    set(gca, 'FontSize', fonter)
    
    pers_temp=eval(['FLS_ARRAY_3Pers.fact', num2str(KK),]);
    var(pers_temp(:,2))
    
    fill([1:1:N_levels(KK), fliplr(1:1:N_levels(KK))], [pers_temp(:,1)',fliplr(pers_temp(:,3)')], fill_trip, 'linewidth', 2, 'edgecolor', fill_trip), hold on    
    plot(1:1:N_levels(KK), zeros(size(1:1:N_levels(KK))),  'r-', 'linewidth', 1.5, 'color', [1,1,1]*0.7), hold on
    plot(1:1:N_levels(KK), pers_temp(:,2), 'k-', 'linewidth', 2.5), hold on
    
    axis([0 N_levels(KK)+1 [-1.1, 1.1]*max(abs(pers_temp(:)))]);
    ylabel(['Red. Space Var ', num2str(KK)], 'FontSize', fonter);
end

%plot the truth
for KK=1:1:N_factors
    subplot(N_factors,1,KK)
    plot(1:1:N_levels(KK), eval(['TRUE_FLS.Levels_', num2str(KK)]), 'r--', 'linewidth', 2.5), hold on
end

%% LOAD the three pers of the imputed values:
load Samples_of_Factor_Array_pers3
% get the indices of the missing values from the Data_INPUT:
inds_Miss=find(isnan(Data_INPUT)==1);

% get the withheld values:
Data_WH=Data_Array_Complete(inds_Miss);
N_WH=length(Data_WH);
%for ease of plotting, sort:
[vals, inds]=sort(Data_WH);

% get the posterior percentiles:
Temper='Samples_of_Factor_Array_pers3(';
for KK=1:1:N_factors
    Temper=[Temper, ':,'];
end
Temper_L=[Temper, '1)'];
Temper_M=[Temper, '2)'];
Temper_H=[Temper, '3)'];

lows=eval(Temper_L);
meds=eval(Temper_M);
highs=eval(Temper_H);

lows_Data_WH=lows(inds_Miss);
meds_Data_WH=meds(inds_Miss);
highs_Data_WH=highs(inds_Miss);


%% plot the withheld values, from smallest to largest:
fill_trip=[0.7, 0.7, 0.95];
figure(4), clf
plot(1:1:N_WH, Data_WH(inds), 'r--', 'linewidth', 2)

% plot the posterior uncertainty bounds:
fill([1:1:N_WH, fliplr(1:1:N_WH)], [lows_Data_WH(inds)',fliplr(highs_Data_WH(inds)')], fill_trip, 'linewidth', 2, 'edgecolor', fill_trip), hold on    

% posterior median:
plot(1:1:N_WH, meds_Data_WH(inds), 'k-', 'linewidth', 2.5)

% withheld values:
plot(1:1:N_WH, Data_WH(inds), 'r--', 'linewidth', 2)

% coverage rate:
inds_nocov=find(lows_Data_WH(inds)-Data_WH(inds)>0 | highs_Data_WH(inds)-Data_WH(inds)<0);
cov_rate=1-length(inds_nocov)/length(inds)
%indicate where the bounds DO NOT cover:
plot(inds_nocov, lows_Data_WH(inds(inds_nocov)), 'b.')
plot(inds_nocov, highs_Data_WH(inds(inds_nocov)), 'b.')

title('Withheld values (red), Posterior Median (black), 90%  CI (shading)') 
xlabel(['Coverage rate of 90% CIs is ', num2str(round(cov_rate*100)), '%. Blue dots = failures.'])

%axis([0 N_WH+1 min(lows_Data_WH)*1.1 max(highs_Data_WH)*1.1 ])