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

clear;

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

% Compile data traces.
disp('Loading data...');
data = {};

% ecRNH demo data
%filename = '2011-10-27 ecRNH demo jump sets'; % filename
%fibername = 'MecRNH_517A_traj_01_multiTraj_01_saved'; % dataset to extract
filename = '2011-11-10 ecRNH 500s splitchannels'; % filename (split channels)
fibername = 'MecRNH_509A_chanB_traj_01_traj_06_multiTraj_01_saved'; % dataset to extract
subsample = 10; % will keep every 'subsample' data point
sampletime = (1/1000) * subsample; % time between samples, in seconds
drop_initial_part_ms = 0; % number of initial ms to drop from each trace

filename = sprintf('data/%s.mat', filename);
load(filename, fibername); % load specified dataset
original_data = eval(fibername);
nalign = 100; % number of final samples to align on (before subsampling)
ntraces = length(original_data);
for trace_index = 1:ntraces
  % Extract observable.
  time_t = original_data{trace_index}(:,1)';
  trap_t = original_data{trace_index}(:,2)';  
  force_t = original_data{trace_index}(:,3)';  
  
  % Discard initial samples.
  indices = find(time_t > (time_t(1) + drop_initial_part_ms*0.001));
  o_t = trap_t(indices);

  % Align last part of data (quenched to folded state).
  mu = mean(o_t(end-nalign+1:end));
  o_t = o_t - mu;
  clear mu;

  % Subsample data.
  o_t = o_t(1:subsample:end);
  
  % Store data.
  data{trace_index} = o_t;
end

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 = false; % infer reversible transition matrices
options.diagonally_dominant = true; % 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

% Generate an initial (poor) model from data.
%initial_model = generate_initial_model(data, nstates, options);

% 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);
  
% 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);

