/*-------------------------------------------------------
 * Package: 	Wavelet Analysis Tool
 * File name: 	wavedata.cc
 *-------------------------------------------------------
*/
//________________________________________________________
// WaveData class is the base class for generic wavelet
// data structures and methods. 

#ifndef _STREAM_H
  #include <stream.h>
  #define _STREAM_H
#endif

#ifndef _WDATA_H
  #include "wavedata.h"
#endif

#ifdef _USE_ROOT
  ClassImp(WaveData)
#endif

/*********************************************************************
 * Default constructor
 *********************************************************************/
    WaveData::WaveData() : N(0), data(NULL) {}

/*********************************************************************
 * Constructor which allocates space for data array with
 * N=n elements of type "wavereal".
 *********************************************************************/
 WaveData::WaveData(int n)
{ 
  if (n <= 0 ) n = 1;
  data = new wavereal[n];
  N = n;
}

/*********************************************************************
 * Copy constructor 
 *********************************************************************/
 WaveData::WaveData(const WaveData& a)
{ 
   N = a.N;
   Rate = a.Rate; 
   data = NULL;
   if (a.N>0) {
      data = new wavereal[a.N];
      for (int i=0; i < N; i++) data[i] = a.data[i];
   }
}

/*********************************************************************
 * Destructor
 *********************************************************************/
 WaveData::~WaveData()
{ 
   if(data != NULL) delete [] data;
}

WaveData& WaveData::operator=(const WaveData& a)
{
  if (this != &a) {
    if (N != a.N) Resize(a.N);
    for (int i=0; i < N; i++) data[i] = a.data[i];
    Rate = a.Rate; 
  }
  return *this;
}

WaveData& WaveData::operator=(const WaveData* a)
{ return *this=*a; }

// assign value c to all array elements
WaveData& WaveData::operator=(const wavereal c)
{
  for (int i=0; i < N; i++) data[i] = c;
  return *this;
}

WaveData& WaveData::operator+=(const WaveData &a)
{
  if (Rate != a.Rate) {
   cout <<
   " Error: attempt to use operator += for data with different sample rate.n";
   cout <<" Data unchanged!n";
  }

  if (N != a.N) {
   cout <<
   " Error: attempt to use operator += for data with different length.n";
   cout <<" Data unchanged!n";
   return *this;
  }
   
  for (int i=0; i < N; i++) data[i] += a.data[i];
  return *this;
}

WaveData& WaveData::operator+=(const WaveData *a)
{ return *this+=*a; }

WaveData& WaveData::operator-=(const WaveData &a)
{
  if (Rate != a.Rate) {
    cout <<
    " Error: attempt to use operator -= for data with different sample rate.n";
    cout <<" Data unchanged!n";
    return *this;
  }

  if (N != a.N) {
    cout <<
    " Error: attempt to use operator -= for data with different length.n";
    cout <<" Data unchanged!n";
    return *this;
  }

  for (int i=0; i < N; i++) data[i] -= a.data[i];
  return *this;
}

WaveData& WaveData::operator-=(const WaveData *a) 
{ return *this-=*a; }

// add constant shift to all elements of array
WaveData& WaveData::operator+=(const wavereal c)
{
  for (int i=0; i < N; i++) data[i] += c;
  return *this;
}

/*********************************************************************
 * subtract constant shift from all elements of array
 *********************************************************************/
WaveData& WaveData::operator-=(const wavereal c)
{
  for (int i=0; i < N; i++) data[i] -= c;
  return *this;
}

/*********************************************************************
 * multiply all elements of data array by constant
 *********************************************************************/
WaveData& WaveData::operator*=(const wavereal c)
{
  for (int i=0; i < N; i++) data[i] *= c;
  return *this;
}

/*********************************************************************
 * Dumps data array to file *fname in ASCII format.
 *********************************************************************/
 void WaveData::Dump(const char *fname, int app)
{
// Dumps data array to file *fname in ASCII format. 
//
  int n=N;
  char *mode = "w";
  if (app == 1) mode = "a";

  FILE *fp;

  if ( (fp = fopen(fname, mode)) == NULL ) {
     cout << " Dump() error: cannot open file " << fname <<". n";
     return;
  };

  for (int i = 0; i < n; i++) fprintf( fp,"%e n", data[i]);
  fclose(fp); 
}
  

/*********************************************************************
 * Dumps data array to file *fname in binary format and type "wavereal".
 *********************************************************************/
 void WaveData::DumpBinary(const char *fname, int app)
{
// Dumps data array to file *fname in binary format and type "wavereal".
//
  int n = N * sizeof(wavereal);
  char *mode = "wb";

  if (app == 1) mode = "ab";

  FILE *fp;

  if ( (fp=fopen(fname, mode)) == NULL ) {
     cout << " DumpBinary() error : cannot open file " << fname <<". n";
     return ;
  }

//  cout << "Writing binary record, size="<< n <<"n";
  fwrite(data, n, 1, fp);
  fclose(fp);
}

/*********************************************************************
 * Dumps data array to file *fname in binary format and type "short".
 *********************************************************************/
 void WaveData::DumpShort(const char *fname, int app)
{
// Dumps data array to file *fname in binary format and type "short".
//
  int n = N;
  char *mode = "wb";
  if (app == 1) mode = "ab";

  FILE *fp;
  if ( (fp = fopen(fname, mode)) == NULL ) {
    cout << " DumpShort() error : cannot open file " << fname <<". n";
    return;
  }

  short *dtemp;
  dtemp=new short[n];

  for ( int i=0; i<n; i++ ) dtemp[i]=short(data[i]) ;

  n = n * sizeof(short);
//  cout << " Writing binary record, size="<< n <<"n";
  fwrite(dtemp, n, 1, fp);
  fclose(fp);
  delete [] dtemp;
}

 void WaveData::ReadBinary(const char *fname)
{
// Read data from file in binary format. Type of data is wavereal.
//

  int step = sizeof(wavereal);
  FILE *fp;
  wavereal d;

  if ( (fp=fopen(fname,"rb")) == NULL ) {
     cout << " ReadBinary() error : cannot open file " << fname <<". n";
     return;
  }


  if(N == 0){              // find the data length
     while(!feof(fp)){ 
	fread(&d,step,1,fp);
	N++;
     }
     N--;
     rewind(fp);
     Resize(N);
  }	

  int n = N * sizeof(wavereal);

  fread(data, n, 1, fp);    // Reading binary record
  fclose(fp);
}

 void WaveData::ReadShort(const char *fname)
{
// Read data from file with given filename in written in binary format as short.
//

  short *dtmp;
  dtmp = new short[N];
  int n = N * sizeof(short);
  FILE *fp;

  if ( (fp=fopen(fname,"rb")) == NULL ) {
     cout << " ReadShort() error : cannot open file " << fname <<". n";
     return;
  }

  cout << " Reading binary record, size="<< n <<"n";
  fread(dtmp, n, 1, fp);
  for (int i = 0; i < N; i++) data[i] = wavereal(dtmp[i]);
  fclose(fp);
}


/*********************************************************************
 * Dumps spectrum of current signal to text file *fname,
 * option=2 corresponds to power spectrum,
 * option=1 corresponds to square root of power spectrum (default)
 * This function does not change the original signal
 *********************************************************************/
 void WaveData::DumpSpectrum(const char *fname, int option)
{
  wavereal x,y,s,df;
  FILE *fp;
  if ( (fp=fopen(fname,"w")) == NULL )
       { cout << "Cannot open file " << fname <<". n"; return;};

  WaveData *td;

  td=new WaveData(N);

  for (int i=0;i<N;i++) td->data[i]=data[i];

  td->FFT();

  x=td->data[0];
  s=sqrt(x*x);
  df=Rate/N;

// print value for frequency=0
  fprintf(fp,"%14.6f %14.6en",0.,s);

// print Fourier without F[0]
  for (int i=1;i<N/2;i++)
  { x=td->data[2*i];
    y=td->data[2*i+1];
    s=2.*(x*x+y*y);	

// option=2 mean power spectrum
// =1 -> sqrt(power) spectrum 
    if (option==1) s=sqrt(x*x+y*y);

    fprintf(fp,"%14.6f %14.6en",df*i,s);
  }						
// F[n/2] is not printed
  
  fclose(fp);
  delete td;
}


/*********************************************************************
 * Resizes data array of the object to a new value N=n.
 *********************************************************************/
 void WaveData::Resize(int n)
{
  if(data != NULL) delete [] data;
  data = new wavereal[n];
  N = n;
}


/*********************************************************************
 * Calculates and returns mean and root mean square of the signal.
 *********************************************************************/
inline void WaveData::getStatistics(wavereal &mean, wavereal &rms)
{
//  Calculates and returns mean and root mean square of the signal.
//
  wavereal v;
  mean = 0.;
  rms = 0.;

  for (int i=0;i<N;i++) {
    v=data[i];
    mean +=v;
    rms += v*v;
  }

  mean /= double(N);
  rms /= double(N);
  rms -= mean*mean;

  rms=sqrt(rms);

  return;
}


/************************************************************************
 * Creates new data set by resampling original data from "WaveData &a"  *
 * with new sample frequency "f". Uses polynomial interpolation scheme  *
 * (Lagrange formula) with np-points, "np" must be even number, by      *
 * default np=6 (corresponds to 5-th order polynomial interpolation).   *
 * This function calls WaveData::Resize() function to adjust            *
 * current data array if it's size does not exactly match new number    *
 * of points.                                                           *
 ************************************************************************/
 void WaveData::Resample(WaveData const &a, wavereal f, int np)
{
  int i1, imax, n1, n2, n, np2 = np/2;
  wavereal s, x, *c, *v;
  c=new wavereal[np];
  v=new wavereal[np];

  Rate = f;
  n = a.N;
  imax = int(a.N * Rate / a.Rate + 0.5); 

  if ( N != imax )  Resize(imax); 

// calculate constant filter part c(k) = -(-1)^k/k!/(np-k-1)!
  for (int i = 0; i < np; i++) {
    int m = 1;
    for (int j = 0; j < np; j++)  if (j != i) m *= (i - j);
    c[i] = 1./double(m);
  }

  for (int i = 0; i < N; i++)
  { 
    x = i*a.Rate/Rate;        

    i1 = int(x);
    x = x - i1 + np2 - 1;

// to treat data boundaries we need to calculate critical numbers n1 and n2
    n1 = i1 - np2 + 1;
    n2 = i1 + np2 + 1 - n;

// here we calculate the sum of products like h(k)*a(k), k=0..np-1,
// where h(k) are filter coefficients, a(k) are signal samples
// h(k) = c(k)*prod (x-i), i != k, c(k) = -(-1)^k/k!/(np-k-1)! 

// get signal part multiplied by constant filter coefficients
    if ( n1 >= 0 )
      if ( n2 <= 0 ) 
// regular case - far from boundaries
        for (int j = 0; j < np; j++) v[j] = c[j]*a.data[i1 + j - np2 + 1];
      
      else {
// right border case
        x = x + n2;
        for (int j = 0; j < np; j++) v[j] = c[j]*a.data[n + j - np];
      }  
    else {
// left border case
      x = x + n1;
      for (int j = 0; j < np; j++)  v[j] = c[j]*a.data[j];
    }

// multiply resulted v[j] by x-dependent factors (x-k), if (k != j)
    for (int k = 0; k < np; k++) {
      for (int j = 0; j < np; j++) if (j != k) v[j] *= x; 
      x -= 1.; }

    s = 0.;
    for (int j = 0; j < np; j++)  s += v[j]; 
        
    data[i] = s;
  }

  delete [] c;
  delete [] v;
}

/******************************************************************************
 * copies data from data array of the object WaveData a to *this
 * object data array. Procedure starts from the element "a_pos" in
 * the source array which is copied to the element "pos" in the
 * destination array. "length" is the number of data elements to copy.
 * By default "length"=0, which means copy all source data or if destination
 * array is shorter, copy data until the destination is full.
 *****************************************************************************/
 void WaveData::cpf(const WaveData &a,
                                int length=0, int a_pos=0, int pos=0)
{ 
  if (length == 0 ) 
    length = ((N - pos) < (a.N - a_pos))? (N - pos) : (a.N - a_pos);
  if( length > (N - pos) ) length = N - pos; 
  if( length > (a.N - a_pos) ) length = a.N - a_pos; 

  for (int i = 0; i < length; i++)
    data[i + pos] = a.data[i + a_pos];

  Rate = a.Rate;
}

/******************************************************************************
 * Adds data from data array of the object WaveData a to *this
 * object data array. Procedure starts from the element "a_pos" in
 * the source array which is added starting from the element "pos" in the
 * destination array. "length" is the number of data elements to add.
 * By default "length"=0, which means add all source data or if destination
 * array is shorter, add data up to the end of destination.
 *****************************************************************************/
inline void WaveData::add( const WaveData &a,
                           int length = 0, int a_pos = 0, int pos = 0 )
{
  if (Rate != a.Rate) {
    cout << " Error: attempt to add data with different sample rate.n";
    return;
  }

  if (length == 0 )
    length = ((N - pos) < (a.N - a_pos))? (N - pos) : (a.N - a_pos);
  if( length > (N - pos) ) length = N - pos; 
  if( length > (a.N - a_pos) ) length = a.N - a_pos; 

  for (int i = 0; i < length; i++)
    data[i + pos] += a.data[i + a_pos];
}


/******************************************************************************
 * Substracts data array of the object WaveData a from *this
 * object data array. Procedure starts from the element "a_pos" in
 * the source array which is subtracted starting from the element "pos" in
 * the destination array. "length" is number of data elements to subtract.
 * By default "length"=0, which means subtract all source data or if the 
 * destination array is shorter, subtract data up to the end of destination.
 ******************************************************************************/
inline void WaveData::sub(const WaveData &a, int length = 0,
                           int a_pos = 0, int pos = 0 )
{
  if (Rate != a.Rate) {
    cout << " Error: attempt to subtract data with different sample rate.n";
    return;
  }

  if ( length == 0 )
    length = ((N - pos) < (a.N - a_pos))? (N - pos) : (a.N - a_pos);
  if( length > (N - pos) ) length = N - pos; 
  if( length > (a.N - a_pos) ) length = a.N - a_pos; 

  for (int i = 0; i < length; i++)
    data[i + pos] -= a.data[i + a_pos];
}


/*****************************************************************
 * Calculates Fourier Transform for real signal using
 * Fast Fourier Transform (FFT) algorithm. Packs resulting 
 * complex data into original array of real numbers.
 * Calls wavefft() function which is capable to deal with any
 * number of data points. FFT(1) means forward transformation,
 * which is default, FFT(-1) means inverse transformation.
 *****************************************************************/
 void WaveData::FFT(int direction=1) 
{ 
  double *a, *b;
  int n2=N/2;

  a=new double[N];
  b=new double[N];

  switch (direction)
  { case 1:

// direct transform
    
    for (int i=0; i<N; i++) { a[i]=data[i]; b[i]=0.;}

    wavefft(a, b, N, N, N, -1);

// pack complex numbers to real array
    for (int i=0;i<n2;i++)
    { data[2*i]=a[i]/N;
      data[2*i+1]=b[i]/N; }

// data[1] is not occupied because imag(F[0]) is always zero and we
// store in data[1] the value of F[N/2] which is pure real for even "N"
      data[1]=a[n2]/N;

// in the case of odd number of data points we store imag(F[N/2])
// in the last element of array data[N]
      if ((N&1) == 1) data[N-1]=b[n2]/N;

      break;

    case -1:
// inverse transform

// unpack complex numbers from real array
      for (int i=1;i<n2;i++)
      { a[i]=data[2*i];
        b[i]=data[2*i+1];
        a[N-i]=data[2*i];
        b[N-i]=-data[2*i+1];
       }

      a[0]=data[0];
      b[0]=0.;

      if ((N&1) == 1)
        { a[n2]=data[1]; b[n2]=data[N-1]; }  // for odd n
      else
        { a[n2]=data[1]; b[n2]=0.; }

      wavefft(a, b, N, N, N, 1);             // call FFT for inverse tranform

      for (int i=0; i<N; i++)  data[i]=a[i]; // copy result from array "a"
  }

  delete [] b;
  delete [] a;
}

/*************************************************************************
 * Stack generates WaveData *this by stacking data from WaveData td.
 * Input data are devided on subsets with with samples "length"
 * then they are added together. Average over the all subsets is saved in
 * *this. 
 *************************************************************************/
inline wavereal WaveData::Stack(const WaveData &td, int length)
{
  wavereal ss, s0, s2;
  Rate  = td.Rate;
  int k = td.N/length;
  int n = k*length;

  if (k == 0) {
    cout <<" Stack() error: data length too short to contain n"
         << length << " samplesn";
    return 0.;
  }

  if (N != length) Resize(length);

// sum (stack) all k periods of frequency f to produce 1 cycle
  s0 = 0.;
  for (int i = 0; i < length; i++) {
    ss = 0.;
    for (int j = i; j < n; j += length) ss += td.data[j];
    data[i] = ss/k;
    s0 += ss;
  }
  s0 /= (k*length);

// remove constant displacement (zero frequency component) 
  s2 = 0.;
  for (int i = 0; i < length; i++) {
    data[i] -= s0;
    s2 += data[i]*data[i];
  }
   s2 /= length;

   return s2;        // return stacked signal power (energy/N)
}

/*************************************************************************
 * Another version of Stack:
 * Input data (starting from sample "start") are devided on "k" subsets 
 * with sections "length"
 * Average over the all sections is saved in *this. 
 *************************************************************************/
inline wavereal WaveData::Stack(const WaveData &td, 
				int length, 
				int start){
  wavereal avr, rms;
  Rate  = td.Rate;

  if(start+length > td.N) length = td.N-start;

  int k = (N<1) ? 0 : length/N;
  if (k == 0) {
    cout <<" Stack() error: data length too short to contain n"
         << length << " samplesn";
    return 0.;
  }

  *this = 0.;
  for (int i = 0; i<k; i++) add(td, N, start+i*N);
  *this *= 1./k;
  getStatistics(avr,rms);                                
  *this -= avr;

  return rms*rms;        // return stacked signal power
}

inline wavereal WaveData::Stack(const WaveData &td, wavereal window)
{ 
// Stack generates WaveData *this by stacking data from WaveData td.
// Input data are devided on subsets with with samples "length"
// then they are added together. Average over the all subsets is saved in
// *this. 
return this->Stack(td, int(td.Rate * window)); }

#ifdef _USE_DMT
WaveData& WaveData::operator=(const TSeries &ts)
{
   Interval ti = ts.getTStep();
   wavereal Tsample = ti.GetSecs();

   int n=ts.getNSample();
   if(n != N) Resize(n);

   if ( Tsample > 0. )
      Rate = 1./Tsample;
   else {
      cout <<" Invalid sampling interval = 0 sec.\n";
   }

   TSeries r(ts.getStartTime(), ti, ts.getNSample());
   r = ts;
   float *vec_ref;
   vec_ref= (float*)(r.refData());

   for ( int i=0; i<n; i++ ) data[i]=wavereal(vec_ref[i]);
   return *this;
}

WaveData& WaveData::operator=(const TSeries *ts)
{ return *this=*ts; }

#endif // use DMT

#ifdef _USE_ROOT

// Stream a WaveData object. Required by ROOT system for
// reading and writing the object to file.
 void WaveData::Streamer(TBuffer &b)
{
  int n;
  if (b.IsReading()) {
    Version_t v = b.ReadVersion();
    Wavelet::Streamer(b);
    b >> n;
    if (n != N) Resize(n);

//cout <<" Reading data array of N="<<N<<" elements."<<"n";

    b.ReadFastArray(data,N);
    b >> Rate;
  } else {
    b.WriteVersion(WaveData::IsA());
    Wavelet::Streamer(b);
    b << N;

//cout <<" Writing data array of N="<<N<<" elements."<<"n";

    b.WriteFastArray(data,N);
    b << Rate;
  }
}


#endif