$$\phantom{******** If you see this on the webpage then the browser could not locate *********}$$
$$\phantom{******** jsMath/easy/load.js or the variable root is set wrong in this file *********}$$
$$\newcommand{\vector}[1]{\left[\begin{array}{c} #1 \end{array}\right]}$$ $$\newenvironment{matrix}{\left[\begin{array}{cccccccccc}} {\end{array}\right]}$$ $$\newcommand{\A}{{\cal A}}$$ $$\newcommand{\W}{{\cal W}}$$

/* Copyright 2008-2014 Research Foundation State University of New York   */

/* This file is part of QUB Express.                                      */

/* QUB Express is free software; you can redistribute it and/or modify    */
/* it under the terms of the GNU General Public License as published by   */
/* the Free Software Foundation, either version 3 of the License, or      */
/* (at your option) any later version.                                    */

/* QUB Express is distributed in the hope that it will be useful,         */
/* but WITHOUT ANY WARRANTY; without even the implied warranty of         */
/* MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the          */
/* GNU General Public License for more details.                           */

/* You should have received a copy of the GNU General Public License,     */
/* named LICENSE.txt, in the QUB Express program directory.  If not, see  */
/* <http://www.gnu.org/licenses/>.                                        */

/*

*/

#ifdef _WIN32
  #include <windows.h>
  #include <time.h>
  #include <process.h>
  #define sleep Sleep
  #define getpid _getpid
#else
  #include <stdlib.h>
  #include <unistd.h>
  #define BOOL int
  #define TRUE 1
  #define FALSE 0
#endif

#include <iostream>
using std::cout;
using std::cerr;
using std::endl;
#include <math.h>
#include <string.h>
#include <vector>

#include <boost/tuple/tuple.hpp>
using boost::tuple;

#include <boost/random/mersenne_twister.hpp>
#include <boost/random/uniform_int.hpp>
#include <boost/random/uniform_real.hpp>
#include <boost/random/exponential_distribution.hpp>
#include <boost/random/normal_distribution.hpp>
#include <boost/random/variate_generator.hpp>

#include "qub_simulate.h"
#include "max_ll_util.h"
#include "ublas_plus.h"
#include "ublas_matrixutil.h"


/*
Stimulus interval iterator

  Each signal is either constant or sampled; constants have sample count = -1.
  The iterator presents an array of values[nstimulus] at the current time,
  and the count of samples until it changes.

*/
class stim_intervals
{
private:
  int nsamp;
  int nstim;
  int *counts;
  float **stims;
  std::vector<int> stor_ixs;
  int *ixs;
  std::vector<int> stor_remain;
  int *remain;
  std::vector<double> stor_values;
public:
  double *values;
  int count;
  
  stim_intervals(int nsample, int nstimulus, int *stim_counts, float **stimuli)
    : nsamp( nsample ),
      nstim( nstimulus ),
      counts( stim_counts ),
      stims( stimuli ),
      stor_ixs( nstim ),
      ixs( nstim ? &(stor_ixs[0]) : NULL ),
      stor_remain( nstim ),
      remain( NULL ), // stay NULL in case no stim // remain( nstim ? &(stor_remain[0]) : NULL ),
      stor_values( nstim ),
      values( nstim ? &(stor_values[0]) : NULL ),
      count( 0 )
  {
    reset();
  }
  
  stim_intervals(const stim_intervals& other)
    : nstim( other.nstim ),
      counts( other.counts ),
      stims( other.stims ),
      stor_ixs( nstim ),
      ixs( nstim ? &(stor_ixs[0]) : NULL ),
      stor_remain( nstim ),
      remain( (nstim && other.remain) ? &(stor_remain[0]) : NULL ),
      stor_values( nstim ),
      values( nstim ? &(stor_values[0]) : NULL ),
      count( other.count )
  {
    if ( nstim ) {
      memcpy(ixs, other.ixs, nstim*sizeof(int));
      memcpy(remain, other.remain, nstim*sizeof(int));
      memcpy(values, other.values, nstim*sizeof(int));
    }
  }
  
  void operator=(const stim_intervals& other)
  {
    if (this == &other)
      return;
    bool resize = (nstim != other.nstim);
    nstim = other.nstim;
    counts = other.counts;
    stims = other.stims;
    if ( resize && nstim ) {
      stor_ixs.resize( nstim );
      ixs = &(stor_ixs[0]);
      stor_remain.resize( nstim );
      remain = other.remain ? &(stor_remain[0]) : NULL;
      stor_values.resize( nstim );
      values = &(stor_values[0]);
    }
  }

  bool operator==(const stim_intervals& other) const {
    if ( nstim != other.nstim || counts != other.counts )
      return false;
    for (int i=nstim-1; i>=0; --i)
      if ( ixs[i] != other.ixs[i] )
	return false;
    return true;
  }
  bool operator!=(const stim_intervals& other) const { return ! (*this == other); }

  void reset() {
    if ( nstim == 0 )
      return;
    bool any_vary = false;
    for (int i=0; i<nstim; ++i)
      if ( counts[i] > 0 )
	any_vary = true;
    count = 1;
    remain = &(stor_remain[0]);
    for (int i=0; i<nstim; ++i) {
      if ( counts[i] == 0 ) {
	ixs[i] = 0;
	values[i] = stims[i][0];
	remain[i] = nsamp;
	if ( any_vary )
	  ++remain[i]; // the first ++(*this) will subtract one
      } else {
	ixs[i] = -1;
	remain[i] = 1;
      }
    }
    if ( any_vary )
      ++(*this);
    else
      count = nsamp;
  }
  
  operator bool () const {
    return ( remain != NULL );
  }
  
  stim_intervals& operator++ ()
  {
    if ( remain == NULL )
      return *this;
    int min_remain = (1<<30);
    for (int i=0; i<nstim; ++i) {
      if ( remain[i] > count ) {
	ixs[i] += count;
	remain[i] -= count;
      } else {
	ixs[i] += count;
	if ( ixs[i] < counts[i] ) {
	  float val = stims[i][ixs[i]];
	  values[i] = val;
	  int j = ixs[i]+1;
	  while ( (j<counts[i]) && (stims[i][j] == val) )
	    ++j;
	  remain[i] = j - ixs[i];
	}
	else {
	  count = 0;
	  remain = NULL;
	  return *this;
	}
      }
      if ( remain[i] < min_remain )
	min_remain = remain[i];
    }
    count = min_remain;
    
    //cerr << count << " x ";
    //for (int i=1; i<nstim; ++i)
    //  cerr << values[i] << "\t";
    //cerr << endl;
    
    return *this;
  }
};



/*
One Markov mechanism:

  nstate: 
  st: state index of the most recently completed dwell
  next_st: state it will go to after that (possibly next_st == st, if dur == maxdur)
  dur: seconds spent in st (length of most recently completed dwell)
  int next_event(Q, expo, uni, maxdur):
      Q: matrix of rate constants
      expo: boost::random::variate_generator, exponentially distributed
      uni:  boost::random::variate_generator, uniform between 0 and 1
      maxdur: max seconds in a dwell; if it cuts an event short, st == next_st.

*/
class single_channel
{
public:
  int nstate;
  int st, next_st;
  double dur;
  
  single_channel(int _nstate, int start_state)
    : nstate( _nstate ), next_st( start_state )
  {}
  
  template<class ExpoRand, class UniRand>
  int next_event(matrix<double>& Q, ExpoRand& expo, UniRand& uni, double maxdur)
  // returns exit_state
  {
    st = next_st;
    double full_duration = expo() / (- Q(st, st));
    if ( full_duration > maxdur ) {
      dur = maxdur;
      // next_st = st;
    }
    else {
      dur = full_duration;
      double staterand = uni() * -Q(st, st);
      double tot = 0.0;
      for (int i=0; i<nstate; ++i) {
	if (i != st) {
	  tot += Q(st, i);
	  if ( tot > staterand ) {
	    next_st = i;
		break;
	  }
	}
      }
    }
    return next_st;
  }
};


// boost::random: exponential, normal, and uniform distributions
typedef boost::mt19937 base_generator_type;
typedef boost::uniform_int<> seed_distribution_type;
typedef boost::variate_generator<base_generator_type&, seed_distribution_type> seed_gen_type;
typedef boost::uniform_real<> uni_distribution_type;
typedef boost::variate_generator<base_generator_type&, uni_distribution_type> uni_gen_type;
typedef boost::exponential_distribution<> expo_distribution_type;
typedef boost::variate_generator<base_generator_type&, expo_distribution_type> expo_gen_type;
typedef boost::normal_distribution<> norm_distribution_type;
typedef boost::variate_generator<base_generator_type&, norm_distribution_type> norm_gen_type;

// library lifetime to avoid duplicate seeds
//   func-lifetime generators with passed-in seed
//   simulate{1}() use global generator (maybe better named)
base_generator_type g_generator( static_cast<boost::mt19937::result_type>( clock() + (int)getpid() ) );


// specialization for 1 channel
int qub_simulate_one(double sampling, int nstate, int *start_state, matrix<double>& P0,
		     matrix<double>& K0, matrix<double>& K1, matrix<int>& L, matrix<int>& V,
		     double base_amp, double base_std, double *exc_amp, double *exc_std,
		     QUBFAST_VAR_NOT_USED int *clazz, double *vRev,
		     int nsample, stim_intervals& stim, int v_signal, float *samples, int *states, int *counts,
		     int seed, StatusCallback pctCB, void *pctObj)
{
  base_generator_type generator( static_cast<boost::mt19937::result_type>(seed) );
  uni_gen_type uni_gen(generator, uni_distribution_type(0.0, 1.0));
  expo_gen_type expo_gen(generator, expo_distribution_type(1.0));
  norm_gen_type norm_gen(generator, norm_distribution_type(0.0, 1.0));

  // pick random start state, if it's undefined (< 0)
  if ( *start_state < 0 ) {
    double rand = uni_gen(); // uniform in [0,1) since P0 sums to 1
    double tot = P0(0,0);
    *start_state = 0;
    while ( (tot < rand) && (*start_state < nstate) )
      tot += P0(0, ++(*start_state));
  }
  
  // set up mechanism and rate constants:
  single_channel chan(nstate, *start_state);
  matrix<double> Q(nstate, nstate);
  
  int at = 0; // output sample index
  int evt_at = 0; // output dwell index
  double sam = 0.0; // fractional samples
  double maxdur = 0.0; // seconds remaining before the stimulus changes (or segment ends)

  for (; stim; ++stim) { // for each "count" samples of constant stimulus "values":
    // update rate constants and seconds remaining
    BuildQ(nstate, Q, K0, K1, L, V, stim.values);
    maxdur = sampling * stim.count;
    // fill with events
    while ( maxdur > 0.0 ) {
      chan.next_event(Q, expo_gen, uni_gen, maxdur);
      maxdur -= chan.dur;
      sam += chan.dur / sampling;
      // chan.dur is not an exact multiple of sampling.  Sometimes too short.
      // See if it (cumulatively) crosses the threshold of one sample:
      int evt_n = (int) (sam + 0.5);
      if ( evt_n ) {           // at least one sample long:
	sam -= evt_n;	            // subtract from fractional accumulator
	// write dwell
	states[evt_at] = chan.st;
	counts[evt_at] = evt_n;
	++evt_at;
	// write gaussian data
	float amp, std;
	if ( v_signal ) {
	  float dV = (float) (1e-3 * (stim.values[v_signal] - vRev[chan.st]));
	  amp = (float) (dV * exc_amp[chan.st]);
	  dV = (float) (dV * exc_std[chan.st]); // now it's scaled exc_std
	  std = (float) (sqrt(base_std*base_std + dV*dV));
	} else {
	  amp = (float) exc_amp[chan.st];
	  std = (float) sqrt(base_std*base_std + exc_std[chan.st]*exc_std[chan.st]);
	}
	for (int i=0; i<evt_n; ++i)
	  samples[at+i] = (float) (base_amp + amp + norm_gen()*std);
	at += evt_n;
	// send progress and stop if requested
	if ( pctCB(pctObj, at*1.0/nsample) ) {
	  //cerr << "break" << endl;
	  return evt_at;
	}
      }
    }
  }
  *start_state = chan.next_st;
  
  if ( at != nsample )
    cerr << "wtf!  at=" << at << ", nsample=" << nsample << endl;
  return evt_at;


}


// specialization for n channels
int qub_simulate_multi(int nchannel, double sampling, int nstate, int *start_state, matrix<double>& P0,
		       matrix<double>& K0, matrix<double>& K1, matrix<int>& L, matrix<int>& V,
		       double base_amp, double base_std, double *exc_amp, double *exc_std,
		       QUBFAST_VAR_NOT_USED int *clazz, double *vRev,
		       int nsample, stim_intervals& stim, int v_signal, float *samples, int *states, int *counts,
		       int seed, QUBFAST_VAR_NOT_USED StatusCallback pctCB, QUBFAST_VAR_NOT_USED void *pctObj)
// states, counts for chans[0] only
{ 
  base_generator_type generator( static_cast<boost::mt19937::result_type>(seed) );
  uni_gen_type uni_gen(generator, uni_distribution_type(0.0, 1.0));
  expo_gen_type expo_gen(generator, expo_distribution_type(1.0));
  norm_gen_type norm_gen(generator, norm_distribution_type(0.0, 1.0));

  // set up nchannel mechanisms and bookkeeping
  std::vector<single_channel> chans;
  std::vector<int> remain; // number of samples remaining in the current dwell
  std::vector<double> sam; // fractional samples accumulated from prior dwells
  for (int i=0; i<nchannel; ++i) {
    // pick random start states, if they're undefined (< 0)
    if ( start_state[i] < 0 ) {
      double rand = uni_gen(); // uniform in [0,1) since P0 sums to 1
      double tot = P0(0,0);
      start_state[i] = 0;
      while ( (tot < rand) && (start_state[i] < nstate) )
	tot += P0(0, ++(start_state[i]));
    }
    chans.push_back(single_channel(nstate, start_state[i]));
    remain.push_back(0); // need to get next event (no samples remaining)
    sam.push_back(0.0);  // no fractional samples accumulated yet
  }
  // one rate matrix for all mechanisms, and one set of output arrays.  (states and counts are for mechanism 0 only.)
  matrix<double> Q(nstate, nstate);
  int at = 0;
  int evt_at = 0;
  double maxdur = 0.0;
  double base_var = base_std*base_std; // you can add and subtract variance (not std)
  std::vector<double> dV(nstate);

  for (; stim; ++stim) { // for each "count" samples of constant stimulus "values":
    // update rate constants and seconds remaining
    BuildQ(nstate, Q, K0, K1, L, V, stim.values);
    maxdur = sampling * stim.count;
    if ( v_signal ) {
      double V_here = stim.values[v_signal];
      for (int i=0; i<nstate; ++i)
	dV[i] = 1e-3 * (V_here - vRev[i]);
    }
    // fill with data points
    for (int s=0; s<stim.count; ++s, maxdur-=sampling) {
      // baseline and "instrumentation" noise, regardless of channels
      double amp = base_amp;
      double variance = base_var;
      for (int i=0; i<nchannel; ++i) {
	if ( ! remain[i] ) { // this channel's dwell is done:
	  while ( ! remain[i] ) { // accumulate events until it crosses the threshold of one sample:
	    chans[i].next_event(Q, expo_gen, uni_gen, maxdur);
	    sam[i] += chans[i].dur / sampling;
	    remain[i] = (int) (sam[i] + 0.5);
	    sam[i] -= remain[i];
	  }
	  if ( i == 0 ) { // record the event, mechanism 0 only.
	    states[evt_at] = chans[i].st;
	    counts[evt_at] = remain[i];
	    ++evt_at;
	  }
	}
	// add amp and noise from this mechanism
	if ( v_signal ) {
	  amp += dV[chans[i].st] * exc_amp[chans[i].st];
	  double scaled_std = dV[chans[i].st] * exc_std[chans[i].st];
	  variance += scaled_std * scaled_std;
	} else {
	  amp += exc_amp[chans[i].st];
	  variance += exc_std[chans[i].st] * exc_std[chans[i].st];
	}
	--remain[i]; // count this sample
      }
      // output one gaussian datapoint.  Should be equivalent to
      // base_amp + base_std*norm_gen() + sum_i(exc_amp[ch[i].st] + exc_std[ch[i].st]*norm_gen())
      // but faster.
      samples[at++] = (float) (amp + sqrt(variance) * norm_gen());
    }
  }
  // output states so you could call again to simulate the continuation
  for (int i=0; i<nchannel; ++i)
    start_state[i] = chans[i].next_st;

  if ( at != nsample )
    cerr << "wtf!  at=" << at << ", nsample=" << nsample << endl;
  return evt_at;

}

  // returns event count
  // P0 == NULL: P0 = QtoPe(Q)
  // start_states[i] < 0: rand from P0
  // start_states is updated to hold exit (next) states
  // stim_counts[i] == 0: stimuli[i][0] holds constant value
extern "C" QUBFAST_API int    qub_simulate(double sampling, int nchannel, int nstate, int *start_states, double *P0,
				double *K0, double *K1, int *L, int *V,
				double base_amp, double base_std, double *exc_amp, double *exc_std,
				int nsample, int nstimulus, int *stim_counts, float **stimuli,
				float *samples, int *states, int *counts,
				StatusCallback pctCB, void *pctObj)
{
  std::vector<int> clazz;
  for (int i=0; i<nstate; ++i)
    clazz.push_back(i);
  seed_gen_type seed_gen(g_generator, seed_distribution_type(-(1<<31), 1<<31));
  int seed = seed_gen();
  return qub_simulate2(sampling, nchannel, nstate, start_states, P0, K0, K1, L, V,
		       base_amp, base_std, exc_amp, exc_std, &(clazz[0]), NULL,
		       nsample, nstimulus, stim_counts, stimuli, 0,
		       samples, states, counts, seed, pctCB, pctObj);
}

QUBFAST_API int    qub_simulate2(double sampling, int nchannel, int nstate, int *start_states, double *P0,
				 double *K0, double *K1, int *L, int *V, 
				 double base_amp, double base_std, double *exc_amp, double *exc_std,
				 int *clazz, double *vRev,
				 int nsample, int nstimulus, int *stim_counts, float **stimuli,
				 int v_signal,
				 float *samples, int *states, int *counts,
				 int seed, StatusCallback pctCB, void *pctObj)
/*
Simulates sampled and state-idealized HMM data and returns the number of dwells written into states and counts.

    @param sampling: interval between samples, in seconds
    @param nchannel: number of identical mechanisms' responses to sum together
    @param nstate: number of model states (vertices)
    @param start_states: array[nstate] of entry state indices, or -1 to pick fresh.  On output, contains exit states.
    @param P0: array[Nstate] of entry probabilities, or NULL
    @param K0: row-major matrix, nstate x nstate; element (i,j) is the pre-exponential rate constant between states i and j
    @param K1: row-major matrix, nstate x nstate; element (i,j) is the exponential rate constant between states i and j
    @param L:  row-major matrix, nstate x nstate; element (i,j) is the pre-exponential stimulus index, or 0
    @param V:  row-major matrix, nstate x nstate; element (i,j) is the exponential stimulus index, or 0
    @param base_amp: constant baseline offset, amplitude of closed (class 0, black)
    @param base_std: baseline std; instrument rather than channel noise
    @param exc_amp: array[Nstate] of (state amp) - base_amp
    @param exc_std: array[Nstate] of sqrt((state std)**2 - base_std**2)
    @param clazz: class index of each state
    @param vRev: reversal potential, in mV, of each state; or NULL if (! v_signal)
    @param nsample: number of data points to return in samples
    @param nstimulus: number of stimulus signals, plus signal 0 which is ignored (so at least 1)
    @param stim_counts: number of datapoints in each signal; -1 if it's constant
    @param stimuli: array[max(1,stim_counts[i])] of datapoints for each signal
    @param v_signal: index of Voltage signal, if exc_amp is to be scaled by 1e-3*(Voltage-vRev)
    @param samples: output array[nsample] for data points
    @param states: output array[nsample] for idealized state index of each dwell
    @param counts: output array[nsample] for number of samples in each dwell
    @param seed: random seed; should be different each invocation
    @param pctCB: callback pctCB(pctObj, complete), 0.0 <= complete <= 1.0; Return nonzero to stop the simulation.
    @param pctObj: data pointer for the callback function, or NULL
  


 */

{
  // copy into boost::matrix 
  matrix<double> _K0(nstate, nstate);       matrix_from_array(_K0, K0, nstate, nstate);
  matrix<double> _K1(nstate, nstate);       matrix_from_array(_K1, K1, nstate, nstate);
  matrix<int>    _L (nstate, nstate);       matrix_from_array(_L,  L,  nstate, nstate);
  matrix<int>    _V (nstate, nstate);       matrix_from_array(_V,  V,  nstate, nstate);
  
  // set up stimulus iterator
  stim_intervals stim(nsample, nstimulus, stim_counts, stimuli);
  
  // copy entry probabilities, or calculate equilibrium at initial stimulus
  matrix<double> _P0(nstate, nstate);
  if ( P0 ) {
    matrix_from_array(_P0, P0, 1, nstate);
  } else {
    matrix<double> Q(nstate, nstate);
    BuildQ(nstate, Q, _K0, _K1, _L, _V, stim.values);
    QtoPe(Q, _P0);
  }
  
  if ( nchannel <= 1 ) {
    return qub_simulate_one(sampling, nstate, start_states, _P0, _K0, _K1, _L, _V,
			    base_amp, base_std, exc_amp, exc_std, clazz, vRev,
			    nsample, stim, v_signal, samples, states, counts,
			    seed, pctCB, pctObj);
  } else {
    return qub_simulate_multi(nchannel, sampling, nstate, start_states, _P0, _K0, _K1, _L, _V,
			      base_amp, base_std, exc_amp, exc_std, clazz, vRev,
			      nsample, stim, v_signal, samples, states, counts,
			      seed, pctCB, pctObj);
  }
}