function [t, x] = gillespie(S, M, h, c, tspan)
% GILLESPIE Gillespie simulator for stochastic reactions
%
% [t, x] = gillespie(S, M, h, c, tspan)
%
% Inputs:
% S      Stoichiometry matrix (matrix of size n x v, where 
%        n is the number of species and 
%        v is the number of reactions).
% M      Initial state of the system (vector of length n)
% h      Hazard function. Takes two arguments x (vector of length n)
%        and c (vector of length v) as input and
%        outputs a vector of length v.
% c      Stochastic rate constants (vector of length v).
% tspan  Time span, a vector giving the inital and and final times,
%        all the intermediate times are ignored.
%
% Outputs:
% t      Vector of event times
% x      Vector of system states
%
% Note that all vectors are column vectors.
% 
% 2016-2019, Henrik Mannerström, Juho Timonen

[n, ~] = size(S);
max_reactions = 1e5;

% This checks that the inputs have the correct shape
gillespie_input_checks(S, M, h, c, tspan)

% Initialize time and state variables
t = zeros(1, max_reactions);
x = zeros(n, max_reactions);
current_state = M(:);
current_time = tspan(1);

% Main loop
for index = 1:max_reactions
    % Record current state and time
    t(index)    = current_time;
    x(:, index) = current_state;
    
    %% Calculate the reaction rate vector
    rates = h(current_state, c);
    
    %% Sample next reaction event time
    current_time = current_time + exprnd(1/sum(rates));
    
    % Are we done?
    if current_time > tspan(end)
        break;
    end
    
    %% Sample index of the next reaction
    idx_reaction = sample_from_discrete(rates);
   
    %% Update current state
    current_state = current_state + S(:, idx_reaction);
end

% Trim the output vectors
t = t(1:index);
x = x(:, 1:index);

end

% No need to touch this
function idx = sample_from_discrete(p)
% Draw an integer from the set {1,2,3,...,length(p)}
% with probability mass distribution given by
% vector p
idx = find(rand < cumsum(p)/sum(p), 1);

end

% No need to touch this
function gillespie_input_checks(S, M, h, c, tspan)
% A helper function for checking the shape of inputs
% given to the gillespie() function
[n,v] = size(S);

% Check input dimensions
if(length(M)~=n)
    error('Length of M should be equal to the number of rows in S!');
end
if(length(c)~=v)
    error('Length of c should be equal to the number of columns in S!');
end
if(length(tspan)<1)
    error('Length of tspan must be at least 2!');
end

% Check that h works correctly
if(~isa(h,'function_handle'))
    error('h must be a function handle!');
end
if(nargin(h)~=2)
    error('h must take two input arguments!');
end
x_test = ones(length(c),1);
h_test = h(x_test, c);
if(length(h_test)~=length(c))
    error('h must return a vector that has same length as c')
end

end