% Last modified: 06/09/2012 by Pete

%% Stochastic Flow Past an Obstacle Observation Script

%% fit

Xf = [ u ; v ] ;        % forecast mean
Xi = [ ui ; vi ] ;      % DO modes
phi = YYt ;             % DO mode coefficient MC particles

% if EnKF, fit Gaussian
if ( EnKF )
    
    % estimate covariance of DO mode coefficients using sample covariance
    covphi = phi'*phi / length(phi) ;
    
    % render Gaussian as single-component GMM for use with GMM-DO filter
    M     = 1 ;             % one component
    pijf  = 1 ;             % weight
    mujf  = zeros(S,1) ;    % mean
    Sigjf = covphi ;        % covariance
    
% else, fit GMM
else
    
    % set GMM fit parameters
    mmin = 01 ;                 % minimum number of mixture components
    mmax = 50 ;                 % maximum number of mixture components
    MaxIter = 2e04 ;            % maximum number of EM iterations
    TolFun = 1e-07 ;            % EM termination tolerance
    Replicates = 50 ;           % number of EM random starts
    fits = cell(mmax,1) ;       % maximum likelihood GMM fits
    BICs = Inf*zeros(mmax,1) ;  % BICs of maximum likelihood GMM fits
    
    % fit phi GMM
    warning('off','stats:gmdistribution:IllCondCov')
    for m = mmin:mmax
        try
            tic
            fits{m} = gmdistribution.fit( phi, m, ...
                'Replicates', Replicates, ...
                'Options', statset('MaxIter',MaxIter,'TolFun',TolFun) ) ;
            BICs(m) = fits{m}.BIC ;
            fprintf('GMM %02d fit: t = %2.2f s, LL = %f\n', m, toc, ...
                -fits{m}.NlogL )
        catch err
            fprintf('GMM %02d fit: t = %2.2f s, error\n', m, toc)
        end
    end
    
    [~,M] = min(BICs) ;         % number of components with minimum BIC
    GMMf  = fits{M} ;           % prior phi GMM
    pijf  = GMMf.PComponents ;  % component weights
    mujf  = GMMf.mu' ;          % component means (note: transposed)
    Sigjf = GMMf.Sigma ;        % component covariances
    
end


%% make observation

% load observations
load( sprintf( '%s/obs/%05d.mat', savedir, 100*t ), 'obs' )

% determine subspace observation matrix, innovation vector, and subspace
% gain matrices
Htil  = H*Xi ;
ytil  = obs - H*Xf ;
Kjtil = zeros(S,n_obs,M) ;
for j = 1:M
    Kjtil(:,:,j) = Sigjf(:,:,j)*Htil'*pinv(Htil*Sigjf(:,:,j)*Htil'+R) ;
end

% determine intermediate GMM component means
muja_hat = zeros(S,M);
for j = 1:M
    muja_hat(:,j) = mujf(:,j) + Kjtil(:,:,j)*(ytil - Htil*mujf(:,j)) ;
end

% update GMM mixture weights
pija = zeros(1,M) ;
HPHt_til = zeros(n_obs,n_obs,M) ;
for j = 1:M
    HPHt_til(:,:,j) = Htil*Sigjf(:,:,j)*Htil' ;
    HPHt_til(:,:,j) = (HPHt_til(:,:,j)+HPHt_til(:,:,j)')/2 ;
    pija(j) = pijf(j)*mvnpdf( ytil, Htil*mujf(:,j), HPHt_til(:,:,j) + R ) ;
end
pija = pija / sum(pija) ;

% update state variable mean
Xa = Xf + Xi*(muja_hat*pija') ;

% update GMM component means
muja = zeros(S,M) ;
for j = 1:M
    muja(:,j) = muja_hat(:,j) - muja_hat*pija' ;
end

% update GMM component covariances
Sigja = zeros(S,S,M) ;
for j = 1:M
    Sigja(:,:,j) = (eye(S,S) - Kjtil(:,:,j)*Htil)*Sigjf(:,:,j) ;
    Sigja(:,:,j) = (Sigja(:,:,j)+Sigja(:,:,j)')/2 ;
end

% draw MC particles from posterior phi GMM
GMMa = gmdistribution( muja', Sigja, pija ) ;
phi = random( GMMa, MC ) ;

% project prior phi GMM into marginalized state space and convolute
% observation noise to find prior observation GMM
piobs = pijf ;
muobs = H*Xf*ones(1,M) + Htil*mujf ;
Sigobs = zeros(n_obs,n_obs,M) ;
for j = 1:M
    Sigobs(:,:,j) = Htil*Sigjf(:,:,j)*Htil' + (modelnoise^2)*eye(n_obs) ;
end
GMMobs = gmdistribution( muobs', Sigobs, piobs ) ;

% calculate log-likelihood of observations
loglikes = [ loglikes ; t , log( pdf( GMMobs , obs' ) ) ] ;

% update FV solver variables
u = Xa(1:n_uvars) ;
v = Xa(n_uvars+1:end) ;
YYt = phi ;

