% Last modified: 06/09/2012 by Pete

% trueshape     true shape of obstacle in system [1,5]
%                 1 = circle,
%                 2 = square,
%                 3 = upstream-pointing triangle,
%                 4 = downstream-pointing triangle,
%                 5 = diamond
% Re            Reynolds number [50,100]
% 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( trueshape, Re, randseed, render )

app.toptic = tic ;
fprintf(['Stochastic flow past an obstacle: ',...
    'model ensemble.\n'])
close all

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

shapenames = { 'circle', ...
               'square', ...
               'upstream-pointing triangle', ...
               'downstream-pointing triangle', ...
               'diamond' } ;

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

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

for i = 1:size(Q,1)
    
    for j = 1:5
        den = 0 ;
        for k = 1:5
            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 models
load( [ savemain, '/true/param.mat' ], 'a', 'b', 'Nx', 'Ny', ...
    'dx', 'dy', 'XP', 'YP', 'XU', 'YU', 'XV', 'YV', 'dt', ...
    'SaveIntrvl', 'Nbcs' )

times = zeros(n_steps,1) ;
P_t = zeros(n_steps,5) ;
PRe = zeros(n_steps,1) ;
RMSE = zeros(n_steps,1) ;

for i = 1:n_steps
    
    times(i) = str2double(steps(i).name(1:5))/100 ;
    
    if ( render )
    
        figure(1), clf
        set(1,'Position',[10,50,1000,1060],'PaperPositionMode','auto')
        warning('off',...
            'MATLAB:hg:patch:ColorDataNotSetForInterpolatedShading')
        
    end
    
    % TRUE STATE ----------------------------------------------------------
    
    % load true flow field state variables and parameters
    load( [ savemain, '/true/', steps(i).name ], 'u_true', 'v_true' )
    load( [ savemain, '/true/param.mat' ], 'Nodeu', 'Nodev' )
    
    % set obstacle velocity values to NaN
    u_true = u_true ; v_true = v_true ;
    u_tmp = [ NaN ; u_true(2:end) ] ;
    v_tmp = [ NaN ; v_true(2:end) ] ;
    
    if ( render )
        
        % render flow field with velocity magnitudes and streamlines
        subplot(5,2,1)
        titlestr = 'true state' ;
        Pete_plotfield
        
        % plot sampling locations
        hold on
        scatter( app.SampleLocs(:,1), app.SampleLocs(:,2), ...
            25, [0.6,0.8,1.0], 'filled')
        hold off
        
    end
    
    % MODEL MEANS ---------------------------------------------------------
    
    % initialize model ensemble mean velocity grids
    ubox_mean = zeros(size(Nodeu)) ;
    vbox_mean = zeros(size(Nodev)) ;
    
    % initialize Re distributions
    Re_x = linspace(app.Re_range(1),app.Re_range(2),1000)' ;
    Re_pdf = zeros(size(Re_x)) ; Re_cdf = zeros(size(Re_x)) ;
    
    % plot model flow fields
    for j = 1:5
        
        % load model state variables and parameters
        load( sprintf( '%s/%d/%s', savemain, j, ...
            steps(i).name ), 'u', 'v', 'ui', 'vi', 'YYt', ...
            'samplestep' )
        load( sprintf( '%s/%d/param.mat', savemain, j ), ...
            'Nodeu', 'Nodev' )
        
        % set obstacle velocity values to NaN
        u = u ; v = v ;
        u_tmp = [ NaN ; u(2:end) ] ;
        v_tmp = [ NaN ; v(2:end) ] ;
        
        % add probability-weighted state variables to ensemble mean
        ubox_mean = ubox_mean + P(samplestep,j) * u(Nodeu) ;
        vbox_mean = vbox_mean + P(samplestep,j) * v(Nodev) ;
        
        % add weighted kernel density to Re ensemble distributions
        Re_pdf = Re_pdf + P(samplestep,j) ...
            * ksdensity( Re * ( u(2,1) + ui(2,1)*YYt(:,1) ), Re_x, ...
            'function', 'pdf' ) ;
        Re_cdf = Re_cdf + P(samplestep,j) ...
            * ksdensity( Re * ( u(2,1) + ui(2,1)*YYt(:,1) ), Re_x, ...
            'function', 'cdf' ) ;
        
        if ( render )
            
        	% render flow field with velocity magnitudes and streamlines
            subplot(5,2,2*j)
            titlestr = shapenames{j} ;
            Pete_plotfield
            if ( j == trueshape )
                trueax = gca ;
            end
            
        end
        
    end
    
    % ENSEMBLE MEAN -------------------------------------------------------
    
    [~,ind] = max(P(samplestep,:)) ;
    load( sprintf( '%s/%d/param.mat', savemain, ind ), 'Nodeu', 'Nodev' )
    u_tmp = zeros(max(Nodeu(:)),1) ; v_tmp = zeros(max(Nodev(:)),1) ;
    for j = 1:length(Nodeu(:))
        u_tmp(Nodeu(j)) = ubox_mean(j) ;
    end
    for j = 1:length(Nodev(:))
        v_tmp(Nodev(j)) = vbox_mean(j) ;
    end
    u_tmp(1) = NaN ; v_tmp(1) = NaN ;
    
    if ( render )
        
        % render flow field with velocity magnitudes and streamlines
        subplot(5,2,3)
        titlestr = 'ensemble mean' ;
        Pete_plotfield
        
    end
    
    % RMSE ----------------------------------------------------------------
    
    % load true flow field state variables and parameters
    load( [ savemain, '/true/', steps(i).name ], 'u_true', 'v_true' )
    load( [ savemain, '/true/param.mat' ], 'Nodeu', 'Nodev' )
    
    % set obstacle velocity values to NaN
    u_true(1) = NaN ; v_true(1) = NaN ;
    
    % calculate RMSE
    uerr = ubox_mean - u_true(Nodeu) ;
    verr = vbox_mean - v_true(Nodev) ;
    err = [uerr(:);verr(:)] ;
    RMSE(i) = sqrt( nansum(err.^2) / sum(~isnan(err)) ) ;
    
    % ENSEMBLE RE DISTRIBUTION --------------------------------------------
    
    % calculate integral of Re pdf between -2.5 and +2.5 from true Re
    PRe(i) = interp1(Re_x,Re_cdf,Re+2.5*[-1,1])*[-1;1] ;
    
    if ( render )
        
        % plot ensemble Re distribution
        subplot(5,2,5)
        plot( Re_x, Re_pdf, 'LineWidth', 1 )
        axis([app.Re_range.*[0.99,1.01],0,0.6])
        set( gca, 'XTick', (app.Re_range(1):10:app.Re_range(2)), ...
            'YTick', (0:0.2:0.6), 'YGrid', 'on', ...
            'Position', get(gca,'Position') ...
            + [1,0,-2,0]*0.04 + [0,1,0,-2]*0.005 )
        title('ensemble Re distribution','FontWeight','bold')
        
    end
    
    % MODEL DISTRIBUTION --------------------------------------------------
    
    % store model distribution
    P_t(i,:) = P(samplestep,:) ;
    
    if ( render )
        
        % plot model distribution
        subplot(5,2,7)
        bar( (1:5), P(samplestep,:), 'BarWidth', 0.5, ...
            'FaceColor', [0.4,1,0], 'EdgeColor', [0,0,0] ) ;
        h = gca ;
        set( h, 'XLim', [0,6], 'YLim', [0,1], ...
            'XTick', [], 'YTick', (0:0.2:1), 'YGrid', 'on', ...
            'Position', get(h,'Position') ...
            + [1,0,-2,0]*0.04 + [0,1,0,-2]*0.005 )
        title( 'model distribution', 'FontWeight', 'bold')
        hh = axes ;
        plot(hh,1,0,'ok',2,0,'sk',3,0,'<k',4,0,'>k',5,0,'dk',...
            'MarkerSize', 12)
        set( hh, 'XLim', [0,6], 'YLim', [0,1], ...
            'Position', get(h,'Position') - [0,0.015,0,0], ...
            'Visible', 'off')
        
    end
    
    % ANNOTATIONS ---------------------------------------------------------
    
    if ( render )
        
        axes('Position',[0.2,0.16,0.2,0.02],'Visible','off')
        colormap(hot), caxis([0,2])
        h = colorbar('Location','SouthOutside') ;
        set(h,'Position',[0.2,0.16,0.2,0.02])
        set(get(h,'Title'),'String','velocity magnitude')
        
        % delineate ensemble plots from candidate model plots
        annotation('line',[0.51,0.51],[0.08,0.96],'LineWidth',1)
        
        % highlight true model amongst candidate models
        if ( exist('trueax','var') )
            annotation('rectangle',...
                get(trueax,'Position')+[-0.025,-0.025,0.045,0.05],...
                'LineStyle','--')
        end
        
    end
    
    % ---------------------------------------------------------------------
    
    % save ensemble mean data
    savefile = [ 'mean/', steps(i).name ] ;
    save( [savemain,'/',savefile], 'ubox_mean', 'vbox_mean' )
    fprintf( 'Saved %s.\n', savefile )
    
    if ( render )
        
        % save raster image (png)
        savefile = sprintf('%d_%02d',trueshape,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', 'PRe', 'RMSE' )
fprintf( 'Saved %s.\n', savefile )

if ( render )

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

    % plot RMSE
    subplot(3,1,1)
    plot(times,RMSE,'b.','LineWidth',1.5)
    box on, xlim([0,T]), ylim([0,0.2])
    set( gca, 'XTick', (0:10:T), 'YTick', (0:0.05:0.2), 'YGrid', 'on' )
    xlabel('time','FontWeight','bold')
    ylabel('RMSE','FontWeight','bold')

    % plot true Re probability
    subplot(3,1,2)
    plot(times,PRe,'b.','LineWidth',1.5)
    box on, xlim([0,T]), ylim([0,1])
    set( gca, 'XTick', (0:10:T), 'YTick', (0:0.2:1), 'YGrid', 'on' )
    xlabel('time','FontWeight','bold')
    ylabel('true Re probability','FontWeight','bold')

    % plot true model probability
    subplot(3,1,3)
    plot(times,P_t(:,trueshape),'b.')
    box on, xlim([0,T]), ylim([0,1])
    set( gca, 'XTick', (0:10:T), 'YTick', (0:0.2:1), 'YGrid', 'on' )
    xlabel('time','FontWeight','bold')
    ylabel('true model probability','FontWeight','bold')
    
    % save raster image (png)
    savefile = sprintf('%d_stats',trueshape) ;
    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) )

