%% CONTROL CODE FOR BAYESIAN ANOVA
% with 1 or 0 observations at each combination of factor levels, 
% and at least one observation at one level of the Nth factor, when the
% levels of all other factors are fixed. 
% (In other words, if the input is an N dimenstional cube, any particular
% colmun or row or . . . will contain at least one observation). 

%SET MATLAB'S CURRENT DIRECTORY TO HERE. 

clear all; close all;


% Number of Samples
Sampler_Its=2600;

%% Set the seed of the random number generators
randn('state', sum((1000+600)*clock))
rand('twister', sum((1000+800)*clock))
% set the priors and the inital values for the MCMC sampler
Prior_Maker
Initial_Vals_Maker


%% load the data
load Data_INPUT

% Calculate relevant quantities:
N_levels=size(Data_INPUT);
N_factors=length(N_levels);
% % %It will also be convenient to have the cumulative sum of the elements of
% % %N_levels:
% % CSNL=cumsum(N_levels);
% % CSNL_Ext=[0, CSNL];
%% Load the priors and the initial value
load PRIORS
load INITIAL_VALS


%% A NOTE ON GAMMA NOTATION. WE USE THE NOTATION OF Gelman et al, "Bayesian
% Data Analysis", WHERE GAMMA PARAMETERS ALPHA, BETA)==(SHAPE, INVERSE SCALE). 
% THE RANDRAW.M CODE USES (A,B)==(SHAPE, SCALE), AND THE CALL IS RANDRAW('GAMMA', [M,B,A], SAMPLESIZE), 
% WHERE M IS THE LOCATION (NOT NEEDED). SO IN THE NOTATION OF GELMAN ET AT, THE CALL IS
% RANDRAW('GAMMA', [0,1/BETA,ALPHA], SAMPLESIZE). 
% For example,
% RANDRAW('GAMMA', [0,1/PRIORS.sigma2(2),PRIORS.sigma2(1)], 1), AND ETC. 



%% ESTABLISH THE SAMPLING ORDER FOR THE SCALAR PARS AND FACTOR LEVELS: 
%The Order of THE SCALAR parameters WILL ALWAYS thus:
% GRAND MEAN (gamma, in the notes). 
% ERROR VARIANCE (sigma^2). 
% The REDUCED SPACE VARIANCE corresponding to each factor, in the order of
% the dimensions of Data_INPUT. 
% DEFINE a structure which will contain these elements. This will be
% updated at each iteration.

% INITIALIZE with values from INITIAL_VALS
CURRENT_PARS.GM=INITIAL_VALS.GM;
CURRENT_PARS.sigma2=INITIAL_VALS.sigma2;
for kk=1:1:N_factors
    CURRENT_PARS=setfield(CURRENT_PARS, ['Red_space_var_', num2str(kk)], eval(['INITIAL_VALS.Red_space_var_',num2str(kk)]));
end

%% DO SOMETHING SIMILAR FOR THE FACTORS:
CURRENT_FLS=struct('Levels_1', []);
for kk=1:1:N_factors
    CURRENT_FLS=setfield(CURRENT_FLS, ['Levels_', num2str(kk)], eval(['INITIAL_VALS.Levels_',num2str(kk)]));
end

% OR LOAD TRUE VALUES - FOR TESTING
% load Factor_Vector.mat
% load Factors_ReducedSpace_Stds.mat
% load Grand_Mean
% load Noise_std.mat
% load Data_Array_Complete.mat

%% DEFINE THE EMPTY PARAMETER ARRAY, which will be filled:
PARS_ARRAY=struct('GM',[],'sigma2',[]);
for kk=1:1:N_factors
    PARS_ARRAY=setfield(PARS_ARRAY, ['Red_space_var_', num2str(kk)], []);
end

PARS_ARRAY(Sampler_Its).GM=[];
%% DEFINE THE EMPTY FACTOR LEVEL ARRAY, which will be filled:
FLS_ARRAY=struct('Levels_1', []);
for kk=1:1:N_factors
    FLS_ARRAY=setfield(FLS_ARRAY, ['Levels_', num2str(kk)], []);
end

FLS_ARRAY(Sampler_Its).Levels_1=[];
%% Find the indices of Data_INPUT where values are observed and missing:
Inds_MISSING =find(isnan(Data_INPUT)==1);
Inds_OBSERVED=find(isnan(Data_INPUT)==0);
% 
% We will define an array of NaNs, Current_Imputed_Values, with dimensions
% given by N_levels. at each iteration of the sampler, we will fill the
% elements corresponding to the missing values of Data_INPUT, and save it. 


% We will also need a 'complete' data array - Data_INPUT with missing
% values imputed - to update the other parameters:
Current_Complete_Data=Data_INPUT;
% The other elements will be filled at the beginning of the first iteration,
% and then updated at each subsequent iteration

%% CALCULATE Arrays of size N_levels, corresponding to each factor. 
% The 'best estimate' of the missing values is then given by the sum of
% these and the Grand Mean.

% Make use of a function, which creates an (N_levels, N_factors)dimensional
% array. 
% Think: an N_factors stack of arrays, each corresponding to the levels of
% one of the factors. 
% Adding along the N_factors+1 dimension, and then adding the Grand Mean,
% gives the best estimates of the missing values, given the factor levels. 

CURRENT_FAC_ARRAY=Levels_Array(CURRENT_FLS);
CURRENT_FAC_ARRAY_Sum=sum(CURRENT_FAC_ARRAY, N_factors+1);

%% RUN THE SAMPLER
for samples=1:1:Sampler_Its
    tic;
    %% SAMPLE the missing values. 
    Current_Imputed_Values=Missing_Imputer(CURRENT_FAC_ARRAY_Sum, CURRENT_PARS, Inds_OBSERVED);
    % Combine these with the values from Data_INPUT:
    Current_Complete_Data(Inds_MISSING)=Current_Imputed_Values(Inds_MISSING);
    
    %% SAMPLE the grand mean:
    New_GM=GrandMean_Updater(Current_Complete_Data, CURRENT_PARS, PRIORS);
    CURRENT_PARS.GM=New_GM;
    clear New_GM
    
    %% SAMPLE the error variance
    New_sigma2=Error_Var_Updater(Current_Complete_Data, CURRENT_FAC_ARRAY_Sum, CURRENT_PARS, PRIORS);
    CURRENT_PARS.sigma2=New_sigma2;
    clear New_sigma2
    
    %% SAMPLE the reduced space variance for the levels of each factor 
    for KK=1:1:N_factors
        New_Reduced_Space_Variance=Reduced_Space_Variance_Updater(CURRENT_FLS, PRIORS, KK);
        CURRENT_PARS=setfield(CURRENT_PARS, ['Red_space_var_', num2str(KK)], New_Reduced_Space_Variance);        
        clear New_Reduced_Space_Variance
    end
    
    %% SAMPLE the levels of the factors
    for KK=1:1:N_factors
        New_Factor_Levels=Factor_Levels_Updater(Current_Complete_Data, CURRENT_PARS, PRIORS, KK);
        CURRENT_FLS=setfield(CURRENT_FLS, ['Levels_', num2str(KK)], New_Factor_Levels);        
        clear New_Factor_Levels
    end
    
   
    %% UPDATE THE FACTOR ARRAY AND SUM:
    CURRENT_FAC_ARRAY=Levels_Array(CURRENT_FLS);
    CURRENT_FAC_ARRAY_Sum=sum(CURRENT_FAC_ARRAY, N_factors+1);

    %% UPDATE THE VARIOUS ARRAYS, SAVE CURRENT IMPUTATIONS OF THE MISSING
    % VALUES
    
    % UPDATE the parameter array:
    PARS_ARRAY(samples)=CURRENT_PARS;
    % UPDATE the factor levels array:
    FLS_ARRAY(samples)=CURRENT_FLS;
    
    % save them:
    save Imputations/PARS_ARRAY PARS_ARRAY
    save Imputations/FLS_ARRAY FLS_ARRAY
    
    % save Current_Imputed_Values, adding an index indicating the sample
    % number:
    save(['Imputations/Imputed_Values_Step' num2str(samples)],'Current_Imputed_Values');    
    
    timertimer=toc;
    disp(['Finished MCMC iteration ', num2str(samples), ' of ', num2str(Sampler_Its), '. Last iteration took ', num2str(timertimer), ' seconds.'])
end

%% SAVE some stuff to the current directory:
    save PARS_ARRAY PARS_ARRAY
    save FLS_ARRAY FLS_ARRAY


