/* -*- c++ -*- */
/*
 * gr-satnogs: SatNOGS GNU Radio Out-Of-Tree Module
 *
 *  Copyright (C) 2016, Libre Space Foundation <http://librespacefoundation.org/>
 *
 *  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 3 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, see <http://www.gnu.org/licenses/>.
 */

#ifdef HAVE_CONFIG_H
#include "config.h"
#endif

#include <gnuradio/io_signature.h>
#include "doppler_correction_cc_impl.h"
#include <satnogs/log.h>
#include <volk/volk.h>

namespace gr {
namespace satnogs {

doppler_correction_cc::sptr
doppler_correction_cc::make(double target_freq,
                            double offset,
                            double sampling_rate,
                            size_t corrections_per_sec)
{
  return gnuradio::get_initial_sptr(
           new doppler_correction_cc_impl(target_freq, offset,
                                          sampling_rate,
                                          corrections_per_sec));
}

/*
 * The private constructor
 */
doppler_correction_cc_impl::doppler_correction_cc_impl(
  double target_freq,
  double offset,
  double sampling_rate,
  size_t corrections_per_sec) :
  gr::sync_block("doppler_correction_cc",
                 gr::io_signature::make(1, 1, sizeof(gr_complex)),
                 gr::io_signature::make(1, 1, sizeof(gr_complex))),
  d_target_freq(target_freq),
  d_offset(offset),
  d_samp_rate(sampling_rate),
  d_update_period(sampling_rate / corrections_per_sec),
  d_est_thrhld(7),
  d_corrections_per_sec(corrections_per_sec),
  d_nco(),
  /* A 3-rd order polynomial curve fitting is more than enough */
  d_doppler_fit_engine(3),
  d_freq_diff(offset),
  d_have_est(false),
  d_freq_est_num(0),
  d_corrections(0),
  d_corrected_samples(0)
{
  message_port_register_in(pmt::mp("freq"));
  message_port_register_in(pmt::mp("reset"));

  /*
   * NOTE:
   * Set the maximum number of samples to be equivalent of half a second.
   * With this way we are sure that at least one frequency message
   * per second will be processed.
   *
   * This is taken into consideration due to the fact that the work()
   * and the input message handler are NOT reentrant.
   */
  set_max_noutput_items(d_samp_rate / 2.0);
  set_alignment(8);

  set_msg_handler(
    pmt::mp("freq"),
    boost::bind(&doppler_correction_cc_impl::new_freq, this, _1));
  set_msg_handler(
    pmt::mp("reset"),
    boost::bind(&doppler_correction_cc_impl::reset, this, _1));

  d_nco.set_freq((2 * M_PI * (-d_freq_diff)) / d_samp_rate);
  /* Allocate the buffer that will hold the predicted frequency differences */
  d_predicted_freqs = new double[d_corrections_per_sec];

  /* Allocate aligned memory for the NCO */
  d_nco_buff = (gr_complex *) volk_malloc(
                 d_update_period * sizeof(gr_complex), 32);
  if (!d_nco_buff) {
    throw std::runtime_error("Could not allocate NCO memory");
  }
}

void
doppler_correction_cc_impl::new_freq(pmt::pmt_t msg)
{
  boost::mutex::scoped_lock lock(d_mutex);
  double new_freq;
  new_freq = pmt::to_double(msg);
  d_freq_diff = new_freq - (d_target_freq - d_offset);
  if (!d_have_est) {
    d_freq_est_num++;
    d_doppler_freqs.push_back(
      freq_drift(nitems_written(0), d_freq_diff));
    if (d_freq_est_num > d_est_thrhld - 1) {
      d_doppler_fit_engine.fit(d_doppler_freqs);
      d_doppler_fit_engine.predict_freqs(d_predicted_freqs,
                                         d_corrections_per_sec,
                                         d_update_period);
      d_have_est = true;
    }
  }
  else {
    d_doppler_freqs.pop_front();
    d_doppler_freqs.push_back(
      freq_drift(nitems_written(0), d_freq_diff));

    /* Fit the doppler drift based on the new estimated frequency */
    d_doppler_fit_engine.fit(d_doppler_freqs);
    /* Predict the frequency differences for the near future */
    d_doppler_fit_engine.predict_freqs(d_predicted_freqs,
                                       d_corrections_per_sec,
                                       d_update_period);
    d_corrections = 0;
  }
}

void
doppler_correction_cc_impl::reset(pmt::pmt_t msg)
{
  boost::mutex::scoped_lock lock(d_mutex);
  d_doppler_freqs.clear();
  d_freq_est_num = 0;
  d_corrections = 0;
  d_have_est = false;
}

/*
 * Our virtual destructor.
 */
doppler_correction_cc_impl::~doppler_correction_cc_impl()
{
  delete[] d_predicted_freqs;
  volk_free(d_nco_buff);
}

int
doppler_correction_cc_impl::work(int noutput_items,
                                 gr_vector_const_void_star &input_items,
                                 gr_vector_void_star &output_items)
{
  const gr_complex *in = (const gr_complex *) input_items[0];
  gr_complex *out = (gr_complex *) output_items[0];
  int produced = 0;
  size_t cnt;

  /*
   * If we do not have an estimation yet, just copy the input to the output.
   * Otherwise perform Doppler correction, using the fitted curve indicating
   * the frequency drift.
   */
  if (d_have_est) {
    while (produced < noutput_items) {
      /*
       * If no samples have been corrected from the current correction step
       * compute and store the NCO buffer with the corresponding frequency
       */
      if (d_corrected_samples == 0) {
        d_nco.set_freq(
          2 * M_PI * (-d_predicted_freqs[d_corrections]) / d_samp_rate);
        d_nco.sincos(d_nco_buff, d_update_period, 1.0);
        d_corrections++;

        /*
         * The doppler estimation may fail/delay. In such a case the block
         * should continue using the predicted frequencies
         */
        if (d_corrections == d_corrections_per_sec) {
          d_doppler_fit_engine.predict_freqs(d_predicted_freqs,
                                             d_corrections_per_sec,
                                             d_update_period);
          d_corrections = 0;
        }
      }

      cnt = std::min(d_update_period - d_corrected_samples,
                     (size_t)(noutput_items - produced));
      /* Perform the doppler shift correction */
      volk_32fc_x2_multiply_32fc(out + produced, in + produced,
                                 d_nco_buff + d_corrected_samples, cnt);

      /* Make the proper advances */
      produced += (int) cnt;
      d_corrected_samples += cnt;

      if (d_corrected_samples == d_update_period) {
        d_corrected_samples = 0;
      }
    }
  }
  else {
    memcpy(out, in, noutput_items * sizeof(gr_complex));
  }

  return noutput_items;
}

} /* namespace satnogs */
} /* namespace gr */