% 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

%------------------------------------------------------------------------
%trunk = regexp(pwd,'.*trunk','match','once') ;
trunk = '/home/phaley/Students/Jing/FV_Matlab_Framework_JingPJH/trunk';
true_dir = [trunk,'/Save_Class2/mat/Bottom_Gravity_Current_Det_006'];
offset = 30;
%plotind = 30:39;
%gmmdo_dir = [trunk,'/Save_GMM04_M001/mat'];
plotind = 30:40;
gmmdo_dir = [trunk,'/Save_GMM04_M005/mat'];
run_dir = 'BGC_';
run_ind = 1:9;
render = 1;
opengl neverselect
addpath ([trunk,'/Cases/Pete_tracer/freezeColors']);
addpath ([trunk,'/Cases/Pete_tracer']);
ymin = 0.6;
have_probs = 1;
truereac = find (run_ind==6);
%------------------------------------------------------------------------

%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 = { '\DeltaS=1.5, S_{init}=exp(2); q_{BC}=D_{inlet}', ...
               '\DeltaS=1.5, S_{init}=exp(8); q_{BC}=D_{inlet}', ...
               '\DeltaS=3.0, S_{init}=cosine; q_{BC}=D_{inlet}', ...
               '\DeltaS=3.0, S_{init}=linear; q_{BC}=D_{inlet}', ...
               '\DeltaS=3.0, S_{init}=exp(2); q_{BC}=D_{inlet}', ...
               '\DeltaS=3.0, S_{init}=exp(8); q_{BC}=D_{inlet}', ...
               '\DeltaS=3.0, S_{init}=cos+exp(5); q_{BC}=D_{inlet}', ...
               '\DeltaS=3.0, S_{init}=linear+exp(5); q_{BC}=D_{inlet}', ...
               '\DeltaS=3.0, S_{init}=cos+linear; q_{BC}=D_{inlet}', ...
               '\DeltaS=3.0, S_{init}=exp(8); q_{BC}=Orig', ...
               '\DeltaS=3.0, S_{init}=cos+exp(5); q_{BC}=Orig', ...
 } ;
%      gmm_m001_001 - dS=1.5; Sinit=exp(2);         qBC=D-inlet
%      gmm_m001_002 - dS=1.5; Sinit=exp(8);         qBC=D-inlet
%      gmm_m001_003 - dS=3.0; Sinit=cos/sin;        qBC=D-inlet
%      gmm_m001_004 - dS=3.0; Sinit=linear;         qBC=D-inlet
%      gmm_m001_005 - dS=3.0; Sinit=exp(2);         qBC=D-inlet
%      gmm_m001_006 - dS=3.0; Sinit=exp(8);         qBC=D-inlet
%      gmm_m001_006 - dS=3.0; Sinit=cos/sin+exp(5); qBC=D-inlet
%      gmm_m001_008 - dS=3.0; Sinit=linear+exp(5);  qBC=D-inlet
%      gmm_m001_009 - dS=3.0; Sinit=cos/sin+linear; qBC=D-inlet
%      gmm_m001_010 - dS=3.0; Sinit=exp(8);         qBC=original
%      gmm_m001_011 - dS=3.0; Sinit=cos/sin+exp(5); qBC=original

Q = [] ;
for i = 1:9
    load(sprintf('%s/BGC_%03d/007_pos.mat',gmmdo_dir,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 = [gmmdo_dir, '/mean' ] ;
if ( ~exist(savedir,'dir') )
   mkdir( savedir )
end

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

% load parameters common to all modelnames
load( fullfile(true_dir,'setup'), 'a', 'b', 'app', ...
    'dx', 'dy', 'XP', 'YP', 'XU', 'YU', 'XV', 'YV', ...
    'Nbcs', 'Nodeu', 'Nodev', 'NodeP', 'h', 'DelS', 'nu_x')
Nx=app.Nx;
Ny=app.Ny;
dt=app.dt;
NodeRho=NodeP;
XP=XP.*(h/1000);
YP=YP.*(h/1000);
XU=XU.*(h/1000);
YU=YU.*(h/1000);
XV=XV.*(h/1000);
YV=YV.*(h/1000);
XRho=XP;
YRho=YP;
a = a*(h/1000);
b = b*(h/1000);

ridgex = [0 33.2667 33.2667 7.414 0].*(h/1000);
ridgey = [0 0 4.25 5.37 5.37].*(h/1000);
ridgeColor = [105 105 105]/255;

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

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

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( fullfile (true_dir,num2str(plotind(i),'%3.3d')), 'u', 'v', 'rho' )
    
    % set obstacle velocity and rho values to NaN
    u_true = u ; v_true = v ; rho_true = rho ;
    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_plotfield2

        % 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( fullfile (gmmdo_dir,[run_dir,num2str(run_ind(j),'%3.3d')],num2str(plotind(i)-offset,'%3.3d')), 'u', 'v', 'rho', 'Yi' )
%        load( sprintf( '%s/%d/%s', savemain, j, steps(i).name ), ...
%            'u', 'v', 'rho', 'YYt', 'i' )
%        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
        if (have_probs)
           u_mean = u_mean + P(i,j)*u_tmp ;
           v_mean = v_mean + P(i,j)*v_tmp ;
           rho_mean = rho_mean + P(i,j)*rho_tmp ;
        end;
        
        if ( render )
            
            % render flow field with rho values and velocity streamlines
            subplot(4,3,3+j)
            titlestr = modelnames{run_ind(j)} ;
            Pete_plotfield2
            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 & have_probs )
        
        % 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_plotfield2
   
        set(gca,'Position',get(gca,'Position')+figshift)
        
    end
    
    if ( have_probs )
       % RMSE ----------------------------------------------------------------
    
       % calculate RMSE
       load( sprintf('%s/%03d.mat', true_dir, plotind(i)), ...
           'u', 'v', 'rho' );
       u_true = u;
       v_true = v;
       rho_true = rho;
       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(i,:) ;
    
       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(i,:) .^ (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( 'IC distribution', 'FontWeight', 'bold' )
           box on, xlim([0,3]), ylim([0,3])
           format_ticks( gca, ...
               {' ',' ',' '}, ...
               {' ',' ',' '}, ...
               (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(i,j)), ...
                   'FontSize', 10, 'FontWeight', 'bold', ...
                   'HorizontalAlignment', 'center', ...
                   'VerticalAlignment', 'middle', 'FitBoxToText', 'on', ...
                   'EdgeColor', 'none' )
           end
        
       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','S/\DeltaS')
        
        % 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
        
       % ---------------------------------------------------------------------
    
       if ( have_probs )
         % save ensemble mean data
         savefile = sprintf('mean/%03d.mat', i);
         save( [gmmdo_dir,'/',savefile], 'u_mean', 'v_mean', 'rho_mean' )
         fprintf( 'Saved %s.\n', savefile )
    
         if ( render )
        
             % save raster image (png)
             savefile = sprintf('%d_%3.3d',truereac,plotind(i)) ;
             savefig([gmmdo_dir,'/',savefile],1,'png')
             fprintf( 'Saved %s.png.\n', savefile )
        
             close(1)
        
         end
       end
    end
    
    fprintf('%d of %d steps done. %0.1f seconds elapsed.\n',...
        i,n_steps,toc)
    
end
return

% save metrics
savefile = 'stats.mat' ;
save( [gmmdo_dir,'/',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([gmmdo_dir,'/',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) )

