%%%%%%%%%%%%%%%%%%%%%
%%%% Main script %%%%
%%%%%%%%%%%%%%%%%%%%%

% This script is used to: 
% 1) read the subsample files containing the firm-destination time series
%    from the origination folder;
% 2) perform the estimation of a block factor model on each data matrix; 
% 3) save the result of the estimation on a user defined destination
%    folder.

%%%%%%%%%%%%%%%%%%%%

%%%%%%% Clear and prepare to cycle over the dataset subsamples

clear

% Set Base_Dir
Base_Dir = cd;
addpath(genpath(Base_Dir));


% Read the names of the folders containing the input files
input_folders = dir('Inputs');
folder_names = {input_folders.name};


% Set the number of subsamples (this needs to be lower or equal to the
% number of files in fold_name)
H = 20


for i = 3:length(folder_names) %% cycle over all the subfolders in the folder 'Input'
    
   fold_name = folder_names{i};
    
    for l = 0:H   %%% insert parfor if you want to run parallel estimates

        disp(['Simulation number: ',num2str(l)])
        file = ['Inputs',filesep,fold_name,filesep,'samples',filesep,'data_sample_gr',num2str(l),'.csv']

        base_import = readtable(file,'ReadVariableNames',true,'TreatAsEmpty','NA'); %% import data
        X_def = base_import{1:end,4:end}; %% import the data matrix contained in each subsample file
        X_char = base_import(1:end,2:3); %% import string information contained in each subsample file

        X_def = X_def';

        group_labels = table2array(X_char(:,2)); %% entries in the 2nd column are block(destination)-specific 
        siren_labels = table2array(X_char(:,1)); %% entries in the 1st column are firm specific (SIREN is the data original data filed)

        X_def_t = X_def; % rolling window restriction (no restriction is applied now)

        ind_not_all_nan = find(sum(isnan(X_def_t)) <= size(X_def_t,1)-2); %% Which series have only two data points?

        X_def_t = X_def_t(:,ind_not_all_nan); % exclude serie with all NaN or all NaN but one and two
        group_labels_t = string(group_labels(ind_not_all_nan));
        siren_labels_t = string(siren_labels(ind_not_all_nan));

        un_siren = unique(siren_labels_t); % store firms identifiers
        sg_labels_t = [siren_labels_t group_labels_t]; % store firm-destination cell identifiers
        groups = unique(group_labels_t); % store block (destination) identifiers


    %%%%%%  Create the structure of the model and store it in a cell-array (index)

        index = {};
        q = 1;
        while(q <= length(groups)) %% cycle over the blocks (destinations)

            block = groups(q); % read the block identifier

            index_bool = strfind(group_labels_t, block); % read the indices of the time series belonging to the block
            index{q,1} = block; % storing the block name
            index{q,2} = find(not(cellfun('isempty',index_bool))); % store the indices in the cell-array
            index{q,3} = 1; %% store the number of factors (r) in the block
            index{q,4} = 1; %% store the number of lags (p) of the block-specific AR process

            q = q + 1; 
        end

    %%%%% Define and store the global parameters 

        P = []; % global parameters for the estimation

        P.r = 1; % number of global factors
        P.p = 1; % number of lags in the global factors VAR
        P.max_iter = 100; % maximal number of EM iterations
        P.conv_thresh = 10^-3; % threshold condition for the EM algorithm


    %%%%%% Estimate the model

        blocks = index;

        res = EM_BDFM(X_def_t,P,blocks); %% estimate the model


    %%%%%% Store the results

        % Create a directory where to save the output files
        mkdir(['..',filesep,'results_analysis',filesep,'Inputs',filesep,fold_name]);

        % String to select the folder where to store
        save_heading = ['..',filesep,'results_analysis',filesep,'Inputs',filesep,fold_name,filesep,'results_',num2str(l)];

        % Store the estimated factors in a table
        F_tab = array2table(res.F,"VariableNames",['common'; cellstr(blocks(:,1))]); 

        % Store the estimated factors' loadings in a table
        C_tab = array2table(res.C,"VariableNames",['common'; cellstr(blocks(:,1))]); 
        C_tab = [array2table(sg_labels_t, "VariableNames",{'SIREN' 'PYOD'}), C_tab]; %% add SIREN (firmID) and PYYOD (countryID)

        % Add the time series means 
        M_tab = [array2table(sg_labels_t, "VariableNames",{'SIREN' 'PYOD'}) array2table(res.Mx')];

        % Add the time series variances
        W_tab = [array2table(sg_labels_t, "VariableNames",{'SIREN' 'PYOD'}) array2table(res.Wx')];

        % Write the outuputs to .csv files
        writetable(F_tab,[save_heading '_F.csv'])
        writetable(C_tab,[save_heading '_C.csv'])
        writetable(M_tab,[save_heading '_M.csv'])
        writetable(W_tab,[save_heading '_W.csv'])

    end
end
