/*
    SPDX-FileCopyrightText: 2014-2017 Max Planck Society.
    All rights reserved.

    SPDX-License-Identifier: BSD-3-Clause
*/

/**
 * @file
 * @date      2014-2017
 * @copyright Max Planck Society
 *
 * @author    Edgar D. Klenske <edgar.klenske@tuebingen.mpg.de>
 * @author    Stephan Wenninger <stephan.wenninger@tuebingen.mpg.de>
 * @author    Raffi Enficiaud <raffi.enficiaud@tuebingen.mpg.de>
 *
 * @brief     Provides mathematical tools needed for the Gaussian process toolbox.
 */

#include "math_tools.h"
#include <stdexcept>
#include <cmath>
#include <cstdint>

namespace math_tools
{

Eigen::MatrixXd squareDistance(const Eigen::MatrixXd &a, const Eigen::MatrixXd &b)
{
    int aRows = a.rows(); // bRows must be identical to aRows
    int aCols = a.cols();
    int bCols = b.cols();

    Eigen::MatrixXd am(aRows, aCols);    // mean-corrected a
    Eigen::MatrixXd bm(aRows, bCols);    // mean-corrected b
    // final result, aCols x bCols
    Eigen::MatrixXd result(aCols, bCols);

    Eigen::VectorXd mean(aRows);

    /*  If the two Matrices have the same address, it means the function was
        called from the overloaded version, and thus the mean only has to be
        computed once.
     */
    if (&a == &b)    // Same address?
    {
        mean = a.rowwise().mean();
        am = a.colwise() - mean;
        bm = am;
    }
    else
    {
        if (aRows != b.rows())
        {
            throw std::runtime_error("Matrix dimension incorrect.");
        }

        mean = static_cast<double>(aCols) / (aCols + bCols) * a.rowwise().mean() +
               static_cast<double>(bCols) / (bCols + aCols) * b.rowwise().mean();

        // The mean of the two Matrices is subtracted beforehand, because the
        // squared error is independent of the mean and this makes the squares
        // smaller.
        am = a.colwise() - mean;
        bm = b.colwise() - mean;
    }

    // The square distance calculation (a - b)^2 is calculated as a^2 - 2*ab * b^2
    // (using the binomial formula) because of numerical stability.

    // fast version
    return
        ((am.array().square().colwise().sum().transpose().rowwise().replicate(bCols).matrix()
          + bm.array().square().colwise().sum().colwise().replicate(aCols).matrix())
         - 2 * (am.transpose()) * bm).array().max(0.0);

    /* // verbose version
    Eigen::MatrixXd a_square =
        am.array().square().colwise().sum().transpose().rowwise().replicate(bCols);

    Eigen::MatrixXd b_square =
        bm.array().square().colwise().sum().colwise().replicate(aCols);

    Eigen::MatrixXd twoab = 2 * (am.transpose()) * bm;

    return ((a_square.matrix() + b_square.matrix()) - twoab).array().max(0);
    */
}

Eigen::MatrixXd squareDistance(const Eigen::MatrixXd &a)
{
    return squareDistance(a, a);
}

Eigen::MatrixXd generate_uniform_random_matrix_0_1(
    const size_t n,
    const size_t m)
{
    Eigen::MatrixXd result = Eigen::MatrixXd(n, m);
    result.setRandom(); // here we get uniform random values between -1 and 1
    Eigen::MatrixXd temp = result.array() + 1; // shift to positive
    result = temp / 2.0; // divide to obtain 0/1 interval
    // prevent numerical problems by enforcing a particular interval
    result = result.array().max(1e-10); // eliminate too small values
    result = result.array().min(1.0); // eliminiate too large values
    return result;
}

Eigen::MatrixXd box_muller(const Eigen::VectorXd &vRand)
{
    size_t n = vRand.rows();
    size_t m = n / 2; // Box-Muller transforms pairs of numbers

    Eigen::ArrayXd rand1 = vRand.head(m);
    Eigen::ArrayXd rand2 = vRand.tail(m);

    /* Implemented according to
     * http://en.wikipedia.org/wiki/Box%E2%80%93Muller_transform
     */

    // enforce interval to avoid numerical issues
    rand1 = rand1.max(1e-10);
    rand1 = rand1.min(1.0);

    rand1 = -2 * rand1.log();
    rand1 = rand1.sqrt(); // this is an amplitude

    rand2 = rand2 * 2 * M_PI; // this is a random angle in the complex plane

    Eigen::MatrixXd result(2 * m, 1);
    Eigen::MatrixXd res1 = (rand1 * rand2.cos()).matrix(); // first elements
    Eigen::MatrixXd res2 = (rand1 * rand2.sin()).matrix(); // second elements
    result << res1, res2; // combine the pairs into one vector

    return result;
}

Eigen::MatrixXd generate_normal_random_matrix(const size_t n, const size_t m)
{
    // if n*m is odd, we need one random number extra!
    // therefore, we have to round up here.
    size_t N = static_cast<size_t>(std::ceil(n * m / 2.0));

    Eigen::MatrixXd result(2 * N, 1);
    // push random samples through the Box-Muller transform
    result = box_muller(generate_uniform_random_matrix_0_1(2 * N, 1));
    result.conservativeResize(n, m);
    return result;
}

std::pair<Eigen::VectorXd, Eigen::VectorXd> compute_spectrum(Eigen::VectorXd &data, int N)
{

    int N_data = data.rows();

    if (N < N_data)
    {
        N = N_data;
    }
    N = static_cast<int>(std::pow(2, std::ceil(std::log(N) / std::log(2)))); // map to nearest power of 2

    Eigen::VectorXd padded_data = Eigen::VectorXd::Zero(N);
    padded_data.head(N_data) = data;

    Eigen::FFT<double> fft;

    // initialize the double vector from Eigen vector. This works by initializing
    // with two pointers: 1) the first element of the data, 2) the last element of the data
    std::vector<double> vec_data(padded_data.data(), padded_data.data() + padded_data.rows() * padded_data.cols());
    std::vector<std::complex<double> > vec_result;
    fft.fwd(vec_result, vec_data); // this is the forward-FFT, from time domain to Fourier domain

    // map back to Eigen vector with the address of the first element and the size
    Eigen::VectorXcd result = Eigen::Map<Eigen::VectorXcd>(&vec_result[0], vec_result.size());

    // the low_index is the lowest useful frequency, depending on the number of actual datapoints
    int low_index = static_cast<int>(std::ceil(static_cast<double>(N) / static_cast<double>(N_data)));

    // Guard: if there are too few actual data points relative to FFT size, the segment would have
    // zero or negative size (low_index >= N/2), which causes an Eigen assertion failure.
    int segment_size = N / 2 - low_index + 1;
    if (segment_size <= 0)
    {
        return std::make_pair(Eigen::VectorXd(), Eigen::VectorXd());
    }

    // prepare amplitudes and frequencies, don't return frequencies introduced by padding
    Eigen::VectorXd spectrum = result.segment(low_index, segment_size).array().abs().pow(2);
    Eigen::VectorXd frequencies = Eigen::VectorXd::LinSpaced(segment_size, low_index, N / 2);
    frequencies /= N;

    return std::make_pair(spectrum, frequencies);
}

Eigen::VectorXd hamming_window(int N)
{
    double alpha = 0.54;
    double beta = 0.46;

    Eigen::VectorXd range = Eigen::VectorXd::LinSpaced(N, 0, 1);
    Eigen::VectorXd window = alpha - beta * (2 * M_PI * range.array()).cos();

    return window;
}

double stdandard_deviation(Eigen::VectorXd &input)
{
    Eigen::ArrayXd centered = input.array() - input.array().mean();
    return std::sqrt(centered.pow(2).sum() / (centered.size() - 1));
}


}  // namespace math_tools





