/*
   Copyright (C) 2015
   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>

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;

    std::vector<complexf> rx_fft_in(N);
    std::vector<complexf> rx_fft_out(N);
    std::vector<complexf> tx_fft_in(N);
    std::vector<complexf> tx_fft_out(N);
    std::vector<complexf> ifft_in(N);
    std::vector<complexf> 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;
}

std::pair<std::vector<complexf>, std::vector<complexf> > AlignSample::get_samples(
                size_t len, size_t rx_delay)
{
    std::pair<std::vector<complexf>, std::vector<complexf> > rval;

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

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

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

        std::copy(m_txsamples.begin(),
                m_txsamples.begin() + len,
                std::back_inserter(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;
}