/*

  Copyright (C) 2019 Gonzalo José Carracedo Carballal

  This program is free software: you can redistribute it and/or modify
  it under the terms of the GNU Lesser General Public License as
  published by the Free Software Foundation, version 3.

  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 Lesser General Public License for more details.

  You should have received a copy of the GNU Lesser General Public
  License along with this program.  If not, see
  <http://www.gnu.org/licenses/>

*/

#define SU_LOG_DOMAIN "audio-inspector"

#include <sigutils/sigutils.h>
#include <sigutils/agc.h>
#include <sigutils/pll.h>
#include <sigutils/sampling.h>
#include <sigutils/iir.h>
#include <sigutils/clock.h>

#include <analyzer/version.h>

#include "inspector/interface.h"
#include "inspector/params.h"
#include "inspector/inspector.h"

#include <string.h>

#define SUSCAN_AUDIO_INSPECTOR_SAMPLE_RATE 44100

struct suscan_audio_inspector_params {
  struct suscan_inspector_gc_params gc;
  struct suscan_inspector_audio_params audio;
};

/*
 * Spike durations measured in symbol times
 * SUSCAN_AUDIO_INSPECTOR_FAST_RISE_FRAC has been doubled to reduce phase noise
 * induced by the non-linearity of the AGC
 */
#define SUSCAN_AUDIO_INSPECTOR_FAST_RISE_FRAC   (100 * 3.9062e-1)
#define SUSCAN_AUDIO_INSPECTOR_FAST_FALL_FRAC   (2 * SUSCAN_AUDIO_INSPECTOR_FAST_RISE_FRAC)
#define SUSCAN_AUDIO_INSPECTOR_SLOW_RISE_FRAC   (10 * SUSCAN_AUDIO_INSPECTOR_FAST_RISE_FRAC)
#define SUSCAN_AUDIO_INSPECTOR_SLOW_FALL_FRAC   (10 * SUSCAN_AUDIO_INSPECTOR_FAST_FALL_FRAC)
#define SUSCAN_AUDIO_INSPECTOR_HANG_MAX_FRAC    (SUSCAN_AUDIO_INSPECTOR_FAST_RISE_FRAC * 5)
#define SUSCAN_AUDIO_INSPECTOR_DELAY_LINE_FRAC  (SUSCAN_AUDIO_INSPECTOR_FAST_RISE_FRAC * 10)
#define SUSCAN_AUDIO_INSPECTOR_MAG_HISTORY_FRAC (SUSCAN_AUDIO_INSPECTOR_FAST_RISE_FRAC * 10)

#define SUSCAN_AUDIO_INSPECTOR_MIN_TS             1e-1

#define SUSCAN_AUDIO_INSPECTOR_BRICKWALL_LEN      200
#define SUSCAN_AUDIO_AM_LPF_SECONDS               .1
#define SUSCAN_AUDIO_AM_ATTENUATION               .25
#define SUSCAN_AUDIO_AM_CARRIER_AVERAGING_SECONDS .2
#define SUSCAN_AUDIO_RAW_GAIN                     1e3
#define SUSCAN_AUDIO_SQUELCH_AVG_SECONDS          1e-2

struct suscan_audio_inspector {
  struct suscan_inspector_sampling_info samp_info;
  struct suscan_audio_inspector_params req_params;
  struct suscan_audio_inspector_params cur_params;

  /* Blocks */
  su_agc_t  agc;          /* AGC, for AM-like modulations */
  su_iir_filt_t filt;     /* Input filter */

  su_iir_filt_t fm_lpf;   /* FM low pass filter */

  su_pll_t pll;           /* Carrier tracking PLL */
  su_ncqo_t lo;           /* Oscillator */
  su_sampler_t sampler;   /* Fixed rate sampler */

  SUFLOAT beta;           /* Coefficient for single pole IIR filter */
  SUCOMPLEX last;         /* Last processed sample (for quad demod) */

  SUFLOAT sql_alpha;      /* Coefficient for SPLPF used for power comparison */
  SUFLOAT fm_power_low;   /* Measure of FM power (LPF) */
  SUFLOAT fm_power_chan;  /* Measure of FM power (full channel) */

  SUFLOAT am_power_carr;  /* Measure of AM power carrier */

  SUFLOAT ssb_power_chan; /* Measure of SSB power */
};

SUPRIVATE void
suscan_audio_inspector_params_initialize(
    struct suscan_audio_inspector_params *params,
    const struct suscan_inspector_sampling_info *sinfo)
{
  memset(params, 0, sizeof(struct suscan_audio_inspector_params));

  params->gc.gc_ctrl = SUSCAN_INSPECTOR_GAIN_CONTROL_AUTOMATIC;
  params->gc.gc_gain = 1;
  params->gc.gc_ts   = 1;

  params->audio.sample_rate = SUSCAN_AUDIO_INSPECTOR_SAMPLE_RATE;
  params->audio.demod       = SUSCAN_INSPECTOR_AUDIO_DEMOD_DISABLED;
  params->audio.cutoff      = SUSCAN_AUDIO_INSPECTOR_SAMPLE_RATE / 2;
}

SUPRIVATE void
suscan_audio_inspector_destroy(struct suscan_audio_inspector *insp)
{
  su_iir_filt_finalize(&insp->filt);

  su_iir_filt_finalize(&insp->fm_lpf);

  su_pll_finalize(&insp->pll);

  su_agc_finalize(&insp->agc);

  su_sampler_finalize(&insp->sampler);

  free(insp);
}

SUINLINE SUBOOL
suscan_audio_inspector_update_agc(
  struct suscan_audio_inspector *self,
  SUFLOAT ts)
{
  su_agc_t new_agc = su_agc_INITIALIZER;
  struct su_agc_params agc_params = su_agc_params_INITIALIZER;
  SUBOOL ok = SU_FALSE;
  SUFLOAT tau;

  if (ts < SUSCAN_AUDIO_INSPECTOR_MIN_TS)
    ts = SUSCAN_AUDIO_INSPECTOR_MIN_TS;

  tau = ts / self->samp_info.bw;

  agc_params.fast_rise_t = tau * SUSCAN_AUDIO_INSPECTOR_FAST_RISE_FRAC;
  agc_params.fast_fall_t = tau * SUSCAN_AUDIO_INSPECTOR_FAST_FALL_FRAC;
  agc_params.slow_rise_t = tau * SUSCAN_AUDIO_INSPECTOR_SLOW_RISE_FRAC;
  agc_params.slow_fall_t = tau * SUSCAN_AUDIO_INSPECTOR_SLOW_FALL_FRAC;
  agc_params.hang_max    = tau * SUSCAN_AUDIO_INSPECTOR_HANG_MAX_FRAC;

  /* TODO: Check whether these sizes are too big */
  agc_params.delay_line_size  = tau * SUSCAN_AUDIO_INSPECTOR_DELAY_LINE_FRAC;
  agc_params.mag_history_size = tau * SUSCAN_AUDIO_INSPECTOR_MAG_HISTORY_FRAC;

  SU_TRY(su_agc_init(&new_agc, &agc_params));

  su_agc_finalize(&self->agc);
  self->agc = new_agc;
  
  ok = SU_TRUE;

done:
  return ok;
}

SUPRIVATE struct suscan_audio_inspector *
suscan_audio_inspector_new(const struct suscan_inspector_sampling_info *sinfo)
{
  struct suscan_audio_inspector *new = NULL;
  SUFLOAT bw;

  SU_ALLOCATE_FAIL(new, struct suscan_audio_inspector);

  new->samp_info = *sinfo;
  bw = sinfo->bw;

  suscan_audio_inspector_params_initialize(&new->cur_params, sinfo);

  SU_TRY_FAIL(suscan_audio_inspector_update_agc(new, new->cur_params.gc.gc_ts));

  /* PLL init, this is an experimental optimum that works rather well for AM */
  su_pll_init(&new->pll, 0, .005f * bw);

  /* Filter init */
  SU_TRY_FAIL(su_iir_bwlpf_init(
      &new->filt,
      5,
      SU_ABS2NORM_FREQ(sinfo->equiv_fs, new->cur_params.audio.cutoff)));

  /* NCQO init, used to sideband adjustment */
  su_ncqo_init(&new->lo, .5 * bw);

  /* FM filters initialization, used for FM squelch */
  SU_TRY_FAIL(su_iir_bwlpf_init(&new->fm_lpf, 5, .5 * bw));

  /* One second time constant, used to remove AM carrier */
  new->beta = SU_SPLPF_ALPHA(
      SUSCAN_AUDIO_AM_CARRIER_AVERAGING_SECONDS * sinfo->equiv_fs);

  new->sql_alpha = SU_SPLPF_ALPHA(
      SUSCAN_AUDIO_SQUELCH_AVG_SECONDS * sinfo->equiv_fs);

  return new;

fail:
  if (new != NULL) {
    suscan_audio_inspector_destroy(new);
    new = NULL;
  }

  return new;
}

/************************** API implementation *******************************/
void *
suscan_audio_inspector_open(const struct suscan_inspector_sampling_info *s)
{
  return suscan_audio_inspector_new(s);
}

SUBOOL
suscan_audio_inspector_get_config(void *private, suscan_config_t *config)
{
  struct suscan_audio_inspector *insp =
      (struct suscan_audio_inspector *) private;

  SU_TRYCATCH(
      suscan_inspector_gc_params_save(&insp->cur_params.gc, config),
      return SU_FALSE);

  SU_TRYCATCH(
      suscan_inspector_audio_params_save(&insp->cur_params.audio, config),
      return SU_FALSE);

  return SU_TRUE;
}

SUBOOL
suscan_audio_inspector_parse_config(void *private, const suscan_config_t *config)
{
  struct suscan_audio_inspector *insp = (struct suscan_audio_inspector *) private;


  SU_TRYCATCH(
      suscan_inspector_gc_params_parse(&insp->req_params.gc, config),
      return SU_FALSE);

  SU_TRYCATCH(
      suscan_inspector_audio_params_parse(&insp->req_params.audio, config),
      return SU_FALSE);

  return SU_TRUE;
}

/* Called inside inspector mutex */
void
suscan_audio_inspector_new_bandwidth(void *private, SUFREQ bw)
{
  struct suscan_audio_inspector *insp =
        (struct suscan_audio_inspector *) private;
  SUFLOAT fs = insp->samp_info.equiv_fs;

  /* Initialize oscillator */
  su_ncqo_set_freq(&insp->lo, SU_ABS2NORM_FREQ(fs, .5 * bw));
}

SUINLINE SUBOOL
suscan_audio_inspector_needs_agc_update(
  const struct suscan_audio_inspector *self)
{
  if (self->req_params.gc.gc_ctrl != self->cur_params.gc.gc_ctrl) {
    /* Changed to automatic */
    if (self->req_params.gc.gc_ctrl == SUSCAN_INSPECTOR_GAIN_CONTROL_AUTOMATIC)
      return SU_TRUE;
  }

  if (self->req_params.gc.gc_ctrl == SUSCAN_INSPECTOR_GAIN_CONTROL_AUTOMATIC) {
    SUFLOAT cur_ts = self->cur_params.gc.gc_ts;
    SUFLOAT req_ts = self->req_params.gc.gc_ts;

    if (req_ts < SUSCAN_AUDIO_INSPECTOR_MIN_TS)
      req_ts = SUSCAN_AUDIO_INSPECTOR_MIN_TS;

    if (!sufeq(cur_ts, req_ts, .5 * SUSCAN_AUDIO_INSPECTOR_MIN_TS))
      return SU_TRUE;
  }

  return SU_FALSE;
}

/* Called inside inspector mutex */
void
suscan_audio_inspector_commit_config(void *private)
{
  struct suscan_audio_inspector *self =
      (struct suscan_audio_inspector *) private;
  su_iir_filt_t filt;
  SUBOOL filt_initialized;
  SUFLOAT fs = self->samp_info.equiv_fs;

  self->last  = 0;

  /* Configure AGC */
  if (suscan_audio_inspector_needs_agc_update(self)) {
    if (!suscan_audio_inspector_update_agc(
      self,
      self->req_params.gc.gc_ts)) {
      SU_ERROR("Failed to update audio AGC\n");
    }
  }

  /* Configure demod */
  if (self->req_params.audio.demod != SUSCAN_INSPECTOR_AUDIO_DEMOD_DISABLED) {
    switch (self->req_params.audio.demod) {
      case SUSCAN_INSPECTOR_AUDIO_DEMOD_FM:
        /*
         * FM transmissions are rather wide (up to 15 kHz), and pilot tones
         * are at around 19 kHz. We prefer to attenuate the pilot tone instead
         * of providing high stability at lower cutoff frequencies.
         */
        filt_initialized = su_iir_bwlpf_init(
            &filt,
            5,
            SU_ABS2NORM_FREQ(fs, self->req_params.audio.cutoff));

        break;

      case SUSCAN_INSPECTOR_AUDIO_DEMOD_RAW:
      case SUSCAN_INSPECTOR_AUDIO_DEMOD_AM:
        /*
         * AM transmissions are around 12 kHz (6 per sideband). In this case,
         * it is okay to provide a filter with lower Q but stable at lower
         * cutoff frequencies.
         */
        filt_initialized = su_iir_bwlpf_init(
            &filt,
            3,
            SU_ABS2NORM_FREQ(fs, self->req_params.audio.cutoff));

        break;

      case SUSCAN_INSPECTOR_AUDIO_DEMOD_LSB:
      case SUSCAN_INSPECTOR_AUDIO_DEMOD_USB:
        /*
         * SSB transmissions are usually very narrow, and require great
         * selectivity, even at low cutoffs. We sacrifice CPU in order
         * to attain this.
         */
        filt_initialized = su_iir_brickwall_lp_init(
            &filt,
            SUSCAN_AUDIO_INSPECTOR_BRICKWALL_LEN,
            SU_ABS2NORM_FREQ(fs, self->req_params.audio.cutoff));
        break;

      default:
        filt_initialized = SU_FALSE;
        break;
    }

    if (!filt_initialized) {
      SU_ERROR("No memory left to initialize audio filter");
    } else {
      su_iir_filt_finalize(&self->filt);
      self->filt = filt;
    }
  }

  /* Set sampling info */
  if (self->req_params.audio.sample_rate > 0) {
    su_sampler_finalize(&self->sampler);
    su_sampler_init(
        &self->sampler,
        SU_ABS2NORM_BAUD(fs, self->req_params.audio.sample_rate));
  }

  self->cur_params = self->req_params;
}

SUSDIFF
suscan_audio_inspector_feed(
    void *private,
    suscan_inspector_t *insp,
    const SUCOMPLEX *x,
    SUSCOUNT count)
{
  SUCOMPLEX last, det_x, output = 0;
  SUSCOUNT i;
  SUCOMPLEX lo;
  SUCOMPLEX ylp;
  struct suscan_audio_inspector *self =
      (struct suscan_audio_inspector *) private;
  SUBOOL raw = self->cur_params.audio.demod == SUSCAN_INSPECTOR_AUDIO_DEMOD_RAW;

  if (self->cur_params.audio.demod == SUSCAN_INSPECTOR_AUDIO_DEMOD_DISABLED)
    return count;

  last = self->last;

  for (i = 0; i < count && suscan_inspector_sampler_buf_avail(insp) > 0; ++i) {
    det_x = SU_C_VALID(x[i]) ? x[i] : 0;

    if (!raw
        && self->cur_params.audio.squelch
        && (self->cur_params.audio.demod
              == SUSCAN_INSPECTOR_AUDIO_DEMOD_LSB
            || self->cur_params.audio.demod
                == SUSCAN_INSPECTOR_AUDIO_DEMOD_USB)) {
      SU_SPLPF_FEED(
          self->ssb_power_chan,
          SU_C_REAL(det_x * SU_C_CONJ(det_x)),
          self->sql_alpha);

      if (self->ssb_power_chan
          * suscan_inspector_get_equiv_fs(insp)
            / suscan_inspector_get_equiv_bw(insp)
            < self->cur_params.audio.squelch_level)
        det_x = 0;
    }

    if (!raw) {
      /* Perform gain control */
      switch (self->cur_params.gc.gc_ctrl) {
        case SUSCAN_INSPECTOR_GAIN_CONTROL_MANUAL:
          det_x = 2 * self->cur_params.gc.gc_gain * det_x;
          break;

        case SUSCAN_INSPECTOR_GAIN_CONTROL_AUTOMATIC:
          det_x  = 2 * su_agc_feed(&self->agc, det_x);
          break;
      }
    }

    switch (self->cur_params.audio.demod) {
      case SUSCAN_INSPECTOR_AUDIO_DEMOD_FM:
        output = SU_C_ARG(det_x * SU_C_CONJ(last)) / M_PI;
        last   = det_x;

        /*
         * FM squelch compares the output in lower frequencies
         * with the output of the full channel.
         */
        if (self->cur_params.audio.squelch) {
          ylp = su_iir_filt_feed(&self->fm_lpf, output);

          SU_SPLPF_FEED(
              self->fm_power_low,
              SU_C_REAL(ylp * SU_C_CONJ(ylp)),
              self->sql_alpha);

          SU_SPLPF_FEED(
              self->fm_power_chan,
              SU_C_REAL(output * SU_C_CONJ(output)),
              self->sql_alpha);

          if (!sufreleq(self->fm_power_chan, self->fm_power_low, 1e-1))
            output = 0;
        }

        break;

      case SUSCAN_INSPECTOR_AUDIO_DEMOD_AM:
        /* Synchronous detection */
        output  = su_pll_track(&self->pll, det_x);

        /* Carrier removal */
        SU_SPLPF_FEED(last, output, self->beta);

        if (self->cur_params.audio.squelch) {
          SU_SPLPF_FEED(
              self->am_power_carr,
              SU_C_REAL(last * SU_C_CONJ(last)),
              self->sql_alpha);

          if (self->am_power_carr < self->cur_params.audio.squelch_level)
            output = 0;
          else
            output -= last;
        } else {
          output -= last;
        }

        /* Volume attenuation */
        output *= SUSCAN_AUDIO_AM_ATTENUATION;
        break;

      case SUSCAN_INSPECTOR_AUDIO_DEMOD_USB:
        lo = su_ncqo_read(&self->lo);
        output  = det_x * lo;
        break;

      case SUSCAN_INSPECTOR_AUDIO_DEMOD_LSB:
        lo = su_ncqo_read(&self->lo);
        output = det_x * SU_C_CONJ(lo);
        break;

      /* Pass thru */
      case SUSCAN_INSPECTOR_AUDIO_DEMOD_RAW:
        output = SUSCAN_AUDIO_RAW_GAIN * det_x;

      default:
        break;
    }

    output *= self->cur_params.audio.volume;

    output = su_iir_filt_feed(&self->filt, output);

    if (su_sampler_feed(&self->sampler, &output)) {
      suscan_inspector_push_sample(insp, output * .75);
      if (suscan_inspector_get_output_length(insp) == insp->sample_msg_watermark) {
        ++i; /* Important! We have consumed this sample too! */
        break;
      }
    }
  }


  self->last = last;

  return i;
}

void
suscan_audio_inspector_close(void *private)
{
  suscan_audio_inspector_destroy((struct suscan_audio_inspector *) private);
}

SUPRIVATE struct suscan_inspector_interface iface = {
    .name = "audio",
    .desc = "Analog audio",
    .open = suscan_audio_inspector_open,
    .get_config = suscan_audio_inspector_get_config,
    .parse_config = suscan_audio_inspector_parse_config,
    .commit_config = suscan_audio_inspector_commit_config,
    .new_bandwidth = suscan_audio_inspector_new_bandwidth,
    .feed = suscan_audio_inspector_feed,
    .close = suscan_audio_inspector_close
};

SUBOOL
suscan_audio_inspector_register(void)
{
  SU_TRYCATCH(
      iface.cfgdesc = suscan_config_desc_new_ex(
          "audio-params-desc-" SUSCAN_VERSION_STRING),
      return SU_FALSE);

  /* Add all configuration parameters */
  SU_TRYCATCH(suscan_config_desc_add_gc_params(iface.cfgdesc), return SU_FALSE);
  SU_TRYCATCH(suscan_config_desc_add_audio_params(iface.cfgdesc), return SU_FALSE);

  SU_TRYCATCH(suscan_config_desc_register(iface.cfgdesc), return SU_FALSE);

  /* Register inspector interface */
  SU_TRYCATCH(suscan_inspector_interface_register(&iface), return SU_FALSE);

  return SU_TRUE;
}
