#define TRACERPLUS (-5)
#include "lasp_sosfilterbank.h"
#include "lasp_mq.h"
#include "lasp_worker.h"
#include "lasp_nprocs.h"



typedef struct Sosfilterbank {
  /// The filter_coefs matrix contains filter coefficients for a SOS filter.
  us filterbank_size;
  us nsections;

  /// The filter coefficients for each of the filters in the Filterbank
  /// The *first* axis is the filter no, the second axis contains the
  /// filter coefficients, in the order, b_0, b_1, b_2, a_0, a_1, a_2, which
  /// corresponds to the transfer function 
  ///                 b_0 + b_1 z^-1 + b_2 z^-2
  ///         H[z] =  -------------------------
  ///                 a_0 + a_1 z^-1 + a_2 z^-2 
  dmat sos; /// sos[filter_no, coeff]

  /// Storage for the current state of the output, first axis correspond to
  /// the filter number axis, the second axis contains state coefficients 
  dmat state;
#ifdef LASP_PARALLEL
  JobQueue* jq;
  Workers* workers;
#endif // LASP_PARALLEL


} Sosfilterbank;

us Sosfilterbank_getFilterbankSize(const Sosfilterbank* fb) {
  fsTRACE(15);
  assertvalidptr(fb);
  return fb->filterbank_size;
  feTRACE(15);

}

int filter_single(void* worker_data,void* job);

static inline us min(us a, us b){
  return a<b?a:b; 
}
static inline us max(us a, us b){
  return a>b?a:b;
}

Sosfilterbank* Sosfilterbank_create(
    const us nthreads_,
    const us filterbank_size,
    const us nsections) {
  fsTRACE(15);

  dbgassert(filterbank_size <= MAX_SOS_FILTER_BANK_SIZE, 
      "Illegal filterbank size. Max size is " 
      annestr(MAX_SOS_FILTER_BANK_SIZE));

  Sosfilterbank* fb = (Sosfilterbank*) a_malloc(sizeof(Sosfilterbank));

  fb->filterbank_size = filterbank_size;
  dbgassert(nsections < MAX_SOS_FILTER_BANK_NSECTIONS,"Illegal number of sections");
  fb->nsections = nsections;

  /// Allocate filter coefficients matrix
  fb->sos = dmat_alloc(filterbank_size, nsections*6);
  fb->state = dmat_alloc(filterbank_size, nsections*2);
  dmat_set(&(fb->state), 0);

  /// Set all filter coefficients to unit impulse response
  vd imp_response = vd_alloc(6*nsections);
  vd_set(&imp_response,0);
  for(us section = 0;section < nsections; section++) {
    // Set b0 coefficient to 1
    setvecval(&imp_response, 0 + 6*section, 1);
    // Set a0 coefficient to 1
    setvecval(&imp_response, 3 + 6*section, 1);
  }

  // Initialize all filters with a simple impulse response, single pass
  for(us filter_no = 0; filter_no < filterbank_size; filter_no++) {
    Sosfilterbank_setFilter(fb,filter_no,imp_response);
  }
  // Check if coefficients are properly initialized
  // print_dmat(&(fb->sos));
  vd_free(&imp_response);

#ifdef LASP_PARALLEL
  fb->jq = NULL;
  fb->workers = NULL;
  dbgassert(nthreads_ <= LASP_MAX_NUM_THREADS, "Illegal number of threads");
  us nthreads;
  us nprocs = getNumberOfProcs();

  if(nthreads_ == 0) {
    nthreads = min(max(nprocs/2,1), filterbank_size);
  } else {
    nthreads = nthreads_;
  }
  iVARTRACE(15, nthreads);

  if(nthreads > 1) {
    if(!(fb->jq = JobQueue_alloc(filterbank_size))) {
      Sosfilterbank_free(fb);
      feTRACE(15);
      return NULL;
    }

    if(!(fb->workers = Workers_create(nthreads,
            fb->jq,
            NULL,
            &filter_single,
            NULL,
            NULL))) {
      Sosfilterbank_free(fb);
      feTRACE(15);
      return NULL;
    }
  }

#endif // LASP_PARALLEL
  feTRACE(15);
  return fb;
}

void Sosfilterbank_setFilter(Sosfilterbank* fb,const us filter_no,
    const vd filter_coefs) {
  fsTRACE(15);
  assertvalidptr(fb);
  assert_vx(&filter_coefs);
  iVARTRACE(15, filter_coefs.n_rows);
  iVARTRACE(15, filter_no);
  dbgassert(filter_no < fb->filterbank_size, "Illegal filter number");
  dbgassert(filter_coefs.n_rows == fb->nsections * 6, 
      "Illegal filter coefficient length");

  dmat *sos = &fb->sos;
  /* dmat *state = &fb->state; */
  us nsections = fb->nsections;

  for(us index=0;index<nsections*6;index++){
    // Copy contents to position in sos matrix
    *getdmatval(sos,filter_no,index) = *getvdval(&filter_coefs,index);
  }

  feTRACE(15);
}

void Sosfilterbank_free(Sosfilterbank* fb) {
  fsTRACE(15);
  assertvalidptr(fb);

  dmat_free(&(fb->sos));
  dmat_free(&(fb->state));

#ifdef LASP_PARALLEL
  if(fb->workers) Workers_free(fb->workers);
  if(fb->jq) JobQueue_free(fb->jq);
#endif // LASP_PARALLEL

  a_free(fb);
  feTRACE(15);

}

typedef struct {
  us filter_no;
  us nsections;
  us nsamples; 
  dmat sos;
  dmat state;
  dmat ys;
} Job;

int filter_single(void* worker_data, void* job_) {
  fsTRACE(15);
  Job* job = (Job*) job_;

  us nsections = job->nsections;
  us nsamples = job->nsamples;
  dmat sos = job->sos;


  for(us section=0;section<nsections;section++) {
    d w1 = *getdmatval(&(job->state),job->filter_no,section*2); 
    d w2 = *getdmatval(&(job->state),job->filter_no,section*2+1); 

    d b0 = *getdmatval(&sos,job->filter_no,section*6+0);
    d b1 = *getdmatval(&sos,job->filter_no,section*6+1);
    d b2 = *getdmatval(&sos,job->filter_no,section*6+2);
    /* d a0 = *getdmatval(&sos,job->filter_no,section*6+3); */
    d a1 = *getdmatval(&sos,job->filter_no,section*6+4);
    d a2 = *getdmatval(&sos,job->filter_no,section*6+5);

    d* y = getdmatval(&(job->ys), 0, job->filter_no);

    for(us sample=0;sample<nsamples;sample++){
      d w0 = *y - a1*w1 - a2*w2;
      d yn = b0*w0 + b1*w1 + b2*w2;
      w2 = w1;
      w1 = w0;
      *y++ = yn;
    }
    *getdmatval(&(job->state),job->filter_no,section*2)   = w1;
    *getdmatval(&(job->state),job->filter_no,section*2+1) = w2; 

  }

  feTRACE(15);
  return 0;
}


dmat Sosfilterbank_filter(Sosfilterbank* fb,const vd* xs) {

  fsTRACE(15);
  assertvalidptr(fb);
  assert_vx(xs);
  dmat state = fb->state;
  dmat sos = fb->sos;

  us nsections = fb->nsections;
  us filterbank_size = fb->filterbank_size;
  us nsamples = xs->n_rows;

  dmat ys = dmat_alloc(nsamples, filterbank_size);
  /// Copy input signal to output array
  for(us filter=0;filter<filterbank_size;filter++) {
    d_copy(getdmatval(&ys,0,filter),getvdval(xs,0),nsamples,1,1);
  }

  Job jobs[MAX_SOS_FILTER_BANK_SIZE];
  /// Implementation is based on Proakis & Manolakis - Digital Signal
  /// Processing, Fourth Edition, p. 550 
  Job job_template = {0, nsections, nsamples, sos, state, ys};

  for(us filter=0;filter<filterbank_size;filter++) {
    /// Obtain state information for current section, and all filters
    jobs[filter] = job_template;
    jobs[filter].filter_no = filter;

#ifdef LASP_PARALLEL
    if(fb->workers) {
      assertvalidptr(fb->jq);
      JobQueue_push(fb->jq, &(jobs[filter]));
    } else {

#endif // LASP_PARALLEL
      /* No workers, we have to do it ourselves */ 
      filter_single(NULL,(void*) &(jobs[filter]));
#ifdef LASP_PARALLEL
    }
#endif // LASP_PARALLEL
  }
#ifdef LASP_PARALLEL
  if(fb->workers) {
    JobQueue_wait_alldone(fb->jq);
  }
#endif // LASP_PARALLEL
  feTRACE(15);
  return ys;
}