% Analyze traces with Bayesian HMM to compute average forces and lifetimes for different trap extensions.

%%


clear;
close all;

addpath ../../matlab/bhmm
javaaddpath ../../matlab/bhmm


%%
%disp('    ...select DataPath');
%DataPath = '/Users/christian/bhmm/kaiser-lysozyme/data';
%DataPath = uigetdir('/Users/christian');
%disp(['       ' DataPath])

%disp('    ...select AnalysisPath');
%AnalysisPath = uigetdir('/Users/christian/bhmm/kaiser-lysozyme/analysis');
%disp(['       ' AnalysisPath])
        

disp('    ...select Files:');
%[files, DataPath] = uigetfile('/Users/christian/bhmm/kaiser-lysozyme/*.mat', 'MultiSelect', 'on');

% DEBUG
files = {'M100824_09_4.3.mat'};
files = {'100926_annotated/M100816_04_4.4.mat'};
DataPath = 'data';

if ~iscell(files)
    files = {files};
end
disp(files');


%%
for file_index = 1:numel(files)
    %%
    close all;
    load([DataPath '/' files{file_index}]);
    disp(files{file_index});

    fibername = files{file_index};
    subsample = 10;
    sampletime= (1./data.fS) * subsample;
%    filename  = [DataPath '/' files{file_index}];
    dataInfo = rmfield(data, 'x'); % some 'meta data'
    
    xx        = data.x;
    
    %%
    data = {};
    ntraces = size(xx,1);
    for trace_index = 1:ntraces
      % Extract x
      o_t = xx{trace_index}';

      % Subsample data.
      o_t = o_t(1:subsample:end);

      % Discard first sample (which is usually at much higher force).
      o_t = o_t(2:end);

      % Store data.
      data{trace_index} = o_t;
    end
    clear 'xx';
    %%
    nstates = 3; % number of states -- fixed beforehand
    nsamples = 1000; % number of models to sample with Bayesian HMM
    nburnin = 50; % number of samples to discard to equilibration

    % Allocate storage for estimates
    estimates = cell(ntraces,1); % estimates{trace_index} is a struct containing estimates of state properties for each trace

    options = bhmm_default_options();

    % Options for transition matrix estimation.
    options.reversible = true; % infer reversible transition matrices
    options.diagonally_dominant = false; % enforce diagonally-dominant transition matrices

    options.verbosity = 2; % set verbosity level
    options.equilibrium = false; % trajectory data is not initially drawn from equilibrium
    options.tau = sampletime; % seconds between observations
    options.observable_units = 'nm'; % units for observable axis
    options.time_units = 's'; % units for time axis

    options.enforce_state_ordering = true; % enforce states are ordered in increasing order
    options.enforce_state_separation = 2.0; % if nonzero, will enforce state means to be separated by specified amount

    % Tell BHMM that we are starting from either states 1 or 2, and not state 3.
    ZERO = 1.0e-8;
    options.log_rho_i = log([0.5 0.5 ZERO]);

    % Set prior transition pseudocounts to inform BHMM that native state should be very metastable.
    options.prior_transition_pseudocounts = ones([nstates, nstates]);
    options.prior_transition_pseudocounts(1,1) = 1000; % self-transition for native state has equivalent to this number of observations
    
% DEBUG: Instead of generating an MLHMM model, we specify an initial guess below.
%    % Determine maximum-likelihood HMM.
%    disp('Determining maximum-likelihood model...');
%    mlmodel = mlhmm(data, nstates, options);
%    plot_state_assignments(data, mlmodel, options);
%    filename = sprintf('plots/%s-mle-states.png', fibername);
%    print('-r300', '-dpng', filename);

    % Use hard-coded model.
    model.nstates = 3;
    model.Tij = [
	0.7682    0.2313    0.0005
	0.2167    0.7824    0.0009
	0.0182    0.0326    0.9492];
    model.Pi = stationary_probability(model.Tij); % maybe optional
    model.logPi = log(model.Pi); % maybe optional    
    model.states = cell(3,1);
    model.states{1} = struct('mu', 0.0, 'sigma', 2.0);
    model.states{2} = struct('mu', 10.0, 'sigma', 5.0);
    model.states{3} = struct('mu', 25.0, 'sigma', 5.0);
    mlmodel = model;   

    % Initialize Bayesian HMM.
    disp('Sampling from Bayesian posterior...');
    models = bhmm(data, mlmodel, nsamples, options);

    % Save models.
    filename = sprintf('analysis/%s-models.mat', fibername);
    save(filename);

    %% Analayze Bayesian models.

    % Plot state assignments.
    %plot_fractional_state_assignments(data, models(nburning+1:end), options);
    plot_state_assignments(data, models(end), options);
    filename = sprintf('plots/%s-stateassignments.png', fibername);
    print('-r300', '-dpng', filename);

    % Plot states.
    options.sameaxis = false;
    analyze_states(models(nburnin+1:end), options);
    filename = sprintf('plots/%s-states.png', fibername);
    print('-r300', '-dpng', filename);
    %print('-depsc', filename);
    %command = sprintf('epstopdf %s', filename);
    %system(command);

    % Plot transition matrix elements.
    options.sameaxis = false;
    analyze_transition_matrix(models(nburnin+1:end), options);
    filename = sprintf('plots/%s-transition-matrix.png', fibername);
    print('-r300', '-dpng', filename);
    %print('-depsc', filename);
    %command = sprintf('epstopdf %s', filename);
    %system(command);

    % Plot rate constants.
    options.sameaxis = false;
    analyze_rates(models(nburnin+1:end), options);
    filename = sprintf('plots/%s-rates.png', fibername);
    print('-r300', '-dpng', filename);
    %print('-depsc', filename);
    %command = sprintf('epstopdf %s', filename);
    %system(command);
    disp(['Done processing ' files{file_index}]);
end