%% function pointEstimate =  estimateMI(X, Z, params, varargin)
%  Estimates mutual information between signal, X, and
%  cellular response, Z, from cell signalling data
%
% The mutual information estimate is built from estimates of the 
% marginal and conditional entorpy:
%
%       I(X;Z)= H(Z) - H(Z|X) 
%
% Entropy estimation is performed using a kNN estimation procedure 
% (see Kraskov et al PRE 2004). 
%
%    Inputs:
%       X - Nx1 signal vector. N is the number of observed cells.
%
%       Z - NxM matrix responses matrix. N is the number of observed cells
%       and M the dimension of the cellular response.
%       
%       params (optional) - matlab structure containing parameters 
%       'k' and 'xdist'. Parameter 'k' is an integer determining the number
%       of nearest neighbours to be used in the calculation (default value 3). 
%       xdist is a vector with specifying the input distribution 
%       (the empirical distribution is used as a default). 
%
%    Outputs:
%       pointEstimate - a point estimate of the mutual information
%       nCells - the number of cell used for the estimation of the margnial
%       entropy
%       tCells - the total number fo cells in the dataset
%%%%
%   USAGE
%   
%   %%construct the parameter structure
%   params.k=3;
%   params.support = unique(X);
%   params.xdist = ones(1,numel(params.support))/length(params.support);
%   %% call the function                    
%   mi_knnEst = estimateMI(X, Z, params);


function [pointEstimate, nCells, tCells ] = estimateMI(X, Z, params )
    
%% check inputs provided by the user 
if (nargin < 2)
    error('Not enough arguments provided')
end

if (size(X,1) ~= size(Z,1))
    error('X and Z should have the same number of rows')
end

% Set optional parameters to their default values if they have not been provided
if (nargin < 3)
    k=3;
    xs = unique(X);% use the uniques values in X as the support for the signal distribution 
    xDistr = hist(X, xs)/length(X); %emprical distribution
    
else
    %check if all necessary fields are there
    if ( ~isfield(params, 'k') ) 
        error('Error in input parameters: parameter ''k'' is missing. ')
    end
    if (~isfield(params, 'xdist'))
        error('Error in input parameters: parameter ''xdistr'' is missing.')
    end
    if ( ~isfield(params, 'support') ) 
        error('rror in input parameters: parameter ''xs'' is missing.')
    end
    %copy values from the parameter structure
    k = params.k;
    xs = params. support;
    xDistr = params.xdist;
    
    xDistr = reshape(xDistr, 1,numel(xDistr)); % make sure its a column vector
    xs = reshape(xs, numel(xs), 1); % make sure its a column vector
    xDistr = xDistr/sum(xDistr); % make sure is normalised
    % check if signal support and signal distribution are consistent
    if (length(xs) ~= length(xDistr))
        error('Error in input parameters: signal distribution is not consistent wit the support provided. Parameters ''xdistr'' and ''xs'' should have the same size.')
    end
    
    if ( sort(xs) ~= unique(X) )
        error('Error in input parameters: singal support provided does not cover the support of data provided.')
    end
end

%% add some noise to the make sure point do not coincide

Z = Z + 1e-5*randn(size(Z));

%% sample according the distribution provided

Xtmp = X; %copy the X vector
Ztmp = Z; %copy the Z vector

numberPerInput = hist(Xtmp,xs); %% number of cells per input concentration
ms = min(numberPerInput); %% sample as many as the number of cells in the least-populated 
Zsampled = [];

while ms > 1 %% repeat until there

sp = randsample(xs,ms,true,xDistr);
hsp = hist(sp,xs);
ZsampledTmp = zeros(length(sp),min(size(Ztmp)));

for itmp = 1:1:length(xs)
    
    entries = find(Xtmp==xs(itmp)); %% find the entries that correspond to the given concentration 
    ii = randsample(entries,hsp(itmp),false); % sample without replacement from these entries
    ZsampledTmp(sp == xs(itmp),:) = Ztmp(ii,:); % store them
    Xtmp(ii) = []; % remove entries from vectors
    Ztmp(ii,:) = [];
end

Zsampled = [Zsampled ; ZsampledTmp];
numberPerInput = hist(Xtmp,xs);
ms = min(numberPerInput);

end

 %% calculate marginal entropy
 
 HZ = getKNNEstimate(Zsampled,k);

%% estimate conditional entropy

HZx = nan(length(xs),1); %initialise vector of entropies
for i=1:1:length(xs) % for each values in the signal support
    Zcond = Z(X==xs(i),:); % get the corresponding responses
    HZx(i) = getKNNEstimate(Zcond,k); %estimate entropy
end

HZX = xDistr*HZx; % get the conditional entropy as the weigthed average of HZx

%% return the estimate of the mutual information:  H(Z) - H(Z|X) 

pointEstimate = HZ - HZX;

nCells = length(Zsampled);
tCells = length(Z);
end

%% return the k-NN estimate of the entropy
function est = getKNNEstimate(x,k) 

N = max(size(x));
d = min(size(x));

[~,dist] = knnsearch(x,x,'k',k+1);
cd = pi^(d/2)/(2^d * gamma(1+d/2));
est = (-psi(k)+ psi(N) + log(cd)+ d/N* sum(log(2*dist(:,end))))/log(2);

end
