% Script to perform Non-Parametric Hidden Markov Models Estimation.
%
% This routine allows to computes the least squares estimators \hat{g}_{K,M}
% of the density of 3 observations of a HMM for all orders K and space dimensions M.
%
% It computes an estimator of the densities and transition matrix by
% incrementing the number of hidden states and using the previous estimate
% as starting point to minimize the empirical contrast function.
%%

%% Random number generator initialization
% Initialize with current time to avoid to get the same results for different instances.
rng('shuffle')

%% Load the simulation parameters from param_simulation.
[kmax, min_m, max_m, approx_type, annotation, stopeval] = param_simulation();

% Load the model parameters from param_HMM.
[Qstar, pistar, param, kast, n, lawtype, lawtypeannotation] = param_HMM();

%% Data initialization and chain generation    
% Chain generation
[X,Y] = func_generateHMM(Qstar, pistar, param, kast, n, lawtype);

% Initialize backup variables.
vartilde = ones(kmax, max_m);		% Variance of estimator (K,M)
contrastvalue = zeros(kmax, max_m);	% Value of the contrast function \gamma_{K,M}
initmethod = cell(kmax, max_m);		% Selected initial point for step (K,M) ("state duplication" or "basis expansion").
									%    Note that for K=1, the value will always be "state duplication".
backup_Mhat = zeros(max_m-min_m+1, max_m, max_m, max_m);
backup_Phat = zeros(max_m-min_m+1, max_m, max_m);
backup_Nhat = zeros(max_m-min_m+1, max_m, max_m);
backup_lhat = zeros(max_m-min_m+1, max_m);
backup_Otilde = zeros(kmax, max_m-min_m+1, max_m, kmax);
backup_Qtilde = zeros(kmax, max_m-min_m+1, kmax, kmax);


%% Estimation : computes \hat{g}_{K,M} for all K and M.
for m = min_m:max_m
    disp( '******************');
    disp(['*     m = ' num2str(m) '     *']);
    disp( '******************');
    
	% Initialize estimators for m-th estimation space.
	[x, basePM, Ostar, fstar, Mth, Mhat, Phat, Nhat, lhat, Qhat, Ohat] = func_prepareForPM(m, Y, approx_type, Qstar, param, kast, lawtype);
	% Save empirical estimator.
	backup_Mhat(m-min_m+1, 1:m, 1:m, 1:m) = Mhat;
	backup_Phat(m-min_m+1, 1:m, 1:m) = Phat;
	backup_Nhat(m-min_m+1, 1:m, 1:m) = Nhat;
	backup_lhat(m-min_m+1, 1:m) = lhat;
    
    for k=1:min(kmax,m-1)
        disp(['----- k = ' num2str(k) ' -----']);
        N = k*(m+k-1);     % number of objective variables/problem dimension for CMAES

        %% Initial point: state duplication
        indiceMin = 1;
        if (k>1)
		    % If k>1, duplicate the first state from the estimator of the model with k-1 states and the same m.
            [Qhat, Ohat] = func_duplicateState(Qtilde, Otilde, 1);
        end
        xmin = zeros(k,N);
        contrastMin = func_empiricalLoss(Mhat, reshape([Ohat; (Qhat(:,1:k-1))'],[N,1]),k);
        for duplicationIndex = 1:max(k-1,1)
		    % Duplicate state "duplicationIndex" from the estimator of the model with k-1 states and the same m.
            if (k>1)
                disp(['=== Duplicates state ' num2str(duplicationIndex) ' ===']);
                [Qhat, Ohat] = func_duplicateState(Qtilde, Otilde, duplicationIndex);
            else
                disp(['=== Initial estimation ===']);
            end

			% Minimization step.
            [xminTemp, contrastTemp] = func_CMAESchrono(m, k, Qhat, Mhat, Ohat, stopeval);
            
            xmin(duplicationIndex,:) = xminTemp;
            contrastAct = contrastTemp;
            if (contrastAct < contrastMin)
				% If this initialization was better, keep it.
				contrastMin = contrastAct;
                indiceMin = duplicationIndex;
            end
        end
        
        %% Initial point: basis expansion
        contrast_expansion = contrastMin + 1;
        if m > min_m
            disp(['=== Expands basis : (k,m) = (' num2str(k) ',' num2str(m-1) ') -> (' num2str(k) ',' num2str(m) ') ===']);
            Ohat = squeeze(backup_Otilde(k, m-min_m, 1:m, 1:k));
            Qhat = squeeze(backup_Qtilde(k, m-min_m, 1:k, 1:k));
            
            [xminTemp, contrastTemp] = func_CMAESchrono(m, k, Qhat, Mhat, Ohat, stopeval);
            
            xmin_expansion = xminTemp;
            contrast_expansion = contrastTemp;
        end
        
        %% Synthesis
        if contrastMin <= contrast_expansion
            Atilde = reshape(xmin(indiceMin,:),[m+k-1,k]);
            disp('     Selected initial point : state duplication');
            initmethod{k,m} = 'state duplication';
        else
            Atilde = reshape(xmin_expansion   ,[m+k-1,k]);
            disp('     Selected initial point : basis expansion');
            initmethod{k,m} = 'basis expansion';
        end
        Otilde = Atilde(1:m,:);

        Qtilde = zeros(k);
        Qtilde(:, 1:k-1) = (Atilde(m+1:m+k-1,:))';
        Qtilde(:, k) = ones(k,1) - sum(Qtilde, 2);

        % Saves estimators
        backup_Otilde(k, m-min_m+1, 1:m, 1:k) = Otilde;
        backup_Qtilde(k, m-min_m+1, 1:k, 1:k) = Qtilde;
        
        vartilde(k,m) = func_theoreticalContrast(Qtilde,Mth,Otilde);
        contrastvalue(k,m) = contrastMin;
    end
end

%% Save simulation results.
mkdir('data');
save(['data/' annotation '_contrast_' approx_type '_' lawtypeannotation '_n=' num2str(n) '_kmax=' num2str(kmax) '_minm=' num2str(min_m) '_maxm=' num2str(max_m)], ...
    'min_m', 'max_m', 'kmax', 'contrastvalue', 'initmethod', 'n', 'vartilde', 'backup_Mhat', 'backup_Phat', ...
    'backup_Nhat', 'backup_lhat', 'backup_Otilde', 'backup_Qtilde', 'approx_type', 'X', 'Y', 'Qstar', 'param', ...
    'kast', 'annotation', 'lawtype', 'lawtypeannotation');
