/*---------------------------------------------------------------------------*
 *                                   IT++			             *
 *---------------------------------------------------------------------------*
 * Copyright (c) 1995-2004 by Tony Ottosson, Thomas Eriksson, Pål Frenger,   *
 * Tobias Ringström, and Jonas Samuelsson.                                   *
 *                                                                           *
 * Permission to use, copy, modify, and distribute this software and its     *
 * documentation under the terms of the GNU General Public License is hereby *
 * granted. No representations are made about the suitability of this        *
 * software for any purpose. It is provided "as is" without expressed or     *
 * implied warranty. See the GNU General Public License for more details.    *
 *---------------------------------------------------------------------------*/

/*! 
  \file 
  \brief Implementation of signal processing functions.
  \author Tony Ottosson, Thomas Eriksson, Pål Frenger, and Tobias Ringström

  1.8

  2004/09/03 09:27:35
*/

#include "base/elmatfunc.h"
#include "base/matfunc.h"
#include "base/stat.h"
#include "base/specmat.h"
#include "base/operators.h"
#include "base/sigfun.h"

#include <iostream>

namespace itpp { 

  void xcorr(const vec &x, vec &out, const int max_lag, const string scaleopt) {

    int m, n;
    double s, M_double, coeff_scale = 0.0;
    int M, N;

    M = x.size();
    M_double = double(M);

    if (max_lag==-1) {
      N = x.size();
    } else {
      N = max_lag + 1;
    }

    out.set_size(2*N-1,false);

    it_assert(N<=x.size(),"xcorr: max_lag cannot be as large as, or larger than, the length of x.");

    for (m=0;m<N;m++) {
      s=0.0;
      for (n=0;n<M-m;n++) {
	s+=x[n]*x[n+m];
      }

      if (m==0) { coeff_scale = s; }

      if (scaleopt=="none"){
	out[N+m-1]=s;
	out[N-m-1]=s;
      }

      else if (scaleopt=="biased") {
	out[N+m-1]=s/M_double;
	out[N-m-1]=s/M_double;
      }

      else if (scaleopt=="unbiased") {
	out[N+m-1]=s/double(M - m);
	out[N-m-1]=s/double(M - m);
      }

      else if (scaleopt=="coeff") {
	out[N+m-1]=s/coeff_scale;
	out[N-m-1]=s/coeff_scale;
      }

      else
	it_error("xcorr: Incorrect scaleopt specified.");

    }

  }

  vec xcorr(const vec &x, const int max_lag, const string scaleopt) {
    vec out;
    xcorr(x,out,max_lag,scaleopt);
    return out;
  }

  void xcorr(vec &x, vec &y, vec &out, const int max_lag, const string scaleopt) {

    //If x and y are of different length, the shortest one is zero-padded:
    int d, D;
    int x_size = x.size();
    int y_size = y.size();
    if (x_size>y_size) {
      it_assert(scaleopt=="none","xcorr: scaleopt must be none for different size vectors x and y");
      y.set_size(x_size,true);
      D = x_size - y_size;
      for (d=0; d<D; d++) {
	y(y_size+d) = 0.0;
      }
    } else if (x_size<y_size) {
      it_assert(scaleopt=="none","xcorr: scaleopt must be none for different size vectors x and y");
      x.set_size(y_size,true);
      D = y_size - x_size;      
      for (d=0; d<D; d++) {
	x(x_size+d) = 0.0;
      }
    }

    int m, n;
    double s_plus, s_minus, M_double, xcorr_0, coeff_scale=0.0;
    int M, N;

    M = x.size();
    M_double = double(M);

    if (max_lag==-1) {
      N=x.size();
    } else {
      N=max_lag+1;
    }

    out.set_size(2*N-1,false);

    it_assert(N<=x.size(),"max_lag cannot be as large as, or larger than, the length of x.");

    if (scaleopt=="coeff") {
      vec x_scale = xcorr(x,0,"none");
      vec y_scale = xcorr(y,0,"none");
      coeff_scale = std::sqrt(x_scale(0)) * std::sqrt(y_scale(0));
    }

    for (m=0;m<N;m++) {
      s_plus=0;
      s_minus=0;
      for (n=0;n<M-m;n++) {
	s_minus+=x[n]*y[n+m];
	s_plus+=x[n+m]*y[n];
      }

      if (m==0) { xcorr_0 = s_plus; }

      if (scaleopt=="none") {
	out[N+m-1]=s_plus;
	out[N-m-1]=s_minus;
      }
      else if (scaleopt=="biased"){
	out[N+m-1]=s_plus/M_double;
	out[N-m-1]=s_minus/M_double; 
      }
      else if (scaleopt=="unbiased"){
	out[N+m-1]=s_plus/double(M-m);
	out[N-m-1]=s_minus/double(M-m); 	
      }
      else if (scaleopt=="coeff") {
	out[N+m-1]=s_plus/coeff_scale;
	out[N-m-1]=s_minus/coeff_scale;
      }
      else
	it_error("Incorrect scaleopt specified.");
    }

    //Set the original size of x and y:
    x.set_size(x_size,true);
    y.set_size(y_size,true);

  }

  vec xcorr(vec &x, vec &y, const int max_lag, const string scaleopt) {
    vec out;
    xcorr(x, y, out, max_lag, scaleopt);
    return out;
  }

  mat cov(const mat &m, bool is_zero_mean)
  {
    int d=m.cols(), n=m.rows(), i, j, k;
    mat R(d, d), m2(n, d);
    vec tmp;

    if (!is_zero_mean) {
      // Compute and remove mean
      for (i=0; i<d; i++) {
	tmp = m.get_col(i);
	m2.set_col(i, tmp - mean(tmp));
      }
    }

    // Calc corr matrix
    R = 0.0;
    for (i=0; i<d; i++)
      for (j=0; j<=i; j++) {
	for (k=0; k<n; k++)
	  R(i,j) += m2(k,i) * m2(k,j);
	R(j,i) = R(i,j); // When i=j this is unnecassary work 
      }
    R /= n;

    return R;
  }

  vec spectrum(const vec &v, int nfft, int noverlap)
  {
    it_assert1(pow2(needed_bits(nfft))==nfft,
	       "nfft must be a power of two in spectrum()!");
    
    vec P(nfft/2+1), w(nfft), wd(nfft);

    P = 0.0;
    w = hanning(nfft);
    double w_energy = nfft==1 ? 1 : (nfft+1)*.375; // Hanning energy
    
    if (nfft > v.size()) {
      P = sqr(abs( fft(to_cvec(elem_mult(zero_pad(v, nfft), w)))(0, nfft/2) ));
      P /= w_energy;
    }
    else {
      int k = (v.size()-noverlap) / (nfft-noverlap), idx = 0;
      for (int i=0; i<k; i++) {
	wd = elem_mult(v(idx, idx+nfft-1), w);
	P += sqr(abs( fft(to_cvec(wd))(0, nfft/2) ));
	idx += nfft - noverlap;
      }
      P /= k * w_energy;
    }
    
    P.set_size(nfft/2+1, true);
    return P;
  }

  vec spectrum(const vec &v, const vec &w, int noverlap)
  {
    int nfft = w.size();
    it_assert1(pow2(needed_bits(nfft))==nfft,
	       "The window size must be a power of two in spectrum()!");
    
    vec P(nfft/2+1), wd(nfft);

    P = 0.0;
    double w_energy = energy(w);
    
    if (nfft > v.size()) {
      P = sqr(abs( fft(to_cvec(elem_mult(zero_pad(v, nfft), w)))(0, nfft/2) ));
      P /= w_energy;
    }
    else {
      int k = (v.size()-noverlap) / (nfft-noverlap), idx = 0;
      for (int i=0; i<k; i++) {
	wd = elem_mult(v(idx, idx+nfft-1), w);
	P += sqr(abs( fft(to_cvec(wd))(0, nfft/2) ));
	idx += nfft - noverlap;
      }
      P /= k * w_energy;
    }
    
    P.set_size(nfft/2+1, true);
    return P;
  }

  vec filter_spectrum(const vec &a, int nfft)
  {
    vec s = sqr(abs(fft(to_cvec(a), nfft)));
    s.set_size(nfft/2+1, true);
    return s;
  }

  vec filter_spectrum(const vec &a, const vec &b, int nfft)
  {
    vec s = sqr(abs(elem_div(fft(to_cvec(a), nfft), fft(to_cvec(b), nfft))));
    s.set_size(nfft/2+1, true);
    return s;
  }

} //namespace itpp