/*!
 * \file
 * \brief Filter design functions
 * \author Tony Ottosson
 *
 * -------------------------------------------------------------------------
 *
 * IT++ - C++ library of mathematical, signal processing, speech processing,
 *        and communications classes and functions
 *
 * Copyright (C) 1995-2008  (see AUTHORS file for a list of contributors)
 *
 * This program 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 2 of the License, or
 * (at your option) any later version.
 *
 * This program 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
 * along with this program; if not, write to the Free Software
 * Foundation, Inc., 51 Franklin St, Fifth Floor, Boston, MA 02110-1301 USA
 *
 * -------------------------------------------------------------------------
 */

#include <itpp/signal/filter_design.h>
#include <itpp/signal/poly.h>
#include <itpp/signal/filter.h>
#include <itpp/signal/transforms.h>
#include <itpp/base/math/elem_math.h>
#include <itpp/base/algebra/ls_solve.h>
#include <itpp/base/matfunc.h>
#include <itpp/base/specmat.h>
#include <itpp/base/math/trig_hyp.h>
#include <itpp/base/converters.h>


namespace itpp {


  void polystab(const vec &a, vec &out)
  {
    cvec r;
    roots(a, r);

    for (int i=0; i<r.size(); i++) {
      if (abs(r(i)) > 1)
	r(i) = std::complex<double>(1.0)/conj(r(i));
    }
    out = real(std::complex<double>(a(0)) * poly(r));
  }

  void polystab(const cvec &a, cvec &out)
  {
    cvec r;
    roots(a, r);

    for (int i=0; i<r.size(); i++) {
      if (abs(r(i)) > 1)
	r(i) = std::complex<double>(1.0)/conj(r(i));
    }
    out = a(0) * poly(r);
  }


  // ----------------------- freqz() ---------------------------------------------------------
  void freqz(const cvec &b, const cvec &a, const int N, cvec &h, vec &w)
  {
    w = pi*linspace(0, N-1, N)/double(N);

    cvec ha, hb;
    hb = fft( b, 2*N );
    hb = hb(0, N-1);

    ha = fft( a, 2*N );
    ha = ha(0, N-1);

    h = elem_div(hb, ha);
  }

  cvec freqz(const cvec &b, const cvec &a, const int N)
  {
    cvec h;
    vec w;

    freqz(b, a, N, h, w);

    return h;
  }


  cvec freqz(const cvec &b, const cvec &a, const vec &w)
  {
    int la = a.size(), lb = b.size(), k = std::max(la, lb);

    cvec h, ha, hb;

    // Evaluate the nominator and denominator at the given frequencies
    hb = polyval( zero_pad(b, k), to_cvec(cos(w), sin(w)) );
    ha = polyval( zero_pad(a, k), to_cvec(cos(w), sin(w)) );

    h = elem_div(hb, ha);

    return h;
  }

  void freqz(const vec &b, const vec &a, const int N, cvec &h, vec &w)
  {
    w = pi*linspace(0, N-1, N)/double(N);

    cvec ha, hb;
    hb = fft_real( b, 2*N );
    hb = hb(0, N-1);

    ha = fft_real( a, 2*N );
    ha = ha(0, N-1);

    h = elem_div(hb, ha);
  }

  cvec freqz(const vec &b, const vec &a, const int N)
  {
    cvec h;
    vec w;

    freqz(b, a, N, h, w);

    return h;
  }


  cvec freqz(const vec &b, const vec &a, const vec &w)
  {
    int la = a.size(), lb = b.size(), k = std::max(la, lb);

    cvec h, ha, hb;

    // Evaluate the nominator and denominator at the given frequencies
    hb = polyval( zero_pad(b, k), to_cvec(cos(w), sin(w)) );
    ha = polyval( zero_pad(a, k), to_cvec(cos(w), sin(w)) );

    h = elem_div(hb, ha);

    return h;
  }



  void filter_design_autocorrelation(const int N, const vec &f, const vec &m, vec &R)
  {
    it_assert(f.size() == m.size(), "filter_design_autocorrelation: size of f and m vectors does not agree");
    int N_f = f.size();

    it_assert(f(0) == 0.0, "filter_design_autocorrelation: first frequency must be 0.0");
    it_assert(f(N_f-1) == 1.0, "filter_design_autocorrelation: last frequency must be 1.0");

    // interpolate frequency-response
    int N_fft = 512;
    vec m_interp(N_fft+1);
    // unused variable:
    // double df_interp = 1.0/double(N_fft);

    m_interp(0) = m(0);
    double inc;

    int jstart = 0, jstop;

    for (int i=0; i<N_f-1; i++) {
      // calculate number of points to the next frequency
      jstop = floor_i( f(i+1)*(N_fft+1) ) - 1;
      //std::cout << "jstart=" << jstart << "jstop=" << jstop << std::endl;

      for (int j=jstart; j<=jstop; j++) {
	inc = double(j-jstart)/double(jstop-jstart);
	m_interp(j) = m(i)*(1-inc) + m(i+1)*inc;
      }
      jstart = jstop+1;
    }

    vec S = sqr(concat( m_interp, reverse(m_interp(2,N_fft)) )); // create a complete frequency response with also negative frequencies

    R = ifft_real(to_cvec(S)); // calculate correlation

    R = R.left(N);
  }


  // Calculate the AR coefficients of order \c n of the ARMA-process defined by the autocorrelation R
  // using the deternined modified Yule-Walker method
  // maxlag determines the size of the system to solve N>= n.
  // If N>m then the system is overdetermined and a least squares solution is used.
  // as a rule of thumb use N = 4*n
  void modified_yule_walker(const int m, const int n, const int N, const vec &R, vec &a)
  {
    it_assert(m>0, "modified_yule_walker: m must be > 0");
    it_assert(n>0, "modified_yule_walker: n must be > 0");
    it_assert(N <= R.size(), "modified_yule_walker: autocorrelation function too short");

    // create the modified Yule-Walker equations Rm * a = - rh
    // see eq. (3.7.1) in Stoica and Moses, Introduction to spectral analysis
    int M = N - m - 1;

    mat Rm;
    if(m-n+1 < 0)
      Rm= toeplitz( R(m, m+M-1), reverse(concat( R(1,std::abs(m-n+1)), R(0,m) ) ) );
    else
      Rm= toeplitz( R(m, m+M-1), reverse(R(m-n+1,m)) );


    vec rh = - R(m+1, m+M);

    // solve overdetermined system
    a = backslash(Rm, rh);

    // prepend a_0 = 1
    a = concat(1.0, a);

    // stabilize polynomial
    a = polystab(a);
  }



  void arma_estimator(const int m, const int n, const vec &R, vec &b, vec &a)
  {
    it_assert(m>0, "arma_estimator: m must be > 0");
    it_assert(n>0, "arma_estimator: n must be > 0");
    it_assert(2*(m+n)<=R.size(), "arma_estimator: autocorrelation function too short");


    // windowing the autocorrelation
    int N = 2*(m+n);
    vec Rwindow = elem_mult(R.left(N), 0.54 + 0.46*cos( pi*linspace(0.0, double(N-1), N)/double(N-1) ) ); // Hamming windowing

    // calculate the AR part using the overdetmined Yule-Walker equations
    modified_yule_walker(m, n, N, Rwindow, a);

    // --------------- Calculate MA part --------------------------------------
    // use method in ref [2] section VII.
    vec r_causal = Rwindow;
    r_causal(0) *= 0.5;

    vec h_inv_a = filter(1, a, concat(1.0, zeros(N-1))); // see eq (50) of [2]
    mat H_inv_a = toeplitz(h_inv_a, concat(1.0, zeros(m)));

    vec b_causal = backslash(H_inv_a, r_causal);

    // calculate the double-sided spectrum
    int N_fft = 256;
    vec H = 2.0*real(elem_div(fft_real(b_causal, N_fft), fft_real(a, N_fft))); // calculate spectrum

    // Do weighting and windowing in cepstrum domain
    cvec cepstrum = log(to_cvec(H));
    cvec q = ifft(cepstrum);

    // keep only causal part of spectrum (windowing)
    q.set_subvector(N_fft/2, N_fft-1, zeros_c(N_fft/2) );
    q(0) *= 0.5;

    cvec h = ifft(exp(fft(q))); // convert back to frequency domain, from cepstrum and do inverse transform to calculate impulse response
    b = real(backslash(to_cmat(H_inv_a), h(0,N-1))); // use Shank's method to calculate b coefficients
  }


  void yulewalk(const int N, const vec &f, const vec &m, vec &b, vec &a)
  {
    it_assert(f.size() == m.size(), "yulewalk: size of f and m vectors does not agree");
    int N_f = f.size();

    it_assert(f(0) == 0.0, "yulewalk: first frequency must be 0.0");
    it_assert(f(N_f-1) == 1.0, "yulewalk: last frequency must be 1.0");


    vec R;
    filter_design_autocorrelation(4*N, f, m, R);

    arma_estimator(N, N, R, b, a);
  }


} // namespace itpp