/*
   Copyright (C) 2017
   Matthias P. Braendli, matthias.braendli@mpb.li

    http://opendigitalradio.org
 */
/*
   This file is part of ODR-DPD.

   ODR-DPD 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 3 of the
   License, or (at your option) any later version.

   ODR-DPD 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 ODR-DPD.  If not, see <http://www.gnu.org/licenses/>.
 */

#include "AlignSample.hpp"
#include <cstring>
#include <exception>
#include <algorithm>
#include <numeric>
#include <boost/math/tools/minima.hpp>

namespace fftw {
    class plan {
        public:
            template <class T>
            plan(size_t N, T& in, T& out, int direction)
            {
                m_plan = fftwf_plan_dft_1d(N,
                        reinterpret_cast<fftwf_complex*>(&in.front()),
                        reinterpret_cast<fftwf_complex*>(&out.front()),
                        FFTW_FORWARD, FFTW_MEASURE);
            }

            plan(const plan& other) = delete;
            const plan& operator=(const plan& other) = delete;

            void execute(void) {
                fftwf_execute(m_plan);
            }

            ~plan() {
                fftwf_destroy_plan(m_plan);
            }
        private:
            fftwf_plan m_plan;
    };
}

void AlignSample::push_tx_samples(complexf* samps, size_t len, double first_sample_time)
{
    std::lock_guard<std::mutex> lock(m_mutex);
    std::copy(samps, samps + len, std::back_inserter(m_txsamples));

    if (m_first_tx_sample_time == 0) {
        m_first_tx_sample_time = first_sample_time;
    }
}

void AlignSample::push_rx_samples(complexf* samps, size_t len, double first_sample_time)
{
    std::lock_guard<std::mutex> lock(m_mutex);
    if (m_first_rx_sample_time == 0 and first_sample_time != 0) {
        std::copy(samps, samps + len, std::back_inserter(m_rxsamples));

        m_first_rx_sample_time = first_sample_time;
        m_num_rx_samples_dropped = 0;
    }
    else if (m_first_rx_sample_time != 0) {
        // We have previously received samples with valid timestamp
        double delta = m_rx_sample_time(m_rxsamples.size()) - first_sample_time;
        if (std::fabs(delta) > 1e-3) {
            fprintf(stderr, "RX sample time %f expected %f. Delta of %f(%zu)."
                    " Resetting RX\n",
                m_rx_sample_time(m_rxsamples.size()),
                first_sample_time,
                delta,
                (size_t)std::fabs(delta * (double)samplerate));

            m_first_rx_sample_time = 0;
            m_num_rx_samples_dropped = 0;
            m_rxsamples.clear();
        }

        std::copy(samps, samps + len, std::back_inserter(m_rxsamples));
    }
    else if (first_sample_time == 0) {
        MDEBUG("RX timestamp missing\n");
        throw std::runtime_error("RX timestamp missing");
    }
}

void AlignSample::reset_rx()
{
    std::lock_guard<std::mutex> lock(m_mutex);
    m_first_rx_sample_time = 0;
    m_num_rx_samples_dropped = 0;
    m_rxsamples.clear();
}

bool AlignSample::ready(size_t min_samples)
{
    std::lock_guard<std::mutex> lock(m_mutex);
    return align() and m_rxsamples.size() > min_samples and m_txsamples.size() > min_samples;
}

CorrelationResult AlignSample::crosscorrelate(size_t len)
{
    double rx_ts = 0;
    double tx_ts = 0;

    /* The size of the FFT is twice as long as the desired
     * correlation length, because the FFT does a circular
     * correlation, and we therefore pad half with zeros
     */
    const size_t N = 2 * len;

    vec_cf rx_fft_in(N);
    vec_cf rx_fft_out(N);
    vec_cf tx_fft_in(N);
    vec_cf tx_fft_out(N);
    vec_cf ifft_in(N);
    vec_cf ifft_out(N);

    fftw::plan rx_fft_plan(N,
            rx_fft_in,
            rx_fft_out,
            FFTW_FORWARD);

    fftw::plan tx_fft_plan(N,
            tx_fft_in,
            tx_fft_out,
            FFTW_FORWARD);

    fftw::plan ifft_plan(N,
            ifft_in,
            ifft_out,
            FFTW_BACKWARD);

    // Do a quick copy, so as to free the mutex
    {
        std::lock_guard<std::mutex> lock(m_mutex);

        if (!align() or
                m_rxsamples.size() < len or
                m_txsamples.size() < len) {
            CorrelationResult result(0);
            return result;
        }

        std::copy(m_rxsamples.begin(), m_rxsamples.begin() + len, rx_fft_in.begin());
        std::copy(m_txsamples.begin(), m_txsamples.begin() + len, tx_fft_in.begin());
        // the other half of the buffers are set to 0

        m_drop_rx_samples(len);
        m_drop_tx_samples(len);

        rx_ts = m_rx_sample_time();
        tx_ts = m_tx_sample_time();
    }

    CorrelationResult result(len);
    result.rx_timestamp = rx_ts;
    result.tx_timestamp = tx_ts;

    // Calculate power
    for (size_t i = 0; i < len; i++) {
        result.rx_power += std::norm(rx_fft_in[i]);
    }
    result.rx_power = std::sqrt(result.rx_power);

    for (size_t i = 0; i < len; i++) {
        result.tx_power += std::norm(tx_fft_in[i]);
    }
    result.tx_power = std::sqrt(result.tx_power);

    // Implement
    // corr(a, b) = ifft(fft(a) * conj(fft(b)))
    // Attention: circular correlation !

    rx_fft_plan.execute();
    tx_fft_plan.execute();

    for (size_t i = 0; i < len; i++) {
        ifft_in[i] = tx_fft_out[i] * std::conj(rx_fft_out[i]);
    }

    ifft_plan.execute();

    for (size_t i = 0; i < len; i++) {
        result.correlation[i] = ifft_out[i];
    }

    return result;
}

void AlignSample::delay_rx_samples(size_t delay)
{
    std::lock_guard<std::mutex> lock(m_mutex);
    if (align() and delay > 0 and m_rxsamples.size() > delay) {
        m_rxsamples.erase(
                m_rxsamples.begin(),
                m_rxsamples.begin() + delay - 1);
    }
}

static std::vector<double> gen_omega(size_t length)
{
    if ((length % 2) == 1) {
        throw std::runtime_error("Needs an even length array.");
    }
    size_t halflength = length/2;

    double factor = 2.0 * M_PI / length;

    std::vector<double> omega(length);
    for (size_t i = 0; i < halflength; i++) {
        omega[i] = factor * i;
    }
    for (size_t i = halflength+1; i < length; i++) {
        omega[i] = factor * ((ssize_t)i - length);
    }

    return omega;
}

/* Find subsample delay in signal versus the reference signal ref and
 * correct it. This assumes that the offset between signal and ref
 * is less than one sample.
 */
static vec_cf align_subsample(vec_cf& signal, vec_cf& ref)
{
    size_t N = signal.size();
    vec_cf fft_sig(N);
    vec_cf rotate_vec(N);
    vec_cf corr_sig(N);

    fftw::plan plan(N, signal, fft_sig, FFTW_FORWARD);
    fftw::plan plan_ifft(N, rotate_vec, corr_sig, FFTW_BACKWARD);
    plan.execute();

    const std::vector<double> omega = gen_omega(N);

    // Calculate the correlation for a lag of tau, which
    // must be between -1 and 1 samples
    auto correlate_point = [&](float tau) {

        // A subsample offset between two signals corresponds, in the frequency
        // domain, to a linearly increasing phase shift, whose slope
        // corresponds to the delay.
        //
        // Here, we build this phase shift in rotate_vec, and multiply it with
        // our signal.
        for (size_t i = 0; i < N; i++) {
            const complexf angle(0, tau * i);
            rotate_vec[i] = std::exp(angle); // e^(j*tau*i)
        }
        // zero-frequency is rotate_vec[0], so rotate_vec[N/2] is the
        // bin corresponding to the [-1, 1, -1, 1, ...] time signal, which
        // is both the maximum positive and negative frequency.
        // I don't remember why we handle it differently.
        rotate_vec[N/2] = cos(M_PI * tau);

        for (size_t i = 0; i < N; i++) {
            rotate_vec[i] *= fft_sig[i];
        }

        plan_ifft.execute(); // corr_sig = IFFT(rotate_vec)

        // Calculate correlation only the real part of the signal:
        // The following implements the elementwise vector operation
        //   corr_real = real(corr_sig) * real(ref)
        // in an awkward way.

        std::vector<float> corr_real(N);
        std::transform(corr_sig.begin(), corr_sig.end(), corr_real.begin(),
                [&](complexf f){ return f.real(); });

        for (size_t i = 0; i < N; i++) {
            complexf f = corr_real[i] * ref[i];
            corr_real[i] = f.real();
        }

        // TODO: replace by the following and verify it's equivalent
        /*
        for (size_t i = 0; i < N; i++) {
            corr_real[i] = corr_sig[i].real() * ref[i].real();
        }
        */

        // TODO why do we only look at the real part? Because it's faster than
        // a complex cross-correlation? Clarify!
        // Compare with
        /*
        for (size_t i = 0; i < N; i++) {
            corr_real[i] = (corr_sig[i] * std::conj(ref[i])).real();
        }
        */

        // The correlation is the sum:
        return std::accumulate(corr_real.begin(), corr_real.end(), 0.0f);
    };

    const float start_arg = -1;
    const float end_arg = 1;

    // Let boost find the tau value with the minimal correlation
    auto tau =
        boost::math::tools::brent_find_minima(
                correlate_point, start_arg, end_arg, 1000).first;

    fprintf(stderr, "Fractional delay is %f\n", tau);

    // Prepare rotate_vec = fft_sig with rotated phase
    for (size_t i = 0; i < N; i++) {
        const complexf angle(0, tau * i);
        rotate_vec[i] = std::exp(angle);
    }
    rotate_vec[N/2] = cos(M_PI * tau);
    for (size_t i = 0; i < N; i++) {
        rotate_vec[i] *= fft_sig[i];
    }

    plan_ifft.execute(); // corr_sig = IFFT(rotate_vec)

    return corr_sig;
}

std::pair<vec_cf, vec_cf> AlignSample::get_samples(
                size_t len)
{
    std::pair<vec_cf, vec_cf> rval;

    std::lock_guard<std::mutex> lock(m_mutex);
    if (align() and
            m_rxsamples.size() > len and
            m_txsamples.size() > len) {

        rval.first.reserve(len);
        rval.second.reserve(len);

        std::copy(m_rxsamples.begin(),
                m_rxsamples.begin() + len,
                std::back_inserter(rval.first));

        std::copy(m_txsamples.begin(),
                m_txsamples.begin() + len,
                std::back_inserter(rval.second));

        rval.first = align_subsample(rval.first, rval.second);
    }

    return rval;
}

void AlignSample::consume(size_t samples)
{
    std::lock_guard<std::mutex> lock(m_mutex);
    if (align() and m_rxsamples.size() > samples and m_txsamples.size() > samples) {
        m_drop_rx_samples(samples);
        m_drop_tx_samples(samples);
    }
}

bool AlignSample::align()
{
    if (std::abs(m_rx_sample_time() - m_tx_sample_time()) < 1e-6) {
        return true;
    }
    else if (m_rx_sample_time() < m_tx_sample_time()) {
        size_t rx_samples_to_skip =
            (m_tx_sample_time() - m_rx_sample_time()) * samplerate;

        if (rx_samples_to_skip > m_rxsamples.size()) {
            return false;
        }

        m_drop_rx_samples(rx_samples_to_skip);
        return true;
    }
    else if (m_rx_sample_time() > m_tx_sample_time()) {
        size_t tx_samples_to_skip =
            (m_rx_sample_time() - m_tx_sample_time()) * samplerate;

        if (tx_samples_to_skip > m_txsamples.size()) {
            return false;
        }

        m_drop_tx_samples(tx_samples_to_skip);
        return true;
    }
    return false;
}