/* utilities.c
 * atsa: ATS analysis implementation
 * Oscar Pablo Di Liscia / Pete Moss / Juan Pampin 
 */

#include "atsa.h"

#include <math.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>

/* various usefull macros */
#define freez(p) if(p){free(p); p=0;}
#define swap(type, a, b) { type tmp = (a); (a) = (b); (b) = tmp; }

/* private function prototypes */
int find_next_val_arr(double* arr, int beg, int size);
int find_next_zero_arr(double* arr, int beg, int size);
int find_prev_val_arr(double* arr, int beg);
void fill_sound_gaps(ATS_SOUND *sound, int min_gap_len);
void trim_partials(ATS_SOUND *sound, int min_seg_len, double min_seg_smr);
void set_av(ATS_SOUND *sound);

/* various conversion functions
 * to deal with dB and dB SPL
 * they take and return double doubles
 */
double amp2db(double amp)
{
  return (20* log10(amp));
}

double db2amp(double db)
{
  return (pow(10, db/20));
}

double amp2db_spl(double amp)
{
  return(amp2db(amp) + ATSA_MAX_DB_SPL);
}

double db2amp_spl(double db_spl)
{
  return(db2amp(db_spl - ATSA_MAX_DB_SPL));
}

/* ppp2
 * ====
 * returns the closest power of two
 * greater than num
 */
unsigned int ppp2(unsigned int num)
{
  unsigned int tmp = 2;

  while(tmp<num) tmp = tmp << 1;
  return(tmp);    
}

/* optimize_sound
 * ==============
 * optimizes an ATS_SOUND in memory before saving
 * anargs: pointer to analysis parameters
 * sound: pointer to ATS_SOUND structure
 */
void optimize_sound(ANARGS *anargs, ATS_SOUND *sound)
{
  double ampmax = 0.0, frqmax = 0.0;
  int frame, partial;

  for(frame=0; frame<sound->frames; frame++)
    for(partial=0; partial<sound->partials; partial++) {
      if(ampmax < sound->amp[partial][frame]) ampmax = sound->amp[partial][frame];
      if(frqmax < sound->frq[partial][frame]) frqmax = sound->frq[partial][frame];
    }
  sound->ampmax = ampmax;
  sound->frqmax = frqmax;

  fill_sound_gaps(sound, anargs->min_gap_len);
  trim_partials(sound, anargs->min_seg_len, anargs->min_seg_SMR);
  set_av(sound);
  /* finally set slot to 1 */
  sound->optimized = 1;
}

/* fill_sound_gaps
 * ===============
 * fills gaps in ATS_SOUND partials by interpolation
 * sound: pointer to ATS_SOUND
 * min_gap_len: minimum gap length, gaps shorter or equal to this
 * value will be filled in by interpolation
 */
void fill_sound_gaps(ATS_SOUND *sound, int min_gap_len)
{
  int i, j, k, next_val, next_zero, prev_val, gap_size;
  double f_inc, a_inc, s_inc, mag; 
  mag = TWOPI / (double)sound->srate;
  fprintf(stderr, "Filling sound gaps...\n");
  for(i = 0 ; i < sound->partials ; i++){
    /* first we fix the freq gap before attack */
    next_val = find_next_val_arr(sound->frq[i], 0, sound->frames);
    if( next_val > 0 ){
      for( j = 0 ; j < next_val ; j++ ){
	sound->frq[i][j] =  sound->frq[i][next_val];
      }
    }
    /* fix the freq gap at end */
    prev_val = find_prev_val_arr(sound->frq[i], sound->frames-1);
    if( prev_val != NIL && prev_val < sound->frames-1 ){
      for( j = prev_val ; j < sound->frames ; j++ ){
	sound->frq[i][j] =  sound->frq[i][prev_val];
      }
    }
    /* now we fix inner gaps of frq, pha, and amp */
    k = find_next_val_arr(sound->amp[i], 0, sound->frames);
    while( k < sound->frames && k != NIL){
      /* find next gap: we consider gaps in amplitude, we fix frequency and phase in parallel */
      next_zero = find_next_zero_arr(sound->amp[i], k, sound->frames);
      if( next_zero != NIL){
	prev_val = next_zero - 1;
	next_val = find_next_val_arr(sound->amp[i], next_zero, sound->frames);
	/* check if we didn't get to end of array */
	if( next_val == NIL) break;
	gap_size = next_val - prev_val;
	//	fprintf(stderr, "par: %d prev_val: %d next_val: %d gap_size %d\n", 
	//		i, prev_val, next_val, gap_size);
	/* check validity of found gap */
	if( gap_size <= min_gap_len){
	  //	  fprintf(stderr, "Filling gap of par: %d\n", i);
	  f_inc = (sound->frq[i][next_val] - sound->frq[i][prev_val]) / gap_size;
	  a_inc = (sound->amp[i][next_val] - sound->amp[i][prev_val]) / gap_size;
	  s_inc = (sound->smr[i][next_val] - sound->smr[i][prev_val]) / gap_size;
	  for( j = next_zero ; j < next_val ; j++){
	    sound->frq[i][j] = sound->frq[i][j-1] + f_inc;
	    sound->amp[i][j] = sound->amp[i][j-1] + a_inc;
	    sound->smr[i][j] = sound->smr[i][j-1] + s_inc;
	    sound->pha[i][j] = sound->pha[i][j-1] - (sound->frq[i][j] * sound->frame_size * mag);
	    /* wrap phase */
	    while(sound->pha[i][j] > PI){ sound->pha[i][j] -= TWOPI; }
	    while(sound->pha[i][j] < (PI * -1.0)){ sound->pha[i][j] += TWOPI; }
	  }
	  /* gap fixed, find next gap */
	  k = next_val;
	} else {
	  /* gap not valid, move to next one */
	  //	  fprintf(stderr, "jumping to next_val: %d\n", next_val);
	  k = next_val;
	}
      } else {
	break;
      }
    }
  }
}

/* trim_partials
 * =============
 * trim short segments from ATS_SOUND partials
 * sound: pointer to ATS_SOUND
 * min_seg_len: minimum segment length, segments shorter or equal
 * to this value will be candidates for trimming
 * min_seg_smr: minimum segment average SMR, segment candidates
 * should have an average SMR below this value to be trimmed
 */
void trim_partials(ATS_SOUND *sound, int min_seg_len, double min_seg_smr)
{
  int i, j, k, seg_beg, seg_end, seg_size, count=0;
  double val=0.0, smr_av=0.0;
  fprintf(stderr, "Trimming short partials...\n");
  for(i = 0 ; i < sound->partials ; i++){
    k = 0;
    while( k < sound->frames ){
      /* find next segment */
      seg_beg = find_next_val_arr(sound->amp[i], k, sound->frames);
      if( seg_beg != NIL){
	seg_end = find_next_zero_arr(sound->amp[i], seg_beg, sound->frames);
	/* check if we didn't get to end of array */
	if( seg_end == NIL) seg_end = sound->frames;
	seg_size = seg_end - seg_beg;
	//	fprintf(stderr, "par: %d seg_beg: %d seg_end: %d seg_size %d\n", 
	//		i, seg_beg, seg_end, seg_size);
	if( seg_size <= min_seg_len ){
	  for( j = seg_beg ; j < seg_end ; j++){
	    if( sound->smr[i][j] > 0.0 ){
	      val += sound->smr[i][j];
	      count++;
	    }
	  }
	  if(count > 0) smr_av = val/(double)count;
	  if( smr_av < min_seg_smr ){
	    /* trim segment, only amplitude and SMR data */
	    //	    fprintf(stderr, "Trimming par: %d\n", i);
	    for( j = seg_beg ; j < seg_end ; j++){
	      sound->amp[i][j] = (double)0.0;
	      sound->smr[i][j] = (double)0.0;
	    }
	  } 
	  k = seg_end;
	} else {
	  /* segment not valid, move to the next one */
	  //	  fprintf(stderr, "jumping to seg_end: %d\n", seg_end);
	  k = seg_end;
	}
      } else {
	break;
      }
    }
  }
}

/* auxiliary functions to fill_sound_gaps and trim_partials */
int find_next_val_arr(double* arr, int beg, int size)
{
  int j, next_val=NIL;
  for( j = beg ; j < size ; j++){
    if( arr[j] > 0.0 ){
      next_val = j;
      break;
    }
  }
  return(next_val);
}

int find_next_zero_arr(double* arr, int beg, int size)
{
  int j, next_zero=NIL;
  for( j = beg ; j < size ; j++){
    if( arr[j] == 0.0 ){
      next_zero = j;
      break;
    }
  }
  return(next_zero);
}

int find_prev_val_arr(double* arr, int beg)
{
  int j, prev_val=NIL;
  for( j = beg ; j >= 0 ; j--){
    if( arr[j] != 0.0 ){
      prev_val = j;
      break;
    }
  }
  return(prev_val);
}

/* set_av
 * ======
 * sets the av structure slot of an ATS_SOUND,
 * it computes the average ampl., freq. and SMR for each partial
 * sound: pointer to ATS_SOUND structure
 */
void set_av(ATS_SOUND *sound)
{
  int i, j, count;
  double val;
  fprintf(stderr,"Computing averages..\n");
  for( i = 0 ; i < sound->partials ; i++){
    /* amp */
    val=0.0;
    count = 0;
    for(j = 0 ; j < sound->frames ; j++){
      if(sound->amp[i][j] > 0.0){
	val += sound->amp[i][j];
	count++;
      }
    }
    if(count > 0){
      sound->av[i].amp =  val/(double)count;
    } else {
      sound->av[i].amp = (double)0.0;
    }
    //    fprintf(stderr,"par: %d amp_av: %f\n", i, (double)sound->av[i].amp);
    /* smr */
    val=0.0;
    count = 0;
    for(j = 0 ; j < sound->frames ; j++){
      if(sound->smr[i][j] > 0.0){
	val += sound->smr[i][j];
	count++;
      }
    }
    if(count > 0){
      sound->av[i].smr =  val/(double)count;
    } else {
      sound->av[i].smr = (double)0.0;
    }
    //    fprintf(stderr,"par: %d smr_av: %f\n", i, (double)sound->av[i].smr);
    /* frq */
    val=0.0;
    count = 0;
    for(j = 0 ; j < sound->frames ; j++){
      if(sound->frq[i][j] > 0.0){
	val += sound->frq[i][j];
	count++;
      }
    }
    if(count > 0){
      sound->av[i].frq =  val/(double)count;
    } else {
      sound->av[i].frq = (double)0.0;
    }
    /* set track# */
    sound->av[i].track = i;
  }

  /* sort partials by increasing frequency */
  qsort(sound->av, sound->partials, sizeof(ATS_PEAK), peak_frq_inc);
}

/* init_sound
 * ==========
 * initializes a new sound allocating memory (WARNING: sound pointer must be allocated first)
 */
void init_sound(ATS_SOUND *sound, int sampling_rate, int frame_size, int window_size, int frames, double duration, int partials, int with_noise)
{
  int i, j, k;

  sound->srate = sampling_rate;
  sound->frame_size = frame_size;
  sound->window_size = window_size;
  sound->frames = frames;
  sound->dur = duration;
  sound->partials = partials;
  sound->max_partials = partials;
  sound->optimized = NIL;
  sound->ampmax = NIL;
  sound->frqmax = NIL;
  sound->av = (ATS_PEAK *)malloc(partials * sizeof(ATS_PEAK));
  sound->time = (double **)malloc(partials * sizeof(double *));
  sound->frq = (double **)malloc(partials * sizeof(double *));
  sound->amp = (double **)malloc(partials * sizeof(double *));
  sound->pha = (double **)malloc(partials * sizeof(double *));
  sound->smr = (double **)malloc(partials * sizeof(double *));
  sound->res = (double **)malloc(partials * sizeof(double *));

  /* allocate memory for time, amp, frq, smr, and res data */
  for(k=0; k<partials; k++) {
    sound->time[k] = (double *)malloc(frames * sizeof(double));
    sound->amp[k] = (double *)malloc(frames * sizeof(double));
    sound->frq[k] = (double *)malloc(frames * sizeof(double));
    sound->pha[k] = (double *)malloc(frames * sizeof(double));
    sound->smr[k] = (double *)malloc(frames * sizeof(double));
    sound->res[k] = (double *)malloc(frames * sizeof(double));
  }

  /* init all array values with 0.0 */
  for(i = 0 ; i < partials ; i++){
    for(j = 0 ; j < frames ; j++){
      sound->amp[i][j] = (double)0.0;
      sound->frq[i][j] = (double)0.0;
      sound->pha[i][j] = (double)0.0;
      sound->smr[i][j] = (double)0.0;
      sound->res[i][j] = (double)0.0;
    }
  }

  /* init critical bands noise array if necessary */
  if(with_noise)
  {
    sound->band_energy = (double **)malloc(ATSA_CRITICAL_BANDS * sizeof(double *));
    for (i = 0; i < ATSA_CRITICAL_BANDS; i++)
      sound->band_energy[i] = (double *)malloc(frames * sizeof(double));
  }
  else
    sound->band_energy = 0;
}  

/* copy_sound
 * ==========
 * make a full copy of an existing sound and return a pointer to it
 */
ATS_SOUND *copy_sound(ATS_SOUND *sound)
{
  int i, k;

  ATS_SOUND* sound_copy = (ATS_SOUND*)malloc(sizeof(ATS_SOUND));

  /* Initialize sound copy (alloc memory) */
  init_sound(sound_copy, sound->srate, sound->frame_size, sound->window_size, sound->frames,
	     sound->dur, sound->partials, sound->band_energy != 0);

  /* Copy data from source sound */
  sound_copy->optimized = sound->optimized;
  sound_copy->ampmax = sound->ampmax;
  sound_copy->frqmax = sound->frqmax;

  memcpy(sound_copy->av, sound->av, sound->partials * sizeof(ATS_PEAK));

  for(k=0; k<sound->partials; k++) {
    memcpy(sound_copy->time[k], sound->time[k], sound->frames * sizeof(double));
    memcpy(sound_copy->amp[k], sound->amp[k], sound->frames * sizeof(double));
    memcpy(sound_copy->frq[k], sound->frq[k], sound->frames * sizeof(double));
    memcpy(sound_copy->pha[k], sound->pha[k], sound->frames * sizeof(double));
    memcpy(sound_copy->smr[k], sound->smr[k], sound->frames * sizeof(double));
    memcpy(sound_copy->res[k], sound->res[k], sound->frames * sizeof(double));
  }

  if(sound->band_energy)
    for (i = 0; i < ATSA_CRITICAL_BANDS; i++)
      memcpy(sound_copy->band_energy[i], sound->band_energy[i], sound->frames * sizeof(double));

  return sound_copy;
}

/* increase_max_partials
 * =====================
 * increases by a given offset the max supported number of partials 
 * by reallocating and copying associated partials data arrays
 * (WARNING: sound must have been initialized first)
 */
void increase_max_partials(ATS_SOUND *sound, unsigned int offset)
{
  int k;
  int new_max_partials = sound->max_partials + offset;

  /* Reallocate each partial data array, and keep old data unchanged */
  sound->av = (ATS_PEAK *)realloc(sound->av, new_max_partials * sizeof(ATS_PEAK));

  sound->time = (double **)realloc(sound->time, new_max_partials * sizeof(double *));

  sound->frq = (double **)realloc(sound->frq, new_max_partials * sizeof(double *));

  sound->amp = (double **)realloc(sound->amp, new_max_partials * sizeof(double *));

  sound->pha = (double **)realloc(sound->pha, new_max_partials * sizeof(double *));

  sound->smr = (double **)realloc(sound->smr, new_max_partials * sizeof(double *));

  sound->res = (double **)realloc(sound->res, new_max_partials * sizeof(double *));

  /* allocate memory for time, amp, frq, smr, and res data for new possible partials */
  for(k=sound->max_partials; k<new_max_partials; k++) {
    sound->time[k] = (double *)malloc(sound->frames * sizeof(double));
    sound->amp[k] = (double *)malloc(sound->frames * sizeof(double));
    sound->frq[k] = (double *)malloc(sound->frames * sizeof(double));
    sound->pha[k] = (double *)malloc(sound->frames * sizeof(double));
    sound->smr[k] = (double *)malloc(sound->frames * sizeof(double));
    sound->res[k] = (double *)malloc(sound->frames * sizeof(double));
  }

  /* save new max number of partials */
  sound->max_partials = new_max_partials;
}

/* free_sound
 * ==========
 * frees sound's memory (WARNING: sound pointer is not deallocated)
 */
void free_sound(ATS_SOUND *sound)
{
  int k;

  freez(sound->av);

  /* data */
  for(k=0; k < sound->max_partials; k++) {
    freez(sound->time[k]);
    freez(sound->amp[k]);
    freez(sound->frq[k]);
    freez(sound->pha[k]);
    freez(sound->smr[k]);
    freez(sound->res[k]);
  }

  /* pointers to data */
  freez(sound->time);
  freez(sound->frq);
  freez(sound->amp);
  freez(sound->pha);
  freez(sound->smr);
  freez(sound->res);

  /* check if we have residual data and free its memory up */
  if(sound->band_energy){
    for(k=0 ; k<ATSA_CRITICAL_BANDS ; k++){
      freez(sound->band_energy[k]);
    }
    freez(sound->band_energy);
  }
}

/* add_partial
 * ============
 * creates a new partial that is the copy of the one with given source index if valid ;
 * invalid source indexes (<0 or >=nb partials) give a new partial with null values 
 * (reallocates partials data arrays as necessary, by an offset of ATSA_INCROFFSET,
 *  recomputes the partial average data and maintains the partial sorting by increasing freq.)
 */
void add_partial(ATS_SOUND *sound, int src_index)
{
  int tgt_index = sound->partials;

  /* reallocate ATS_SOUND if not enough room for the new partial */
  if (sound->partials == sound->max_partials)
    increase_max_partials(sound, ATSA_INCROFFSET);

  /* copy source partial at the end of the list if its index is valid */
  if (src_index >= 0 && src_index < sound->partials) {
    memcpy(sound->time[tgt_index], sound->time[src_index], sound->frames * sizeof(double));
    memcpy(sound->frq[tgt_index], sound->frq[src_index], sound->frames * sizeof(double));
    memcpy(sound->amp[tgt_index], sound->amp[src_index], sound->frames * sizeof(double));
    memcpy(sound->pha[tgt_index], sound->pha[src_index], sound->frames * sizeof(double));
    memcpy(sound->smr[tgt_index], sound->smr[src_index], sound->frames * sizeof(double));
    memcpy(sound->res[tgt_index], sound->res[src_index], sound->frames * sizeof(double));
  } else {
    /* other wise, set it to 0 everywhere */
    memset(sound->time[tgt_index], 0, sound->frames * sizeof(double));
    memset(sound->frq[tgt_index], 0, sound->frames * sizeof(double));
    memset(sound->amp[tgt_index], 0, sound->frames * sizeof(double));
    memset(sound->pha[tgt_index], 0, sound->frames * sizeof(double));
    memset(sound->smr[tgt_index], 0, sound->frames * sizeof(double));
    memset(sound->res[tgt_index], 0, sound->frames * sizeof(double));
  }

  /* update the number of partials */
  sound->partials++;

  /* update average data and resort partials by increasing frequency */
  set_av(sound);
}

/* remove_partials
 * ===============
 * removes the partials of given indexes
 * (recompute the partial average data and maintain the partial sorting by increasing freq).
 */
void remove_partials(ATS_SOUND *sound, int* rem_indexes, int nb_rem_indexes)
{
  int part, kept, ind;
  int* kept_indexes;
  int nb_kept;

  /* if nothing to do, do nothing */
  if (!rem_indexes || nb_rem_indexes <= 0)
      return;

  /* determine the indexes of the partials to keep, sorted by increasing index */
  kept_indexes = (int*)malloc(sound->partials*sizeof(int));
  nb_kept = 0;
  for (part=0; part<sound->partials; part++) {
      for (ind = 0; ind < nb_rem_indexes && rem_indexes[ind] != part; ind++);
      if (ind == nb_rem_indexes) /* Is this partial to be kept ? */
	  kept_indexes[nb_kept++] = part;
  }

  /* for each partial to keep : */
  part = 0;
  for (ind = 0; ind < nb_kept; ind++) {

      /* if the partial index changes, move the partial */
      kept = kept_indexes[ind];
      if (kept != part) {
	  swap(double*, sound->time[part], sound->time[kept]);
	  swap(double*, sound->frq[part],  sound->frq[kept]);
	  swap(double*, sound->amp[part],  sound->amp[kept]);
	  swap(double*, sound->pha[part],  sound->pha[kept]);
	  swap(double*, sound->smr[part],  sound->smr[kept]);
	  swap(double*, sound->res[part],  sound->res[kept]);
      }

      /* increment kept partials insertion index */
      part++;
  }

  /* free now useless kept partials index array */
  freez(kept_indexes);

  /* update the number of partials */
  sound->partials = nb_kept;

  /* update average data and resort partials by increasing frequency */
  set_av(sound);
}