% Last modified: 06/09/2012 by Pete

% truereac      true reaction equation of tracer in system [1,9]
% Re            Reynolds number
% randseed      rng seed for observation noise realization (set to -1 to
%               suppress addition of noise to observations)
% render        1 = generate figures,
%               0 = do not generate figures

function Pete_RunEns( truereac, Re, randseed, render )

app.toptic = tic ;
fprintf(['Stochastic microorganism tracer system: ',...
    'reac ensemble.\n'])
close all

trunk = regexp(pwd,'.*trunk','match','once') ; Pete_src
app.tag = sprintf('tracer_truereac%d_Re%d_randseed%d', ...
    truereac, Re, randseed) ;
savemain = [trunk,'/Save/',app.tag] ;
addpath(genpath([trunk,'/Cases/Pete_tracer']))
load( [ savemain, '/1/param.mat' ], 'app' )

modelnames = { 'g_{+}(t) , d_{-}(r)', ...
               'g_{+}(t) , d_{0}', ...
               'g_{+}(t) , d_{+}(r)', ...
               'g_{0} , d_{-}(r)', ...
               'g_{0} , d_{0}', ...
               'g_{0} , d_{+}(r)', ...
               'g_{-}(t) , d_{-}(r)', ...
               'g_{-}(t) , d_{0}', ...
               'g_{-}(t) , d_{+}(r)' } ;

Q = [] ;
for i = 1:9
    load( sprintf('%s/loglikes%d.mat',savemain,i), 'loglikes' )
    Q = [ Q , loglikes(:,2) ] ;
end

% calculate distributions over modelnames
P = [ (1/9)*ones(1,9) ; zeros(size(Q)) ] ;

for i = 1:size(Q,1)
    
    for j = 1:9
        den = 0 ;
        for k = 1:9
            den = den + exp(Q(i,k)-Q(i,j))*P(i,k) ;
        end
        P(i+1,j) = P(i,j) / den ;
    end
    P(i+1,:) = P(i+1,:) / sum(P(i+1,:)) ;
    
end

% create new directory for ensemble means
savedir = [ savemain, '/mean' ] ;
if ( ~exist(savedir,'dir') )
    mkdir( savedir )
end

% find all save steps
steps = dir( [ savemain, '/true/*0.mat' ] ) ;
n_steps = size(steps,1) ;

% load parameters common to all modelnames
load( [ savemain, '/true/param.mat' ], 'a', 'b', 'Nx', 'Ny', ...
    'dx', 'dy', 'XP', 'YP', 'XU', 'YU', 'XV', 'YV', 'XRho', 'YRho', ...
    'dt', 'SaveIntrvl', 'Nbcs' )

times = zeros(n_steps,1) ;
P_t = zeros(n_steps,9) ;
RMSE = zeros(n_steps,1) ;

figshift = [-0.10,0.05,0,0] ;

for i = 1:n_steps
    
    times(i) = str2double(steps(i).name(1:5))/100 ;
    
    if ( render )
        
        figure(1), clf
        set(1,'Position',[10,50,1800,1060],'PaperPositionMode','auto')
        warning('off',...
            'MATLAB:hg:patch:ColorDataNotSetForInterpolatedShading')
        
    end
    
    % TRUE STATE ----------------------------------------------------------
    
    % load true state variables and parameters
    load( [ savemain, '/true/', steps(i).name ], ...
        'u_true', 'v_true', 'rho_true' )
    load( [ savemain, '/true/param.mat' ], ...
        'Nodeu', 'Nodev', 'NodeRho' )
    
    % set obstacle velocity and rho values to NaN
    u_true = u_true ; v_true = v_true ; rho_true = rho_true ;
    u_tmp = [ NaN ; u_true(2:end) ] ;
    v_tmp = [ NaN ; v_true(2:end) ] ;
    rho_tmp = [ NaN ; rho_true(2:end) ] ;
    
    if ( render )
        
        % render flow field with rho values and velocity streamlines
        subplot(4,3,1)
        titlestr = 'true state' ;
        Pete_plotfield
        
        % plot sampling locations
        hold on
        scatter( app.SampleLocs(:,1), app.SampleLocs(:,2), ...
            25, [1.0,0.4,0.2], 'filled')
        hold off
        
        set(gca,'Position',get(gca,'Position')+figshift)
        
    end
    
    % MODEL MEANS ---------------------------------------------------------
    
    % initialize reac ensemble mean variables
    u_mean = zeros(size(u_tmp)) ;
    v_mean = zeros(size(v_tmp)) ;
    rho_mean = zeros(size(rho_tmp)) ;
    
    % plot reac flow fields
    for j = 1:9
        
        % load reac state variables and parameters
        load( sprintf( '%s/%d/%s', savemain, j, steps(i).name ), ...
            'u', 'v', 'rho', 'YYt', 'samplestep' )
        load( sprintf( '%s/%d/param.mat', savemain, j ), ...
            'Nodeu', 'Nodev', 'NodeRho' )
        
        % set obstacle velocity and rho values to NaN
        u = u ; v = v ; rho = rho ;
        u_tmp = [ NaN ; u(2:end) ] ;
        v_tmp = [ NaN ; v(2:end) ] ;
        rho_tmp = [ NaN ; rho(2:end) ] ;
        
        % add probability-weighted state variables to ensemble mean
        u_mean = u_mean + P(samplestep,j)*u_tmp ;
        v_mean = v_mean + P(samplestep,j)*v_tmp ;
        rho_mean = rho_mean + P(samplestep,j)*rho_tmp ;
        
        if ( render )
            
            % render flow field with rho values and velocity streamlines
            subplot(4,3,3+j)
            titlestr = modelnames{j} ;
            Pete_plotfield
            if ( j == truereac )
                trueax = gca ;
            end
            
            set(gca,'Position',get(gca,'Position')+figshift)
            
        end
        
    end
    
    % ENSEMBLE MEAN -------------------------------------------------------
    
    % load true state geometry masks
    load( [ savemain, '/true/param.mat' ], ...
            'Nodeu', 'Nodev', 'NodeRho' )
    
    if ( render )
        
        % render flow field with rho values and velocity streamlines
        u_tmp = u_mean ; v_tmp = v_mean ; rho_tmp = rho_mean ;
        subplot(4,3,2)
        titlestr = 'ensemble mean' ;
        Pete_plotfield
        
        set(gca,'Position',get(gca,'Position')+figshift)
        
    end
    
    % RMSE ----------------------------------------------------------------
    
    % calculate RMSE
    load( [ savemain, '/true/', steps(i).name ], ...
        'u_true', 'v_true', 'rho_true' )
    all_dif = ([u_mean;v_mean;rho_mean]-[u_true;v_true;rho_true]).^2 ;
    RMSE(i) = sqrt( nansum(all_dif) / sum(~isnan(all_dif)) ) ;
    
    % MODEL DISTRIBUTION --------------------------------------------------
    
    % store ensemble distribution
    P_t(i,:) = P(samplestep,:) ;
    
    if ( render )
        
        % plot instantaneous reac probabilities
        subplot(4,3,3)
        
        % exponentially scale reac probabilities such that uniform prior
        % corresponds to mid-point of colorscale
        P_scaled = P(samplestep,:) .^ (log(1/2)/log(1/9)) ;
        pcolor( (0:3), flipud((0:3)'), ...
            padarray(reshape(P_scaled,3,3)',[1,1],0,'post') )
        cmap = colormap(hsv) ; colormap(cmap(1:86,:))
        shading faceted, caxis([0,1]), freezeColors
        
        % set subplot properties
        title( 'reac distribution', 'FontWeight', 'bold' )
        box on, xlim([0,3]), ylim([0,3])
        format_ticks( gca, ...
            {'d_{-}(r)','d_{0}','d_{+}(r)'}, ...
            {'g_{-}(t)','g_{0}','g_{+}(t)'}, ...
            (0.5:1:2.5), (0.5:1:2.5), 0, 0, 0.04, ...
            'FontWeight', 'bold' ) ;
        
        set(gca,'Position',get(gca,'Position')+figshift)
        
        % label grid
        pos = get( gca, 'Position' ) ;
        gridc = pos(1,[1,2]) + pos(1,[3,4])/2 ;
        scale = 0.85 ;
        pos = [ gridc - scale*pos(1,[3,4])/2, scale*pos(1,[3,4]) ] ;
        gridd = pos(1,[3,4])/3 ;
        set( gca, 'Position', pos )
        for j = 1:9
            shift = [ mod(j-1,3)-1, 2-ceil(j/3) ] ;
            annotation( 'textbox', [gridc+gridd.*shift,eps,eps], ...
                'String', sprintf('%2.1f%%', 100*P(samplestep,j)), ...
                'FontSize', 10, 'FontWeight', 'bold', ...
                'HorizontalAlignment', 'center', ...
                'VerticalAlignment', 'middle', 'FitBoxToText', 'on', ...
                'EdgeColor', 'none' )
        end
        
    end
    
    % ANNOTATIONS ---------------------------------------------------------
    
    if ( render )
        
        axes('Position',[0.42,0.04,0.2,0.02]+figshift,'Visible','off')
        colormap(pmkmp(256,'CubicL'))
        shading interp, caxis([0.0,1.0]), freezeColors
        h = colorbar('Location','SouthOutside') ;
        set(h,'Position',[0.42,0.04,0.2,0.02]+figshift)
        set(get(h,'Title'),'String','concentration')
        
        % delineate ensemble plots from candidate reac plots
        annotation('line',[0.10,0.93]+figshift(1),...
            [0.74,0.74]+figshift(2),'LineWidth',1)
        
        % highlight true reac amongst candidate modelnames
        if ( exist('trueax','var') )
            annotation('rectangle',...
                get(trueax,'Position')+[-0.015,-0.025,0.025,0.055],...
                'LineStyle','--')
        end
        
    end
    
    % ---------------------------------------------------------------------
    
    % save ensemble mean data
    savefile = [ 'mean/', steps(i).name ] ;
    save( [savemain,'/',savefile], 'u_mean', 'v_mean', 'rho_mean' )
    fprintf( 'Saved %s.\n', savefile )
    
    if ( render )
        
        % save raster image (png)
        savefile = sprintf('%d_%02d',truereac,times(i)) ;
        savefig([savemain,'/',savefile],1,'png')
        fprintf( 'Saved %s.png.\n', savefile )
        
        close(1)
        
    end
    
    fprintf('%d of %d steps done. %0.1f seconds elapsed.\n',...
        i,n_steps,toc(app.toptic))
    
end

% save metrics
savefile = 'stats.mat' ;
save( [savemain,'/',savefile], 'times', 'P_t', 'RMSE' )
fprintf( 'Saved %s.\n', savefile )

if ( render )

    figure(2), clf
    set(2,'Position',[10,50,500,750],'PaperPositionMode','auto')

    % plot RMSE
    subplot(2,1,1)
    plot(times,RMSE,'b.')
    box on, xlim([0,T]), ylim([0,0.12])
    set( gca, 'XTick', (0:20:T), 'YTick', (0:0.03:0.12), 'YGrid', 'on' )
    xlabel('time','FontWeight','bold')
    ylabel('RMSE','FontWeight','bold')

    % plot true reac probability
    if ( 1 <= truereac && truereac <= 9 )
        subplot(2,1,2)
        plot(times,P_t(:,truereac),'b.')
        box on, xlim([0,T]), ylim([0,1])
        set( gca, 'XTick', (0:20:T), 'YTick', (0:0.2:1), 'YGrid', 'on' )
        xlabel('time','FontWeight','bold')
        ylabel('true reac probability','FontWeight','bold')
    end

    % % plot other reac probabilities
    % subplot(3,1,3)
    % specs = { '.r', 'or', '+r', ...
    %           '.g', 'og', '+g', ...
    %           '.b', 'ob', '+b' } ;
    % for i = 1:9
    %     if ( i ~= truereac )
    %         plot(times,P_t(:,i),specs{i})
    %         hold on
    %     end
    % end
    % box on, xlim([0,T]), ylim([0,0.25])
    % set( gca, 'XTick', (0:20:T), 'YTick', (0:0.05:0.25), 'YGrid', 'on' )
    % xlabel('time','FontWeight','bold')
    % ylabel('reac probability','FontWeight','bold')
    
    % save raster image (png)
    savefile = sprintf('%d_stats',truereac) ;
    savefig([savemain,'/',savefile],2,'png')
    fprintf( 'Saved %s.png.\n', savefile )
    
    close(2)
    
end

fprintf( 'Phew! Total time = %02d:%02d.\n', ...
    floor(toc(app.toptic)/3600), floor(mod(toc(app.toptic),3600)/60) )

