$$\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>
#else
  #include <stdlib.h>
  #define BOOL int
  #define TRUE 1
  #define FALSE 0
#endif

#include <string.h>
#include <iostream>
#include <fstream>

#define _USE_MATH_DEFINES
#include <math.h>
#include <string.h>
#include <set>
#include <map>
#include <vector>
#include <algorithm>

#include "qubx_idealize.h"
#include "../qubfast/max_ll_util.h"
#include "../qubfast/ublas_matrixutil.h"
#include "../qubfast/callbk_reportstream.h"
#include "../qubfast/running_mean.h"
#include "../qubfast/sampled_gaussian.h"
#include "../qubfast/qub_constraints.h"
#include "../qubfast/nmsimplex.h"
#include "../qubfast/qubx_model_storage.h"
#include "../qubfast/qub_kalman.h"

using std::cout;
using std::cerr;
using std::endl;
using std::set;
using std::map;
using std::pair;
using std::copy;
using std::min;
using std::max;

class qubx_idlz_data;


#define QUBX_IDLZ_AA_MAXITER 25
#define QUBX_IDLZ_AMPS_MAXITER 120

// like the telephone keypad:
#define UNSET_VALUE 86738.82583

#ifdef _WIN32
  #include "../qubfast/gettimeofday.h"
#else
  #include <sys/time.h>
#endif
void print_timestamp(const char *lbl)
{
  struct timeval tm;
  gettimeofday(&tm, NULL);
  cerr << "--- " << tm.tv_sec << " : " << (tm.tv_usec / 1000) << "." << (tm.tv_usec % 1000) << " --- " << lbl << endl;
}


typedef map<std::vector<double>, int> stimmap;
void mplex_idl(int Nsignal, int *dwellCounts, int **classeses, float **durationses,
	       double **ampses, int *out_dwellCount, int *out_classes, float *out_durations,
	       stimmap& stimcls_of)
{
  if ( ! Nsignal )
    return;
  // the key to stimmap
  std::vector<double> amps(1+Nsignal);
  // output
  int &Nd = *out_dwellCount;
  Nd = 0;
  // per-signal stuff
  std::vector<int> finger(Nsignal, 0);
  std::vector<int> clss(Nsignal);
  std::vector<float> remain(Nsignal);
  for (int i=0; i<Nsignal; ++i) {
    if ( ! dwellCounts[i] )
      return;
    clss[i] = classeses[i][0];
    remain[i] = durationses[i][0];
    if ( clss[i] >= 0 )
      amps[i+1] = ampses[i][ clss[i] ];
    else
      amps[i+1] = 0.0; // gap at beginning of stimulus: should complain and fail; instead silently defaulting to zero
  }
  while ( 1 ) {
    // until any signal ends
    bool done = false;
    for (int i=0; i<Nsignal; ++i)
      if ( finger[i] >= dwellCounts[i] )
	done = true;
    if ( done )
      break;
    // find signal with shortest remaining
    float tm = remain[0];
    //int changing = 0;
    for (int i=0; i<Nsignal; ++i) {
      if ( remain[i] < tm ) {
	tm = remain[i];
	//changing = i;
      }
    }
    // look up stimclass
    stimmap::iterator sti = stimcls_of.find(amps);
    if ( sti == stimcls_of.end() ) {
      int sc = (int) stimcls_of.size();
      stimcls_of[amps] = sc;
    }
    // append the event
    out_classes[Nd] = stimcls_of[amps];
    out_durations[Nd] = tm;
    ++Nd;
    // shorten all signals remain, and move the finger for any that are 0
    for (int i=0; i<Nsignal; ++i) {
      remain[i] -= tm;
      if ( remain[i] <= 0.0 ) {
	if ( (++finger[i]) < dwellCounts[i] ) {
	  remain[i] = durationses[i][ finger[i] ];
	  clss[i] = classeses[i][ finger[i] ];
	  if ( clss[i] >= 0 )
	    amps[i+1] = ampses[i][ clss[i] ];
	    // else: gap: keep the last event's amp
	}
      }
    }
  }
}


class qubx_idlz_data_piece
{
public:
  int stimclass;
  float *data; // pointer into mac_data->data, or NULL
  int *pt_cls; // pointer into mac_data->data_out, or NULL if count == 0
  float *baseline; // ""
  int count;
  int first;
  
  qubx_idlz_data_piece(int stimclass_, float *data_, int *pt_cls_, float *baseline_, int count_, int first_)
    : stimclass(stimclass_), data(data_), pt_cls(pt_cls_), baseline(baseline_), count(count_), first(first_)
  {}
  int count_dwells()
  {
    if ( ! pt_cls ) return 1;
    int n = 1;
    int c = pt_cls[0];
    for (int i=1; i<count; ++i)
      if ( pt_cls[i] != c ) {
	c = pt_cls[i];
	++n;
      }
    return n;
  }

  int get_dwells(int *ff, int *ll, int *cc)
  {
    if ( ! pt_cls ) {
      *ff = first;
      *ll = first + count - 1;
      *cc = -1;
      return 1;
    }
    int i_out = 0;
    int c = pt_cls[0];
    int f = first;
    int n = 1;
    ff[0] = f;
    cc[0] = c;
    for (int i=1; i<count; ++i)
      if ( pt_cls[i] == c ) {
	++n;
      } else {
	ll[i_out] = f + n - 1;
	++i_out;
	f = ff[i_out] = f + n;
	n = 1;
	c = pt_cls[i];
	cc[i_out] = c;
      }
    ll[i_out] = f + n - 1;
    return ( i_out + 1 );
  }
};

class qubx_idlz_data_seg
{
public:
  std::vector<qubx_idlz_data_piece> pieces;
  int first, count;
  std::vector<double> sc_frac;
  // algo responsibility:
  std::vector<double> amp, std;
  
  qubx_idlz_data_seg()
    : first(0), count(0)
  {}
};

class qubx_idlz_data
{
public:
  double sampling;
  int Nstim, Nsig, Ndata, data_cap;
  float *data;
  int *pt_cls;
  float *baseline;
  std::vector<qubx_idlz_data_seg> segments;
  std::vector<double> stimclasses;
  stimmap stimcls_of;
  std::vector<double> sc_frac;
  
  qubx_idlz_data(double sampling_, int Nstimsig_, int data_cap_)
    : sampling(sampling_), Nstim(0), Nsig(Nstimsig_+1), Ndata(0), data_cap(data_cap_)
    , data(new float[data_cap_]), pt_cls(new int[data_cap_]), baseline(new float[data_cap_])
  {
    memset(baseline, 0, data_cap_*sizeof(float));
  }
  
  virtual ~qubx_idlz_data()
  {
    delete [] data;
    delete [] pt_cls;
    delete [] baseline;
  }
  
  void reset(double sampling_, int Nstimsig, int newcap)
  {
    Nstim = Ndata = 0;
    sampling = sampling_;
    Nsig = Nstimsig + 1;
    segments.clear();
    stimclasses.clear();
    stimcls_of.clear();
    ensure_capacity(newcap);
  }
  
  void ensure_capacity(int newcap)
  {
    if ( data_cap < newcap ) {
      while ( data_cap < newcap )
	data_cap *= 2;
      float *newdata = new float[data_cap];
      int *new_pt_cls = new int[data_cap];
      float *new_base = new float[data_cap];
      memcpy(newdata, data, Ndata*sizeof(float));
      memcpy(new_pt_cls, pt_cls, Ndata*sizeof(int));
      memcpy(new_base, baseline, Ndata*sizeof(float));
      memset(new_base+Ndata, 0, (data_cap - Ndata) * sizeof(float));
      
      for (std::vector<qubx_idlz_data_seg>::iterator segi=segments.begin(); segi!=segments.end(); ++segi) {
	for (std::vector<qubx_idlz_data_piece>::iterator piece = segi->pieces.begin(); piece != segi->pieces.end(); ++piece) {
	  if ( piece->data ) {
	    piece->data = newdata + (piece->data - data);
	    piece->pt_cls = new_pt_cls + (piece->pt_cls - pt_cls);
	    piece->baseline = new_base + (piece->baseline - baseline);
	  }
	}
      }

      delete [] data;
      delete [] pt_cls;
      delete [] baseline;
      data = newdata;
      pt_cls = new_pt_cls;
      baseline = new_base;
      memset(baseline, 0, data_cap*sizeof(float));
    }
  }

  int append_segment(int first)
  {
    int ix = (int) segments.size();
    segments.resize(ix+1);
    segments[ix].first = first;
    return ix;
  }
  
  void append_chunk(int seg_ix, int Nd, float *I,
		    int *dwellCounts, int **classeses, float **durationses, double **ampses)
  {
    if ( I )
      ensure_capacity(Ndata + Nd);
    float *chunk_data = I ? (data + Ndata) : NULL;
    int *chunk_pt_cls = I ? (pt_cls + Ndata) : NULL;
    float *chunk_baseline = I ? (baseline + Ndata) : NULL;
    if ( I ) {
      float *di = chunk_data;
      for (int i=0; i<Nd; ++i, ++di)
	*di = I[i];
      Ndata += Nd;
    }
    
    if ( Nsig < 2 ) {
      segments[seg_ix].pieces.push_back(qubx_idlz_data_piece(0, chunk_data, chunk_pt_cls, chunk_baseline, Nd,
							     segments[seg_ix].first + segments[seg_ix].count));
      segments[seg_ix].count += Nd;
      return;
    }

    int max_mplex = 0;
    for (int i=0; i<(Nsig-1); ++i)
      max_mplex += dwellCounts[i];
    std::vector<int> mplex_clss(max_mplex);
    std::vector<float> mplex_durs(max_mplex);
    int mplex_count = 0;
    mplex_idl(Nsig-1, dwellCounts, classeses, durationses, ampses, &mplex_count, & mplex_clss[0], & mplex_durs[0], stimcls_of);
    
    double sampling_ms = sampling * 1e3;
    int counts_so_far = 0;
    int points_so_far = 0;
    for (int i=0; i<mplex_count; ++i) {
      int count = (int) (mplex_durs[i] / sampling_ms + 0.5);
      int points = I ? count : 0;
      segments[seg_ix].pieces.push_back(qubx_idlz_data_piece(mplex_clss[i],
						     points ? (chunk_data + points_so_far) : NULL,
						     points ? (chunk_pt_cls + points_so_far) : NULL,
						     points ? (chunk_baseline + points_so_far) : NULL,
						     count,
						     segments[seg_ix].first + segments[seg_ix].count + counts_so_far));
      counts_so_far += count;
      points_so_far += points;
    }
    segments[seg_ix].count += counts_so_far;
  }

  double* get_stimclasses()
  {
    if ( stimclasses.size() )
      return & stimclasses[0];
    
    Nstim = (int) stimcls_of.size();
    if ( ! Nstim )
      Nstim = 1;
    stimclasses.resize(Nstim*Nsig);
    for (stimmap::iterator stimi = stimcls_of.begin(); stimi != stimcls_of.end(); ++stimi) {
      int i = stimi->second * Nsig;
      copy(stimi->first.begin(), stimi->first.end(), & stimclasses[i]);
    }
    // calculate which fraction of points are in each stimclass, for each seg and whole data
    sc_frac.resize(Nstim);
    memset(& sc_frac[0], 0, Nstim*sizeof(double));
    int total = 0;
    for (std::vector<qubx_idlz_data_seg>::iterator segi=segments.begin(); segi!=segments.end(); ++segi) {
      segi->sc_frac.resize(Nstim);
      memset(& segi->sc_frac[0], 0, Nstim*sizeof(double));
      int seg_total = 0;
      for (std::vector<qubx_idlz_data_piece>::iterator piece = segi->pieces.begin(); piece != segi->pieces.end(); ++piece) {
	if ( piece->data ) {
	  seg_total += piece->count;
	  segi->sc_frac[piece->stimclass] += piece->count;
	}
      }
      if ( seg_total ) {
	for (int i=0; i<Nstim; ++i) {
	  sc_frac[i] += segi->sc_frac[i];
	  segi->sc_frac[i] /= seg_total;
	  total += seg_total;
	}
      }
    }
    if ( total )
      for (int i=0; i<Nstim; ++i)
	sc_frac[i] /= total;
    return & stimclasses[0];
  }
};




extern "C" MAXILL_API void* qubx_idlz_data_create(double sampling, int Nstimsig, int data_cap)
{
  return new qubx_idlz_data(sampling, Nstimsig, data_cap);
}

extern "C" MAXILL_API void qubx_idlz_data_free(void *data)
{
  delete (qubx_idlz_data *) data;
}

extern "C" MAXILL_API void qubx_idlz_data_reset(void *data, double sampling, int Nstimsig, int data_cap)
{
  ((qubx_idlz_data *) data)->reset(sampling, Nstimsig, data_cap);
}

extern "C" MAXILL_API int qubx_idlz_data_append_segment(void *data, int first)
{
  return ((qubx_idlz_data *) data)->append_segment(first);
}

extern "C" MAXILL_API void qubx_idlz_data_append_chunk(void *data, int seg_ix, int Nd,
								float *I, int *dwellCounts,
								int **classeses, float **durationses, double **ampses)
{
  ((qubx_idlz_data *) data)->append_chunk(seg_ix, Nd, I, dwellCounts, classeses, durationses, ampses);
}

extern "C" MAXILL_API double* qubx_idlz_data_get_stimclasses(void *data, int *Nstimclass)
// returns double[Nstimclass * Nsignal], Nsignal = Nstimsig + 1;
{
  qubx_idlz_data *idldata = (qubx_idlz_data *) data;
  double *stimclasses = idldata->get_stimclasses();
  *Nstimclass = idldata->Nstim;
  return stimclasses;
}

extern "C" MAXILL_API double* qubx_idlz_data_get_stimclass_frac(void *data)
{
  qubx_idlz_data *idldata = (qubx_idlz_data *) data;
  return & idldata->sc_frac[0];
}

extern "C" MAXILL_API int qubx_idlz_data_get_dwellcount(void *data, int iseg, int *ncls)
{
  qubx_idlz_data *idldata = (qubx_idlz_data *) data;
  int count = 0;
  qubx_idlz_data_seg& seg = idldata->segments[iseg];
  int maxcls = -1;
  for (std::vector<qubx_idlz_data_piece>::iterator piece = seg.pieces.begin(); piece != seg.pieces.end(); ++piece) {
    count += piece->count_dwells();
    if ( piece->data )
      for (int i=0; i<piece->count; ++i)
	if ( maxcls < piece->pt_cls[i] )
	  maxcls = piece->pt_cls[i];
  }
  *ncls = maxcls + 1;
  return count;
}

extern "C" MAXILL_API void qubx_idlz_data_get_dwells(void *data, int iseg, int *ff, int *ll, int *cc, double *aa, double *ss)
{
  qubx_idlz_data *idldata = (qubx_idlz_data *) data;
  int count = 0;
  qubx_idlz_data_seg& seg = idldata->segments[iseg];
  for (std::vector<qubx_idlz_data_piece>::iterator piece = seg.pieces.begin(); piece != seg.pieces.end(); ++piece)
    count += piece->get_dwells(ff+count, ll+count, cc+count);
  int maxcls = -1;
  for (int i=0; i<count; ++i)
    if ( cc[i] > maxcls )
      maxcls = cc[i];
  memcpy(aa, & seg.amp[0], (maxcls+1)*sizeof(double));
  memcpy(ss, & seg.std[0], (maxcls+1)*sizeof(double));
}

extern "C" MAXILL_API int qubx_idlz_data_get_baseline(void *data, int iseg, int offset, int count, float *baseline)
{
  qubx_idlz_data *idldata = (qubx_idlz_data *) data;
  if ( (iseg < (int) idldata->segments.size()) && idldata->segments[iseg].pieces.size() ) {
    memcpy(baseline, idldata->segments[iseg].pieces[0].baseline+offset, count*sizeof(float));
    return count;
  } else {
    return 0;
  }
}

void qubx_idlz_set_seg_amps(qubx_stim_amps *ampm, qubx_idlz_data *data)
{
  qubx_model *model = ampm->model;

  for (std::vector<qubx_idlz_data_seg>::iterator segi=data->segments.begin(); segi!=data->segments.end(); ++segi) {
    matrix<double> sc_frac_per_cls(model->Nclass, ampm->Nstim);
    memset(& sc_frac_per_cls(0,0), 0, model->Nclass*ampm->Nstim*sizeof(double));
    std::vector<int> cls_total(model->Nclass, 0);
    for (std::vector<qubx_idlz_data_piece>::iterator piece=segi->pieces.begin(); piece!=segi->pieces.end(); ++piece) {
      if ( ! piece->data )
	continue;
      int sc = piece->stimclass;
      
      for (int i=0; i<piece->count; ++i) {
	int cls = piece->pt_cls[i];
	cls_total[cls] += 1;
	sc_frac_per_cls(cls, sc) += 1.0;
      }
    }
    for (int c=0; c<model->Nclass; ++c) {
      if ( cls_total[c] ) {
	for (int sc=0; sc<ampm->Nstim; ++sc) {
	  sc_frac_per_cls(c, sc) /= cls_total[c];
	}
      }
    }

    segi->amp.resize(model->Nclass);
    segi->std.resize(model->Nclass);
    for (int c=0; c<model->Nclass; ++c) {
      segi->amp[c] = UNSET_VALUE;
      double amp = 0.0;
      double std = 0.0;
      double weight = 0.0;
      for (int sc=0; sc<ampm->Nstim; ++sc) {
	double w = sc_frac_per_cls(c,sc);
	weight += w;
	amp += w * ampm->sc_amp[sc][c];
	std += w * ampm->sc_std[sc][c];
      }
      if ( weight ) {
	segi->amp[c] = amp;
	segi->std[c] = std;
      }
    }
  }
}


class qubx_idlz_halfamp_state
{
public:
  qubx_model *model;
  qubx_stim_amps *ampm;
  std::vector<double> thresh_v;
  std::vector<int> cls_ix_v;
  double *thresh;
  int *cls_ix;
  int cls_last;
  bool omp_stop;
  
  qubx_idlz_halfamp_state(qubx_model *model_, qubx_stim_amps *ampm_)
    : model(model_), ampm(ampm_), thresh_v(model_->Nclass), cls_ix_v(model_->Nclass),
      thresh(&(thresh_v[0])), cls_ix(&(cls_ix_v[0])), omp_stop(false)
  {}
  
  void setup(int sc)
  {
    double *amp = ampm->sc_amp[sc];
    
    std::vector< pair<double, int> > cls_amp_ix;
    for (int c=0; c<model->Nclass; ++c)
      cls_amp_ix.push_back( pair<double, int>(amp[c], c) );
    sort(cls_amp_ix.begin(), cls_amp_ix.end());
    for (int c=0; c<(model->Nclass-1); ++c) {
      thresh[c] = (cls_amp_ix[c+1].first + cls_amp_ix[c].first) / 2.0;
      cls_ix[c] = cls_amp_ix[c].second;
    }
    cls_last = cls_amp_ix[model->Nclass-1].second;
  }
  
  int get_cls(double current)
  {
    for (int c=0; c<(model->Nclass-1); ++c)
      if ( current < thresh[c] )
	return cls_ix[c];
    return cls_last;
  }
};

extern "C" MAXILL_API int qubx_idlz_halfamp(qubx_model *model, void *data, qubx_stim_amps *ampm)
{
  //print_timestamp("halfamp...");
  qubx_idlz_data *idldata = (qubx_idlz_data *) data;
  qubx_idlz_halfamp_state state(model, ampm);

  model->progress = 0;
  int rtn = 0; // success

  #pragma omp parallel
  if ( model->IeqFv ) {
    int idata = 0;
    for (std::vector<qubx_idlz_data_seg>::iterator segi=idldata->segments.begin(); segi!=idldata->segments.end(); ++segi) {
      for (std::vector<qubx_idlz_data_piece>::iterator piece=segi->pieces.begin(); piece!=segi->pieces.end(); ++piece) {
	if ( ! piece->data )
	  continue;
	#pragma omp single
	{
	  state.setup(piece->stimclass);
	}
	
	#pragma omp for
	for (int i=0; i<piece->count; ++i)
	  piece->pt_cls[i] = state.get_cls(piece->data[i]);

	#pragma omp single
	{
	  idata += piece->count;
	  model->progress = (int) (idata * 100.0 / (double) idldata->Ndata);
	  state.omp_stop = state.omp_stop || model->stop_flag;
	}
	if ( state.omp_stop ) {
	  rtn = -10;
	  break;
	}
      }
    }
  } else {
    #pragma omp single
    {
      state.setup(0);
    }
    for (int block=0; block<idldata->Ndata; block += 65536) {
      int block_end = block + 65536;
      if ( idldata->Ndata < block_end )
	block_end = idldata->Ndata;
      #pragma omp single
      {
	model->progress = (int) (block * 100.0 / (double) idldata->Ndata);
	state.omp_stop = state.omp_stop || model->stop_flag;
      }
      if ( state.omp_stop ) {
	rtn = -10;
	break;
      }
      #pragma omp for
      for (int i=block; i<block_end; ++i)
	idldata->pt_cls[i] = state.get_cls(idldata->data[i]);
    }
  }
  
  qubx_idlz_set_seg_amps(ampm, idldata);

  //print_timestamp("done.");
  return rtn;
}



double qubx_idlz_skm_AA_func_simplex(void *obj, double *params);
double qubx_idlz_skm_amp_func_simplex(void *obj, double *params);

class qubx_idlz_skm_state
{
public:
  qubx_model *model;
  qubx_stim_amps *ampm;
  qubx_stim_rates *ratesm;
  qubx_idlz_data *data;
  qubx_model_storage *st;
  double sampling;
  double resolution;
  int b_reest_amps, b_reest_rates, b_track_baseline;
  double baseline_std;
  std::vector< pointy_matrix<double> > AA_log, AA_r;
  std::vector< pointy_col_vector<double> > sc_amp_r, sc_std_r;
  matrix<char> best_prior; // (t, state)
  int t, t0_seg;
  matrix<double> P0_prev;
  std::vector< matrix<double> > P_from;
  qub_kalman_const *kalman_const;
  std::vector<qub_kalman_state *> kalman_prev, kalman_from;
  matrix<float> best_baseline;
  SampledGaussian bl_measure;
  int prev_sc;
  std::vector<double> amp_of_state;
  std::vector< std::vector< SampledGaussian > > measure;
  bool omp_stop;
  double progress_iterweight, progress_before;

  qubx_idlz_skm_state(qubx_model *model_, qubx_stim_amps *ampm_, qubx_stim_rates *ratesm_, qubx_idlz_data *data_,
		      double resolution_, int reest_amps_, int reest_rates_, int track_baseline_, double baseline_std_)
    : model(model_), ampm(ampm_), ratesm(ratesm_), data(data_),
      st((qubx_model_storage*) model->storage), sampling(data->sampling), resolution(resolution_),
      b_reest_amps(reest_amps_), b_reest_rates(reest_rates_), b_track_baseline(track_baseline_), baseline_std(baseline_std_),
      best_prior(data->Ndata, model->Nstate), t(0),
      P0_prev(1, model->Nstate), P_from(model->Nstate),
      kalman_const(NULL), amp_of_state(model_->Nstate), measure(data->Nstim),
      omp_stop(false)
  {
    for (int i=0; i<model->Nstate; ++i)
      P_from[i].resize(1, model->Nstate);
    for (int sc=0; sc<data->Nstim; ++sc)
      measure[sc].resize(model->Nstate);

    AA_log.resize(data->Nstim);
    AA_r.resize(data->Nstim);
    sc_amp_r.resize(data->Nstim);
    sc_std_r.resize(data->Nstim);
    
    if ( model->Nstate ) {
      for (int sc=0; sc<data->Nstim; ++sc) {
	AA_log[sc].resize(model->Nstate, model->Nstate);
	AA_r[sc].resize(model->Nstate, model->Nstate);
	sc_amp_r[sc].resize(st->Nclass);
	sc_std_r[sc].resize(st->Nclass);
      }
      qubx_stim_rates_setup_QQ(ratesm);
      qubx_stim_rates_setup_AA(ratesm);
    }
    
    if ( b_track_baseline ) {
      kalman_const = qub_kalman_create(1, 1);
      // kalman_const.A[0][0] defaults to 1.0 (identity) deterministic process
      kalman_const->H[0][0] = 1.0; // baseline -> measurement
      kalman_const->Q[0][0] = baseline_std * baseline_std; // baseline variance
      kalman_const->R[0][0] = model->var0; // measurement variance
      for (int i=0; i<st->Nstate; ++i) {
	kalman_prev.push_back(qub_kalman_state_create(kalman_const));
	for (int j=0; j<st->Nstate; ++j)
	  kalman_from.push_back(qub_kalman_state_create(kalman_const));
      }
      best_baseline.resize(data->Ndata, model->Nstate);
      bl_measure.setup(0.0, baseline_std, resolution, 6.0, 1024, true, true);
    }
  }
  
  virtual ~qubx_idlz_skm_state()
  {
    for (size_t i=0; i<kalman_prev.size(); ++i)
      qub_kalman_state_free(kalman_prev[i]);
    for (size_t i=0; i<kalman_from.size(); ++i)
      qub_kalman_state_free(kalman_from[i]);
    if ( kalman_const )
      qub_kalman_free(kalman_const);
  }
  
  void setup(int sc)
  {
    for (int j=0; j<model->Nstate; ++j)
      measure[sc][j].setup(ampm->sc_amp[sc][model->clazz[j]],
			   ampm->sc_std[sc][model->clazz[j]],
			   resolution, 6.0, 1024, true, true);
    for (int i=0; i<model->Nstate; ++i)
      for (int j=0; j<model->Nstate; ++j)
	if ( ratesm->AA[sc][i][j] > 0.0 )
	  AA_log[sc].m(i,j) = log(ratesm->AA[sc][i][j]);
	else
	  AA_log[sc].m(i,j) = -1e2;
  }

  void load_sc_amps(int sc)
  {
    for (int i=0; i<model->Nstate; ++i)
      amp_of_state[i] = ampm->sc_amp[sc][model->clazz[i]];
  }
  
  void do_P0(int sc)
  {
    qubx_stim_rates_storage *ratesst = (qubx_stim_rates_storage *) ratesm->storage;
    prev_sc = sc;
    load_sc_amps(sc);
    if ( model->usePeq ) {
      QtoPe(ratesst->QQ[sc].m, P0_prev);
    } else {
      memcpy(&(P0_prev(0,0)), model->P0, model->Nstate * sizeof(double));
    }
    for (int i=0; i<model->Nstate; ++i)
      if ( P0_prev(0,i) > 0 )
	P0_prev(0,i) = log(P0_prev(0,i));
      else
	P0_prev(0,i) = -1e2;
  }
  
  void reset()
  {
    t = 0;
  }
    
  void start_seg(int sc)
  {
    do_P0(sc);
    t0_seg = t;
  }
  
  void first_sample_of_seg(double y, int s)
  {
    if ( b_track_baseline ) {
      double baseline_prior = y - amp_of_state[s];
      qub_kalman_state_reset(kalman_prev[s]);
      kalman_prev[s]->X[0] = baseline_prior;
      y -= kalman_prev[s]->X[0];
      best_baseline(t,s) = (float) kalman_prev[s]->X[0];
    }
    P0_prev(0,s) = measure[prev_sc][s](y); // TODO: += or =  ?
  }
  
  void next_sample_to_lattice(double y, int sc, int prev_s, int next_s)
  {
    double logP_baseline = 0.0;
    if ( b_track_baseline ) {
      double baseline_prior = y - amp_of_state[next_s];
      qub_kalman_state *into = kalman_from[prev_s*st->Nstate + next_s];
      qub_kalman_state_next_into(into, kalman_prev[prev_s], &baseline_prior);
      y -= into->X[0];
      logP_baseline = bl_measure(into->X[0] - kalman_prev[prev_s]->X[0]);
    }
    
    P_from[prev_s](0, next_s) = P0_prev(0,prev_s) + AA_log[prev_sc].m(prev_s, next_s) + measure[sc][next_s](y)
                                + logP_baseline;;
  }
  
  void lattice_to_pick(int next_s)
  {
    double best_P = P_from[0](0,next_s);
    int best_i = 0;
    for (int prev_s=1; prev_s<model->Nstate; ++prev_s) {
      if ( P_from[prev_s](0, next_s) > best_P ) {
	best_P = P_from[prev_s](0, next_s);
	best_i = prev_s;
      }
    }
    best_prior(t, next_s) = (char) best_i;
    P0_prev(0, next_s) = best_P;
    if ( b_track_baseline ) {
      qub_kalman_state_copy(kalman_prev[next_s], kalman_from[best_i*st->Nstate + next_s]);
      best_baseline(t, next_s) = (float) kalman_prev[next_s]->X[0];
    }
  }
  
  void end_sample(int sc)
  {
    prev_sc = sc;
    load_sc_amps(sc);
    ++t;
  }

  void skip_samples(int n, int sc)
  {
    // these numbers are too small for IEEE floating point, so we'll multiply with a proportional vector:
    double logshift = P0_prev(0,0);
    for (int i=1; i<model->Nstate; ++i)
      if ( logshift < P0_prev(0,i) )
	logshift = P0_prev(0,i);
    logshift = - logshift - 1;
    for (int i=0; i<model->Nstate; ++i)
      P0_prev(0,i) = exp(P0_prev(0, i) + logshift);
    
    qubx_stim_rates_storage *ratesst = (qubx_stim_rates_storage *) ratesm->storage;
    ratesst->setup_A(sc, n*sampling);
    P0_prev = prod(P0_prev, ratesst->A.m);
    
    for (int i=0; i<model->Nstate; ++i)
      P0_prev(0, i) = log(P0_prev(0, i)) - logshift;
  }
  
  double end_seg()
  {
    int s_final = 0;     // s = argmax_j( P0_prev(0,j) );
    double best_P = P0_prev(0, s_final);
    for (int j=1; j<model->Nstate; ++j) {
      if ( best_P < P0_prev(0, j) ) {
	best_P = P0_prev(0, j);
	s_final = j;
      }
    }
    int s = s_final;
    for (int tt=t-1; tt>t0_seg; --tt) {
      data->pt_cls[tt] = s;
      s = best_prior(tt, s);
    }
    data->pt_cls[t0_seg] = s;

    if ( b_track_baseline ) {
      int *ss = data->pt_cls + t-1;
      for (int tt=t-1; tt>t0_seg; --tt) {
	data->baseline[tt] = best_baseline(tt,*(ss--));
      }
      data->baseline[t0_seg] = best_baseline(t0_seg,s);
    }
    return best_P;
  }

  void avg_sc_AA()
  {
    for (int sc=0; sc<ratesm->Nstim; ++sc)
      if ( data->sc_frac[sc] )
	AA_r[sc].clear();
    
    matrix<int> sc_from_count(ratesm->Nstim, model->Nstate);
    sc_from_count.clear();
    
    for (std::vector<qubx_idlz_data_seg>::iterator segi=data->segments.begin(); segi!=data->segments.end(); ++segi) {
      for (std::vector<qubx_idlz_data_piece>::iterator piece=segi->pieces.begin(); piece!=segi->pieces.end(); ++piece) {
	if ( ! piece->data )
	  continue;
	int sc = piece->stimclass;
	int from = piece->pt_cls[0];
	for (int i=1; i<piece->count; ++i) {
	  AA_r[sc].m(from, piece->pt_cls[i]) += 1.0;
	  sc_from_count(sc, from) += 1;
	  from = piece->pt_cls[i];
	}
      }
    }
    for (int sc=0; sc<ratesm->Nstim; ++sc) {
      if ( data->sc_frac[sc] ) {
	for (int i=0; i<model->Nstate; ++i) {
	  if ( sc_from_count(sc, i) ) {
	    double scale = 1.0 / sc_from_count(sc, i);
	    for (int j=0; j<model->Nstate; ++j)
	      AA_r[sc].m(i,j) *= scale;
	  }
	}
      }
    }
  }
    
  bool reest_k() {
    avg_sc_AA();
    qubx_stim_rates_setup_constraints(ratesm);
    simplex(qubx_idlz_skm_AA_func_simplex, this, ratesm->cns->fPars, ratesm->cns->NfPar, 1e-9, 1.0, QUBX_IDLZ_AA_MAXITER);
    return true;
  }

  double AA_func_simplex(double *params)
  {
    memcpy(ratesm->cns->fPars, params, ratesm->cns->NfPar*sizeof(double));
    qubx_constrained_free_to_pars(ratesm->cns);
    qubx_stim_rates_pars_to_model(ratesm);
    return AA_eval();
  }

  double AA_eval()
  {
    double ssd = 0.0;
    for (int sc=0; sc<ratesm->Nstim; ++sc) {
      if ( data->sc_frac[sc] ) {
	double sc_ssd = 0.0;
	for (int i=0; i<model->Nstate; ++i) {
	  for (int j=0; j<model->Nstate; ++j) {
	    double diff = ratesm->AA[sc][i][j] - AA_r[sc].m(i,j);
	    sc_ssd += diff * diff;
	  }
	}
	ssd += data->sc_frac[sc] * sc_ssd;
      }
    }
    return ssd;
  }

  void avg_sc_amps()
  {
    if ( model->IeqFv && model->iVoltage ) {
      std::vector<RunningMean> avger(ampm->Nstim * model->Nclass);
      
      for (std::vector<qubx_idlz_data_seg>::iterator segi=data->segments.begin(); segi!=data->segments.end(); ++segi) {
	for (std::vector<qubx_idlz_data_piece>::iterator piece=segi->pieces.begin(); piece!=segi->pieces.end(); ++piece) {
	  if ( ! piece->data )
	    continue;
	  int sc_off = piece->stimclass * model->Nclass;
	  for (int i=0; i<piece->count; ++i)
	    avger[sc_off + piece->pt_cls[i]].add( piece->data[i] - piece->baseline[i] );
	}
      }
      for (int sc=0; sc<ampm->Nstim; ++sc) {
	int sc_off = sc * model->Nclass;
	for (int cls=0; cls<model->Nclass; ++cls) {
	  sc_amp_r[sc].p[cls] = avger[sc_off + cls].mean;
	  sc_std_r[sc].p[cls] = avger[sc_off + cls].std();
	}
      }
    } else { // no v-sensitive amps
      std::vector<RunningMean> avger(model->Nclass);
      for (int i=0; i<data->Ndata; ++i)
	avger[data->pt_cls[i]].add( data->data[i] - data->baseline[i] );
      for (int cls=0; cls<model->Nclass; ++cls) {
	sc_amp_r[0].p[cls] = avger[cls].mean;
	sc_std_r[0].p[cls] = avger[cls].std();
      }
    }
  }
  
  bool reest_amps() {
    avg_sc_amps();
    if ( model->IeqFv && model->iVoltage ) {
      qubx_stim_amps_model_to_pars(ampm);
      qubx_constrained_setup(ampm->cns);
      simplex(qubx_idlz_skm_amp_func_simplex, this, ampm->cns->fPars, ampm->cns->NfPar, 1e-9, 1.0, QUBX_IDLZ_AMPS_MAXITER);
    } else { // no v-sensitive amps: all collected in _r[sc=0]
      memcpy(st->amp.p, sc_amp_r[0].p, st->Nclass*sizeof(double));
      for (int c=0; c<st->Nclass; ++c)
	st->var.p[c] = sc_std_r[0].p[c] * sc_std_r[0].p[c];
      qubx_stim_amps_setup_scs(ampm);
    }
    return true;
  }
  
  double amp_eval(double *params)
  {
    memcpy(ampm->cns->fPars, params, ampm->cns->NfPar * sizeof(double));
    qubx_constrained_free_to_pars(ampm->cns);
    qubx_stim_amps_pars_to_model(ampm);
    qubx_stim_amps_setup_scs(ampm);
    
    double ssd = 0.0;
    for (int sc=0; sc<data->Nstim; ++sc) {
      if ( data->sc_frac[sc] ) {
	double sc_ssd = 0.0;
	for (int i=0; i<model->Nclass; ++i) {
	  double diff = ampm->sc_amp[sc][i] - sc_amp_r[sc].p[i];
	  sc_ssd += diff * diff;
	  diff = ampm->sc_std[sc][i] - sc_std_r[sc].p[i];
	  sc_ssd += diff * diff;
	}
	ssd += data->sc_frac[sc] * sc_ssd;
      }
    }
    return ssd;
  }
};

double qubx_idlz_skm_AA_func_simplex(void *obj, double *params)
{
  return ((qubx_idlz_skm_state *) obj)->AA_func_simplex(params);
}

double qubx_idlz_skm_amp_func_simplex(void *obj, double *params)
{
  return ((qubx_idlz_skm_state *) obj)->amp_eval(params);
}


double qubx_idlz_skm_once(qubx_model *model, qubx_idlz_data *idldata, qubx_idlz_skm_state& state)
{
  double ll = 0.0;
  int idata = 0;
  state.reset();
  
  //print_timestamp("setup sc...");
#pragma omp parallel reduction (+: ll)
  {
    #pragma omp for
    for (int sc=0; sc<idldata->Nstim; ++sc)
      state.setup(sc);
    
    //print_timestamp("seg/piece/point loop...");
    for (std::vector<qubx_idlz_data_seg>::iterator segi=idldata->segments.begin(); segi!=idldata->segments.end(); ++segi) {
      bool first_sample = true;
      for (std::vector<qubx_idlz_data_piece>::iterator piece=segi->pieces.begin(); piece!=segi->pieces.end(); ++piece) {
	if ( ! state.omp_stop ) {
	  int sc = piece->stimclass;
	  int n = piece->count;
	  if ( piece == segi->pieces.begin() ) {
	    #pragma omp single
	    {
	      state.start_seg(sc);
	    }
	  }
	  if ( ! piece->data ) {
	  #pragma omp single
	    {
	      state.skip_samples(n, sc);
	    }
	  } else if ( n ) {
	    int i = 0;
	    if ( first_sample ) {
	    #pragma omp for
	      for (int s=0; s<model->Nstate; ++s)
		state.first_sample_of_seg(piece->data[0], s);
	    #pragma omp single
	      {
		state.end_sample(sc);
		++idata;
	      }
	      i = 1;
	      first_sample = false;
	    }
	    for (; i<n && ! state.omp_stop; ++i) {
	      double y = piece->data[i];
	      #pragma omp for
	      for (int j=model->Nstate*model->Nstate-1; j>=0; --j) {
		state.next_sample_to_lattice(y, sc, j/model->Nstate, j%model->Nstate);
	      }  
	      #pragma omp for
	      for (int j=model->Nstate-1; j>=0; --j) {
		state.lattice_to_pick(j);
	      }
	      #pragma omp single
	      {
		state.end_sample(sc);
		state.omp_stop = state.omp_stop || model->stop_flag;
		++idata;
		if ( ! (idata % 16384) )
		  model->progress = (int) (state.progress_before + 100 * idata * state.progress_iterweight / idldata->Ndata);
	      }
	    }
	  }
	}
      }
      #pragma omp single
      {
	if ( ! state.omp_stop )
	  ll += state.end_seg();
      }
    }
  }

  if ( state.omp_stop ) {
    for (int i=0; i<idldata->Ndata; ++i)
      idldata->pt_cls[i] = 0;
  } else {
    // before converting state index to class, summarize the state-to-state transitions:
    if ( state.b_reest_rates ) {
      //print_timestamp("reestimating Ks...");
      state.reest_k();
    }
    
    //print_timestamp("recovering class sequence...");
    int *clazz = state.st->clazz.p;
    for (int i=0; i<idldata->Ndata; ++i) {
      //cerr << i << ": " << idldata->pt_cls[i] << endl;
      idldata->pt_cls[i] = clazz[idldata->pt_cls[i]];
    }
    
    // now get stats on class amp and std
    if ( state.b_reest_amps ) {
      //print_timestamp("reestimating amps...");
      state.reest_amps(); // bool std?);
    }
  }
  return ll;
}



extern "C" MAXILL_API int qubx_idlz_skm(qubx_model *model, void *data,
					qubx_stim_amps *ampm, qubx_stim_rates *ratesm, double resolution, int method,
					int reest_amps, int reest_rates, int track_baseline, double baseline_std,
					double *ll_out)
{
  *ll_out = 0.0;
  model->progress = 0;
  if ( model->stop_flag )
     return -10;

  //print_timestamp("setup skm...");
  
  qubx_idlz_data *idldata = (qubx_idlz_data *) data;
  qubx_idlz_skm_state state(model, ampm, ratesm, idldata, resolution, reest_amps, reest_rates, track_baseline, baseline_std);
  int rtn = 0; // success

  if ( method == 2/*SKM*/ ) {
    double prev_ll = -1e100;

    int maxiter = 7;
    state.progress_iterweight = 1.0 / maxiter;
    double conv_ll = 1e-8 * idldata->Ndata;
    for (int iter=0; iter<maxiter && !model->stop_flag; ++iter) {
      state.progress_before = iter * state.progress_iterweight;
      double ll = qubx_idlz_skm_once(model, idldata, state);
      model->progress = (int) (state.progress_before + state.progress_iterweight);
      cerr << "ll: " << ll << endl;
      if ( (ll - prev_ll) < conv_ll )
	break;
      *ll_out = prev_ll = ll;
    }
  } else {
    state.progress_iterweight = 1.0;
    state.progress_before = 0.0;
    *ll_out = qubx_idlz_skm_once(model, idldata, state);
  }

  if ( model->stop_flag ) {
    rtn = -10;
  } else {
    //print_timestamp("averaging seg amps...");
    qubx_idlz_set_seg_amps(ampm, idldata);
  }
  //print_timestamp("done.");
  return rtn;
}


// note: don't use openmp; it makes it slower (at least on my 2-core).