/**
 * @file core/data/mfe_impl.hpp
 * @author Omar Shrit
 *
 * Implementation file for MFE (Mel-Frequency Energy) filters.
 *
 * mlpack is free software; you may redistribute it and/or modify it under the
 * terms of the 3-clause BSD license.  You should have received a copy of the
 * 3-clause BSD license along with mlpack.  If not, see
 * http://www.opensource.org/licenses/BSD-3-Clause for more information.
 */
#ifndef MLPACK_CORE_DATA_MFE_IMPL_HPP
#define MLPACK_CORE_DATA_MFE_IMPL_HPP

#include "mfcc.hpp"

namespace mlpack {

template<typename eT>
inline arma::Col<eT> HammingWindow(size_t len)
{
  return (0.54 - (0.46 * arma::cos(2.0 * M_PI * arma::linspace<arma::Col<eT>>(0,
    len - 1, len) / (len - 1))));
}

template<typename eT>
inline eT HzToMel(eT hz)
{
  return 2595 * std::log10(1.0 + hz / 700.0);
}

template<typename eT>
inline eT MelToHz(eT mel)
{
  return 700.0 * (std::pow(10.0, mel / 2595.0) - 1.0);
}

template<typename eT>
inline arma::Mat<eT> MelFilterbank(size_t numFilters,
                                   size_t nFFT,
                                   size_t sampleRate,
                                   double lowFreq,
                                   double highFreq)
{
  size_t numBins = nFFT / 2 + 1;
  eT melLow = HzToMel(static_cast<eT>(lowFreq));
  eT melHigh = HzToMel(static_cast<eT>(highFreq));

  size_t numPoints = numFilters + 2;
  arma::Col<eT> melPoints = arma::linspace<arma::Col<eT>>(melLow, melHigh,
      numPoints);

  arma::Col<eT> hzPoints = 700.0 * (arma::exp10(melPoints / 2595.0) - 1.0);

  arma::Col<eT> binFreqHz = arma::regspace<arma::Col<eT>>(0, numBins - 1)
      * sampleRate / nFFT;

  arma::Mat<eT> melFilterbank(numFilters, numBins);

  for (size_t i = 0; i < numFilters; ++i)
  {
    eT left   = hzPoints(i);
    eT center = hzPoints(i + 1);
    eT right  = hzPoints(i + 2);

    if (highFreq <= lowFreq)
      Log::Fatal << "MFE(): highFreq must be strictly greater than lowFreq."
                 << std::endl;

    arma::Col<eT> goUp = (binFreqHz - left) / (center - left);
    arma::Col<eT> goDown = (right - binFreqHz) / (right - center);

    melFilterbank.row(i) = arma::clamp(arma::min(goUp, goDown), 0, 1).t();
  }

  return melFilterbank;
}

template<typename eT>
inline void MFE(const arma::Mat<eT>& inputSignal,
                arma::Mat<eT>& mfe,
                size_t sampleRate,
                size_t numMelFilters,
                double windowLength,
                double windowStep,
                size_t nFFT,
                double lowFreq,
                double highFreq,
                const typename std::enable_if_t<
                    std::is_floating_point<eT>::value>*)
{
  if (highFreq == 0.0f)
    highFreq = sampleRate / 2.0f;

  size_t lengthInSamples = static_cast<size_t>(windowLength * sampleRate
      / 1000.0f);
  size_t stepsInSamples = static_cast<size_t>(windowStep * sampleRate
      / 1000.0f);

  if (nFFT == 0)
    nFFT = NextPowerOf2(lengthInSamples);

  if (nFFT < lengthInSamples)
    Log::Fatal << "MFE(): nFFT cannot be lower than window length in samples. "
               << "nFFT needs to be >= windowLength x sampleRate." << std::endl;

  size_t numBins = nFFT / 2 + 1;

  arma::Mat<eT> filterBanks = MelFilterbank<eT>(numMelFilters, nFFT, sampleRate,
      lowFreq, highFreq);

  size_t totalWindows = inputSignal.n_cols * ((inputSignal.n_rows
      - lengthInSamples) / stepsInSamples + 1);

  // Needs to be matching windows length in SlidingWindow
  arma::Col<eT> hWindow = HammingWindow<eT>(lengthInSamples);

  mfe.set_size(numMelFilters, totalWindows);

  size_t colOffset = 0;

  for (size_t i = 0; i < inputSignal.n_cols; ++i)
  {
    arma::Mat<eT> slidingWindows;
    arma::Mat<eT> power;

    SlidingWindow(inputSignal.col(i), slidingWindows, hWindow, lengthInSamples,
        stepsInSamples, nFFT);

    PowerSpectrum(slidingWindows, power, nFFT);

    // Adding a small value 1e-10 so we do not take the log of 0, in case the
    // multiplication results in zero.
    mfe.cols(colOffset, colOffset + power.n_cols - 1) =
        arma::log((filterBanks * power) + 1e-10);

    colOffset += power.n_cols;
  }
}

template<typename eT>
inline void MFE(const arma::Mat<eT>& inputSignal,
                arma::Mat<eT>& mfe,
                size_t sampleRate,
                size_t numMelFilters,
                double windowLength,
                double windowStep,
                size_t nFFT,
                double lowFreq,
                double highFreq,
                const typename std::enable_if_t<
                    !std::is_floating_point<eT>::value>*)
{
  static_assert(std::is_floating_point<eT>::value,
      "MFE(): input matrix must have a floating-point element type "
      "(float or double). Integer matrices are not supported because "
      "FFT, log, and other operations require floating-point "
      "arithmetic.");
}

inline size_t NextPowerOf2(size_t n)
{
  size_t p = 1;
  while (p < n && p < (SIZE_MAX / 2))
    p <<= 1;
  return p;
}

template<typename eT>
inline void PowerSpectrum(const arma::Mat<eT>& windows, arma::Mat<eT>& power,
    size_t nFFT)
{
  size_t numBins = nFFT / 2 + 1;
  // spectrum is (nFFT x numWindows).
  arma::Mat<std::complex<eT>> spectrum = arma::fft(windows, nFFT);

  power = square(abs(spectrum.rows(0, numBins - 1)));
}

template<typename MatType, typename eT>
inline void SlidingWindow(const MatType& signal,
                          arma::Mat<eT>& windows,
                          arma::Col<eT>& hWindow,
                          size_t windowLength,
                          size_t windowStep,
                          size_t nFFT)
{
  if (signal.n_elem < windowLength)
  {
    windows.zeros(nFFT, 1);
    windows.col(0).subvec(0, signal.n_elem - 1) = signal;
  }
  else
  {
    size_t numWindows = (signal.n_elem - windowLength) / windowStep + 1;
    windows.set_size(nFFT, numWindows);
    windows.rows(windowLength, nFFT - 1).zeros();

    for (size_t i = 0; i < numWindows; ++i)
    {
      size_t start = i * windowStep;
      windows.col(i).subvec(0, windowLength - 1) =
          signal.subvec(start, start + windowLength - 1) % hWindow;
    }
  }
}

} // namespace mlpack

#endif
