% Bottom Gravity Current Bayesian Model Learning

% render        1 = generate figures,
%               0 = do not generate figures

% Modified from Pete's biology reaction term learning case (06/09/2012) by Jing
% Last modified: Aug 2014
%------------------------------------------------------------------------
trunk = regexp(pwd,'.*trunk','match','once') ;
true_dir = [trunk,'/Save/mat/BGC_Det_True1'];
plotind = 30:39;

gmmdo_dir = [trunk,'/Save/mat'];
ML_dir = [gmmdo_dir, '/ML_1'];
if (~exist(ML_dir,'dir')), mkdir(ML_dir), end;
run_dir = 'BGC_DO_GMM_';
run_ind = {'101','102','103','104','105','107','108','112'};
True_Model = 2;
Num_Model = length(run_ind);
offset = 0;

opengl neverselect
addpath([trunk,'/Cases/Bottom_Gravity_Current/freezeColors']);
addpath([trunk,'/Cases/Bottom_Gravity_Current/pmkmp']);

ymin = 0.6;
render = true;
have_probs = true;

%------------------------------------------------------------------------

modelnames = {...
    '\DeltaS=0.5, S_{init}=exp(2)', ...
    '\DeltaS=1.5, S_{init}=exp(8)', ...
    '\DeltaS=1.5, S_{init}=cosine', ...
    '\DeltaS=1.5, S_{init}=linear', ...
    '\DeltaS=3.0, S_{init}=exp(2)', ...
    '\DeltaS=1.5, S_{init}=cos+exp(5)', ...
    '\DeltaS=1.5, S_{init}=linear+exp(5)', ...
    '\DeltaS=1.5, S_{init}=exp(8), 0 Vert Diffusivity', ...
    };

% Load log-likelihood data
Q = [];
for i = 1 : Num_Model
    load(sprintf('%s/%s%s/%03d_pos.mat', gmmdo_dir, run_dir, run_ind{i},...
        plotind(end)-offset), 'loglikes');
    Q = [Q, loglikes(:,2)];
end;

% Calculate distributions over modelnames by Bayes formula
P = [(1/Num_Model)*ones(1,Num_Model); zeros(size(Q))];
for i = 1 : size(Q,1)
    for j = 1 : Num_Model
        den = 0;
        for k = 1:Num_Model
            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;

% find all save steps
n_steps = length(plotind);

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

tan_theta = tan(app.theta*pi/180);
app.ObsLocs = [...       
    2000, (A1-2000)*tan_theta + B1 + 25;
    2000, (A1-2000)*tan_theta + B1 + 58;
    2000, (A1-2000)*tan_theta + B1 + 91;
    2000, (A1-2000)*tan_theta + B1 + 124;
    2000, (A1-2000)*tan_theta + B1 + 157;
    2000, (A1-2000)*tan_theta + B1 + 190;
    3000, (A1-3000)*tan_theta + B1 + 25;
    3000, (A1-3000)*tan_theta + B1 + 65;
    3000, (A1-3000)*tan_theta + B1 + 105;
    3000, (A1-3000)*tan_theta + B1 + 145;
    3000, (A1-3000)*tan_theta + B1 + 185;
    3000, (A1-3000)*tan_theta + B1 + 225;
    4000, (A1-4000)*tan_theta + B1 + 30;
    4000, (A1-4000)*tan_theta + B1 + 80;
    4000, (A1-4000)*tan_theta + B1 + 130;
    4000, (A1-4000)*tan_theta + B1 + 180;
    4000, (A1-4000)*tan_theta + B1 + 230;
    4000, (A1-4000)*tan_theta + B1 + 280]/1000;

% Set up the patch for the bathymetry
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 = plotind;
P_t = zeros(n_steps, Num_Model);
RMSE = zeros(n_steps, 3);
Model_Ent = 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 = u_true; u_tmp(1:Nbcs) = NaN;
    v_tmp = v_true; v_tmp(1:Nbcs) = NaN;
    rho_tmp = rho_true; rho_tmp(1:Nbcs) = NaN;
    
    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.ObsLocs(:,1),app.ObsLocs(:,2),20,[1.0,1.0,1.0],'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 : Num_Model
        
        % Compute entropy of distribution over different models
        if (P(i,j) > eps)
            Model_Ent(i) = Model_Ent(i) - P(i,j)*log(P(i,j));
        end;
        
        % load model state variables and parameters
        load(fullfile(gmmdo_dir,[run_dir,run_ind{j}],...
            num2str(plotind(i)-offset,'%3.3d')), 'u', 'v', 'rho', 'Yi');
        
        % set obstacle velocity and rho values to NaN
        u_tmp = u; u_tmp(1:Nbcs) = NaN;
        v_tmp = v; v_tmp(1:Nbcs) = NaN;
        rho_tmp = rho; rho_tmp(1:Nbcs) = NaN;
        
        % 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{j};
            Pete_plotfield2;
            if (~isempty(True_Model))&&(j == True_Model)
                trueax = gca;
            end;
            
            set(gca,'Position',get(gca,'Position')+figshift);
            
        end;
        
    end;
    
    % ENSEMBLE MEAN -------------------------------------------------------
    
    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 = 'Bayesian 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)));
        RMSE(i,1) = sqrt(sum(u_mean(Nbcs+1:end)-u_true(Nbcs+1:end)).^2/(length(u_true)-Nbcs));
        RMSE(i,2) = sqrt(sum(v_mean(Nbcs+1:end)-v_true(Nbcs+1:end)).^2/(length(v_true)-Nbcs));
        RMSE(i,3) = sqrt(sum(rho_mean(Nbcs+1:end)-rho_true(Nbcs+1:end)).^2/(length(rho_true)-Nbcs));
        
        % 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 = zeros(1,9);
            P_scaled(1:Num_Model) = P(i,:).^(log(1/2)/log(1/9));
            P_scaled(Num_Model+1:end) = NaN;
            
            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('Model 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 : Num_Model
                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
        Handle = colorbar('Location','SouthOutside') ;
        set(Handle,'Position',[0.42,0.04,0.2,0.02]+figshift)
        set(get(Handle,'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('%03d.mat', plotind(i));
            save([ML_dir,'/',savefile], 'u_mean', 'v_mean', 'rho_mean');
            fprintf('Saved %s.\n', savefile);
            if render
                % save raster image (png)
                savefile = sprintf('%3.3d',plotind(i));
                savefig([ML_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([ML_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(3,1,1)
    plot(times,RMSE(:,1)/RMSE(1,1),'b', times,RMSE(:,2)/RMSE(1,2),'g',...
        times,RMSE(:,3)/RMSE(1,3),'r');
    grid on, grid minor;
    legend('u velocity', 'v velocity', 'Salinity');
    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 (~isempty(True_Model))
        subplot(3,1,2);
        plot(times,P_t(:,True_Model),'b');
        grid on, grid minor;
        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 model probability','FontWeight','bold');
    end;
    
    % Plot entropy of distribution over different models
    subplot(3,1,3);
    plot(times,Model_Ent,'b');
    grid on, grid minor;
    box on, ylim([1,2.5]);
    %set( gca, 'XTick', (0:20:T), 'YTick', (0:0.2:1), 'YGrid', 'on' )
    xlabel('time','FontWeight','bold');
    ylabel('Entropy of dist over models','FontWeight','bold');
    
    % % 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('stats');
    savefig([ML_dir,'/',savefile],2,'png');
    fprintf('Saved %s.png.\n', savefile);
    close(2);    
end;

toc;
