% This example covers sampling from the infinite Gaussian mixture model
%
% Copyright 2014, Author: Frank Wood, fwood@robots.ox.ac.uk

%% add gmm sampling routines to path
addpath(genpath('./crp'))

%% load iris classes test data (thanks to the UCI machine learning repo)
data = csvread('iris.csv');

class_labels = data(:,end);
data = data(:,1:(end-1));
data = data-repmat(mean(data),size(data,1),1);

% mix up the classes
rp = randperm(size(data,1));

data = data(rp,:);
class_labels = class_labels(rp);     
    


%% set up variables for sampling code
true_number_of_classes = 3;
number_of_observations = size(data,1);
full_dimensionality = size(data,2); 
pca_projection_dimensions_to_retain = full_dimensionality;

% do PCA 
[U,S,V] = svd(cov(data));
reduced_dimensionality_classes = data*U(:,1:pca_projection_dimensions_to_retain);
% report how much of the variance is accounted for by first few dims
percent_of_variance_accounted_for = diag(S)/sum(diag(S));
disp(['Percent of variance accounted for by first ' num2str(pca_projection_dimensions_to_retain) ' eigenvectors (PCA) : ' num2str(sum(percent_of_variance_accounted_for(1:pca_projection_dimensions_to_retain)))]);
% for diagnostic purposes it helps to look at the projection, at minimally
% the first two components
figure(1)
scatter(reduced_dimensionality_classes(:,1),reduced_dimensionality_classes(:,2),[],class_labels)
title('True data')

%% 
% OK, now we have (possibly) reduced dimensionality data.  Let's set up the
% priors for model estimation.  The model requires us to specify:
%
%   \Lambda_0^{-1} : our prior for observation variability about each per 
%                    class mean 
%   v_0 : how much confidence we have in this prior (must be greater than
%   pca_projection_dimensions_to_retain for mathematical reasons)
%   mu_0 : the prior mean a particular observation
%   k_0 : how much confidence we have in this prior
%   a_0, b_0 : gamma(a_0, b_0) on the CRP concentration parameter alpha
%
% Even though these parameters have reasonably clear interpretations, it still 
% can be difficult to hand-specify them (although this is really the way they should
% be specified, i.e. carefully designed to match your understanding about how much 
% variability the data has, etc.  If you don't want to go through with this
% reasonable heuristic fall-back is to repeatedly sample from the model's
% prior until the samples "look like" your data.  I _highly_ recommend that
% you run the section below many, many times to develop an intuition for
% how the model hyperparameters affect things.  Run it multiple times for
% the same hyperparameter values to see the "distribution," then change
% hyperparameter values and do it again to gauge the effect of different
% hyperparameters on the prior.
%%%%%%%%%%%%%%% REPEAT THIS
figure(2)
mu_0 = zeros(pca_projection_dimensions_to_retain,1);
k_0 = .1;
lambda_0 = 1*eye(pca_projection_dimensions_to_retain);
v_0 = 10+pca_projection_dimensions_to_retain+1;
a_0 = 1;
b_0 = 3;
[samples labels means covariances] = sample_igmm_prior(number_of_observations,a_0,b_0,mu_0,lambda_0,k_0,v_0);
scatter(samples(:,1),samples(:,2),[],labels);
%%%%%%%%%%%%%%% END REPEAT


%%
% Starting here is the real meat of the mixture model sampler.  _One_ out of many possible
% usage scenarios are considered:  here a small subset of the overall data is
% processed using batch estimation and then sequential techniques are used to integrate
% the remainder of the data into the model estimate.  The following code could easily
% be rearranged to use batch estimation to estimate the posterior given all data 
% or to use sequential estimation to estimate the posterior given all data one by one.

% ... once you have hyperparameters that work for your data it's time to start model
% estimation.  If there is an enormous amount of data some subsampling of the
% data can be done first.  Here we'll treat all the datapoints via batch
% estimation.  Doing this is _not necessary_ and particle filter estimation
% can be used from the start 

batch_points = size(data,1);  %floor(size(data,1)/2);
batch_training_data = reduced_dimensionality_classes(1:batch_points,:);
batch_training_data_labels = class_labels(1:batch_points,:);
sequential_training_data = reduced_dimensionality_classes(batch_points+1:end,:);
sequential_training_data_labels = class_labels(batch_points+1:end,:);

% num_sweeps is the number of steps in the Gibbs sampler.  The sampler
% should be run to convergence.  Convergence of MCMC methods is very
% difficult to assess automatically -- visually checking that the
% model_scores have stabilized is a practical and practiced way to do it
% but can actually be quite misleading 
num_sweeps = 200;
trace_plot_number = 3;
progress_plot_number = 4;
alpha_0 = 10; 
[class_id_samples, num_classes_per_sample, model_scores, alpha_record] = collapsed_gibbs_sampler(...
    batch_training_data', num_sweeps, a_0, b_0, mu_0, k_0, v_0, ...
    lambda_0, alpha_0, trace_plot_number, progress_plot_number);

%% "samples" from before the sampler has burned in should be thrown away
burned_in_index = 51;
class_id_samples = class_id_samples(:,burned_in_index:end);
num_classes_per_sample = num_classes_per_sample(burned_in_index:end);
model_scores = model_scores(burned_in_index:end);
alpha_record = alpha_record(burned_in_index:end);

% at this point, if we want, we can run some diagnostics to see whether or
% not the model is doing what we want it to do, for instance, we can check
% out the estimated marginal distribution over the number of classes
figure(5)
histogram = histc(num_classes_per_sample,1:21-.5);
bar(1:20,histogram);
lh = line([4 4],[0 max(histogram)+1]);
set(lh,'Color',[0 1 0],'LineWidth',2);
title('Estimated distribution over number of classes')
legend(lh,'True # classes')
% realize that class_id_samples is a collection of equally weighted class
% assignments. The model_scores are proportional to the log probability of the
% model given the data.

%%
% now that we have handled all of the "in-sample" datapoints, we now should
% integrate the "out of sample" data.  We could  just as well run the 
% Gibbs sampler again, but this time with all of the data instead of 
% just the first "in-sample" datapoints.  Instead to do this we switch to
% sequential posterior estimation through particle filtering.
% *Important note* -- the next two lines of code are tied fairly tightly to
% the previous code.  If you want to disconnect the particle filter from the Gibbs
% sampler in order to run only the particle filter then num_particles can
% be set to any value (higher is good, infinite is exact).  The particle
% filter suffers from one drawback, namely, alpha is instead
% treated as a parameter instead of a random variable.  This is a difficult 
% technical issue.
num_particles = num_sweeps-burned_in_index+1;
[class_assignments, assignment_weights, number_of_classes_in_each_particle] = particle_filter(reduced_dimensionality_classes', ...
    num_particles, a_0, b_0, mu_0, k_0, v_0, ...
    lambda_0,mean(alpha_record), class_id_samples);

% You're done!  Now all of the data has been integrated into the model as we now have 
% (weighted) set of data assigned to classes.  Now comes the hard part.  One
% can simply take the MAP assignment (the one with the highest
% weight) but that throws away basically the entire point of having gone
% through all of this in the first point.  Instead, analyses should be run
% on all assignments and then averaged: something like this:

% for n=1:num_particles
%     result(n) = perform_averagable_analysis(spike_sorting(n),
%     per_spike_side_information_like_spike_timing_and_or_stimulus);
% end
% average_result = result.*assignment_weights;

% just for the fun of it though, let's look at the MAP sample
map_class_assignment_index = find(assignment_weights == max(assignment_weights),1);
map_assigment = class_assignments(map_class_assignment_index,:);
figure(6)
scatter(reduced_dimensionality_classes(:,1),reduced_dimensionality_classes(:,2),[],map_assigment);
title('MAP Class Assigment')
% if you compare this to figure 1 you will notice that the colors assigned
% to each class are probably different -- that's because the distribution
% is invariant to permutations of the labeling (the labeling of which
% neuron each spike is attributed to).  This is _extremely consequential_
% to the design of analyses that use the entire posterior distribution.
% Although it may _happen_ that all of the spike sortings have roughly the
% same labeling scheme, that is definitely _not_ something you should ever
% count on in the use of this spike sorting technique

% and we can this analysis, but this time for weighted
% samples
E_number_of_classes = sum(double(number_of_classes_in_each_particle) .* assignment_weights);
MAP_number_of_classes = number_of_classes_in_each_particle(map_class_assignment_index);
disp(['Expected number of classes (' num2str(E_number_of_classes) ') vs. MAP number of classes (' num2str(MAP_number_of_classes) ')'])
