/* GStreamer
 * Copyright (C) 2007 Sebastian Dröge <slomo@circular-chaos.org>
 *           (C) 2015 Wim Taymans <wim.taymans@gmail.com>
 *
 * gstaudioquantize.c: quantizes audio to the target format and optionally
 *                     applies dithering and noise shaping.
 *
 * This library is free software; you can redistribute it and/or
 * modify it under the terms of the GNU Library General Public
 * License as published by the Free Software Foundation; either
 * version 2 of the License, or (at your option) any later version.
 *
 * This library 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
 * Library General Public License for more details.
 *
 * You should have received a copy of the GNU Library General Public
 * License along with this library; if not, write to the
 * Free Software Foundation, Inc., 51 Franklin St, Fifth Floor,
 * Boston, MA 02110-1301, USA.
 */

/* TODO: - Maybe drop 5-pole noise shaping and use coefficients
 *         generated by dmaker
 *         http://shibatch.sf.net
 */
#ifdef HAVE_CONFIG_H
#include "config.h"
#endif

#include <gst/gst.h>
#include <string.h>
#include <math.h>

#include "gstaudiopack.h"
#include "audio-quantize.h"

typedef void (*QuantizeFunc) (GstAudioQuantize * quant, const gpointer src,
    gpointer dst, gint count);

struct _GstAudioQuantize
{
  GstAudioDitherMethod dither;
  GstAudioNoiseShapingMethod ns;
  GstAudioQuantizeFlags flags;
  GstAudioFormat format;
  guint quantizer;
  guint stride;
  guint blocks;

  guint shift;
  guint32 mask, bias;

  /* last random number generated per channel for hifreq TPDF dither */
  gpointer last_random;
  guint32 random_state;
  /* contains the past quantization errors, error[channels][count] */
  guint error_size;
  gpointer error_buf;
  /* buffer with dither values */
  guint dither_size;
  gpointer dither_buf;
  /* noise shaping coefficients */
  gpointer coeffs;
  gint n_coeffs;

  QuantizeFunc quantize;
};

#define ADDSS(res,val) \
        if (val > 0 && res > 0 && G_MAXINT32 - res <= val){             \
          res = G_MAXINT32;                                             \
        } else if (val < 0 && res < 0 && G_MININT32 - res >= val){      \
          res = G_MININT32;                                             \
        } else                                                          \
          res += val;

static void
gst_audio_quantize_quantize_memcpy (GstAudioQuantize * quant,
    const gpointer src, gpointer dst, gint samples)
{
  if (src != dst)
    memcpy (dst, src, samples * sizeof (gint32) * quant->stride);
}

/* Quantize functions for gint32 as intermediate format */
static void
gst_audio_quantize_quantize_int_none_none (GstAudioQuantize * quant,
    const gpointer src, gpointer dst, gint samples)
{
  audio_orc_int_bias (dst, src, quant->bias, ~quant->mask,
      samples * quant->stride);
}

/* 32 bit xorshift PRNG, see https://en.wikipedia.org/wiki/Xorshift */
static inline guint32
gst_fast_random_uint32 (guint32 * state)
{
  guint64 x = *state;
  x ^= x << 13;
  x ^= x >> 17;
  x ^= x << 5;
  return (*state = x);
}

static inline gint32
gst_fast_random_int32 (guint32 * state)
{
  return (gint32) gst_fast_random_uint32 (state);
}

/* Assuming dither == 2^n,
 * returns one of 2^(n+1) possible random values:
 * -dither <= retval < dither */
#define RANDOM_INT_DITHER(state, dither)                                       \
  (- dither + (gst_fast_random_int32 (state) & ((dither << 1) - 1)))

static void
setup_dither_buf (GstAudioQuantize * quant, gint samples)
{
  gboolean need_init = FALSE;
  gint stride = quant->stride;
  gint i, len = samples * stride;
  guint shift = quant->shift;
  guint32 bias;
  gint32 dither, *d;

  if (quant->dither_size < len) {
    quant->dither_size = len;
    quant->dither_buf = g_realloc (quant->dither_buf, len * sizeof (gint32));
    need_init = TRUE;
  }

  bias = quant->bias;
  d = quant->dither_buf;

  switch (quant->dither) {
    case GST_AUDIO_DITHER_NONE:
      if (need_init) {
        for (i = 0; i < len; i++)
          d[i] = 0;
      }
      break;

    case GST_AUDIO_DITHER_RPDF:
      dither = 1 << (shift);
      for (i = 0; i < len; i++)
        d[i] = bias + RANDOM_INT_DITHER (&quant->random_state, dither);
      break;

    case GST_AUDIO_DITHER_TPDF:
      dither = 1 << (shift - 1);
      for (i = 0; i < len; i++)
        d[i] =
            bias + RANDOM_INT_DITHER (&quant->random_state,
            dither) + RANDOM_INT_DITHER (&quant->random_state, dither);
      break;

    case GST_AUDIO_DITHER_TPDF_HF:
    {
      gint32 tmp, *last_random = quant->last_random;

      dither = 1 << (shift - 1);
      for (i = 0; i < len; i++) {
        tmp = RANDOM_INT_DITHER (&quant->random_state, dither);
        d[i] = bias + tmp - last_random[i % stride];
        last_random[i % stride] = tmp;
      }
      break;
    }
  }
}

static void
gst_audio_quantize_quantize_int_dither_none (GstAudioQuantize * quant,
    const gpointer src, gpointer dst, gint samples)
{
  setup_dither_buf (quant, samples);

  audio_orc_int_dither (dst, src, quant->dither_buf, ~quant->mask,
      samples * quant->stride);
}

static void
setup_error_buf (GstAudioQuantize * quant, gint samples, gint extra)
{
  gint stride = quant->stride;
  gint len = (samples + extra) * stride;

  if (quant->error_size < len) {
    quant->error_buf = g_realloc (quant->error_buf, len * sizeof (gint32));
    if (quant->error_size == 0)
      memset ((gint32 *) quant->error_buf, 0, stride * extra * sizeof (gint32));
    quant->error_size = len;
  }
}

static void
gst_audio_quantize_quantize_int_dither_feedback (GstAudioQuantize * quant,
    const gpointer src, gpointer dst, gint samples)
{
  guint32 mask;
  gint i, len, stride;
  const gint32 *s = src;
  gint32 *dith, *d = dst, v, o, *e, err;

  setup_dither_buf (quant, samples);
  setup_error_buf (quant, samples, 1);

  stride = quant->stride;
  len = samples * stride;
  dith = quant->dither_buf;
  e = quant->error_buf;
  mask = ~quant->mask;

  for (i = 0; i < len; i++) {
    o = v = s[i];
    /* add dither */
    err = dith[i];
    /* remove error */
    err -= e[i];
    ADDSS (v, err);
    v &= mask;
    /* store new error */
    e[i + stride] = e[i] + (v - o);
    /* store result */
    d[i] = v;
  }
  memmove (e, &e[len], sizeof (gint32) * stride);
}

#define SHIFT 10
#define REDUCE 8
#define RROUND (1<<(REDUCE-1))
#define SREDUCE 2
#define SROUND (1<<(SREDUCE-1))

static void
gst_audio_quantize_quantize_int_dither_noise_shape (GstAudioQuantize * quant,
    const gpointer src, gpointer dst, gint samples)
{
  guint32 mask;
  gint i, j, k, len, stride, nc;
  const gint32 *s = src;
  gint32 *c, *dith, *d = dst, v, o, *e, err;

  nc = quant->n_coeffs;

  setup_dither_buf (quant, samples);
  setup_error_buf (quant, samples, nc);

  stride = quant->stride;
  len = samples * stride;
  dith = quant->dither_buf;
  e = quant->error_buf;
  c = quant->coeffs;
  mask = ~quant->mask;

  for (i = 0; i < len; i++) {
    v = s[i];
    /* combine and remove error */
    err = 0;
    for (j = 0, k = i; j < nc; j++, k += stride)
      err -= e[k] * c[j];
    err = (err + SROUND) >> (SREDUCE);
    ADDSS (v, err);
    o = v;
    /* add dither */
    err = dith[i];
    ADDSS (v, err);
    /* quantize */
    v &= mask;
    /* store new error with reduced precision */
    e[k] = (v - o + RROUND) >> REDUCE;
    /* store result */
    d[i] = v;
  }
  memmove (e, &e[len], sizeof (gint32) * stride * nc);
}

#define MAKE_QUANTIZE_FUNC_NAME(name)                                   \
gst_audio_quantize_quantize_##name

static const QuantizeFunc quantize_funcs[] = {
  (QuantizeFunc) MAKE_QUANTIZE_FUNC_NAME (int_none_none),
  (QuantizeFunc) MAKE_QUANTIZE_FUNC_NAME (int_dither_feedback),
  (QuantizeFunc) MAKE_QUANTIZE_FUNC_NAME (int_dither_noise_shape),
  (QuantizeFunc) MAKE_QUANTIZE_FUNC_NAME (int_dither_noise_shape),
  (QuantizeFunc) MAKE_QUANTIZE_FUNC_NAME (int_dither_noise_shape),
  (QuantizeFunc) MAKE_QUANTIZE_FUNC_NAME (int_dither_none),
  (QuantizeFunc) MAKE_QUANTIZE_FUNC_NAME (int_dither_feedback),
  (QuantizeFunc) MAKE_QUANTIZE_FUNC_NAME (int_dither_noise_shape),
  (QuantizeFunc) MAKE_QUANTIZE_FUNC_NAME (int_dither_noise_shape),
  (QuantizeFunc) MAKE_QUANTIZE_FUNC_NAME (int_dither_noise_shape),
  (QuantizeFunc) MAKE_QUANTIZE_FUNC_NAME (int_dither_none),
  (QuantizeFunc) MAKE_QUANTIZE_FUNC_NAME (int_dither_feedback),
  (QuantizeFunc) MAKE_QUANTIZE_FUNC_NAME (int_dither_noise_shape),
  (QuantizeFunc) MAKE_QUANTIZE_FUNC_NAME (int_dither_noise_shape),
  (QuantizeFunc) MAKE_QUANTIZE_FUNC_NAME (int_dither_noise_shape),
  (QuantizeFunc) MAKE_QUANTIZE_FUNC_NAME (int_dither_none),
  (QuantizeFunc) MAKE_QUANTIZE_FUNC_NAME (int_dither_feedback),
  (QuantizeFunc) MAKE_QUANTIZE_FUNC_NAME (int_dither_noise_shape),
  (QuantizeFunc) MAKE_QUANTIZE_FUNC_NAME (int_dither_noise_shape),
  (QuantizeFunc) MAKE_QUANTIZE_FUNC_NAME (int_dither_noise_shape),
};

/* Same as error feedback but also add 1/2 of the previous error value.
 * This moves the noise a bit more into the higher frequencies. */
static const gdouble ns_simple_coeffs[] = {
  -0.5, 1.0
};

/* Noise shaping coefficients from[1], moves most power of the
 * error noise into inaudible frequency ranges.
 *
 * [1]
 * "Minimally Audible Noise Shaping", Stanley P. Lipshitz,
 * John Vanderkooy, and Robert A. Wannamaker,
 * J. Audio Eng. Soc., Vol. 39, No. 11, November 1991. */

static const gdouble ns_medium_coeffs[] = {
  0.6149, -1.590, 1.959, -2.165, 2.033
};

/* Noise shaping coefficients by David Schleef, moves most power of the
 * error noise into inaudible frequency ranges */
static const gdouble ns_high_coeffs[] = {
  -0.340122, 0.876066, -1.72008, 2.61339, -3.31399, 3.27918, -2.92975, 2.08484,
};


static void
gst_audio_quantize_setup_noise_shaping (GstAudioQuantize * quant)
{
  gint i, n_coeffs = 0;
  gint32 *q;
  const gdouble *coeffs;

  switch (quant->ns) {
    case GST_AUDIO_NOISE_SHAPING_HIGH:
      n_coeffs = 8;
      coeffs = ns_high_coeffs;
      break;

    case GST_AUDIO_NOISE_SHAPING_MEDIUM:
      n_coeffs = 5;
      coeffs = ns_medium_coeffs;
      break;

    case GST_AUDIO_NOISE_SHAPING_SIMPLE:
      n_coeffs = 2;
      coeffs = ns_simple_coeffs;
      break;

    case GST_AUDIO_NOISE_SHAPING_ERROR_FEEDBACK:
      break;

    case GST_AUDIO_NOISE_SHAPING_NONE:
    default:
      break;
  }

  if (n_coeffs) {
    quant->n_coeffs = n_coeffs;
    q = quant->coeffs = g_new0 (gint32, n_coeffs);
    for (i = 0; i < n_coeffs; i++)
      q[i] = floor (coeffs[i] * (1 << SHIFT) + 0.5);
  }
  return;
}

static void
gst_audio_quantize_setup_dither (GstAudioQuantize * quant)
{
  /* Some non-zero number */
  quant->random_state = 0xc2d6038f;

  switch (quant->dither) {
    case GST_AUDIO_DITHER_TPDF_HF:
      quant->last_random = g_new0 (gint32, quant->stride);
      break;
    case GST_AUDIO_DITHER_RPDF:
    case GST_AUDIO_DITHER_TPDF:
      quant->last_random = NULL;
      break;
    case GST_AUDIO_DITHER_NONE:
    default:
      quant->last_random = NULL;
      break;
  }
  return;
}

static void
gst_audio_quantize_setup_quantize_func (GstAudioQuantize * quant)
{
  gint index;

  if (quant->shift == 0) {
    quant->quantize = (QuantizeFunc) MAKE_QUANTIZE_FUNC_NAME (memcpy);
    return;
  }

  index = 5 * quant->dither + quant->ns;
  quant->quantize = quantize_funcs[index];
}

static gint
count_power (guint v)
{
  gint res = 0;
  while (v > 1) {
    res++;
    v >>= 1;
  }
  return res;
}

/**
 * gst_audio_quantize_new: (skip):
 * @dither: a #GstAudioDitherMethod
 * @ns: a #GstAudioNoiseShapingMethod
 * @flags: #GstAudioQuantizeFlags
 * @format: the #GstAudioFormat of the samples
 * @channels: the amount of channels in the samples
 * @quantizer: the quantizer to use
 *
 * Create a new quantizer object with the given parameters.
 *
 * Output samples will be quantized to a multiple of @quantizer. Better
 * performance is achieved when @quantizer is a power of 2.
 *
 * Dithering and noise-shaping can be performed during quantization with
 * the @dither and @ns parameters.
 *
 * Returns: a new #GstAudioQuantize. Free with gst_audio_quantize_free().
 */
GstAudioQuantize *
gst_audio_quantize_new (GstAudioDitherMethod dither,
    GstAudioNoiseShapingMethod ns, GstAudioQuantizeFlags flags,
    GstAudioFormat format, guint channels, guint quantizer)
{
  GstAudioQuantize *quant;

  g_return_val_if_fail (format == GST_AUDIO_FORMAT_S32, NULL);
  g_return_val_if_fail (channels > 0, NULL);

  quant = g_slice_new0 (GstAudioQuantize);
  quant->dither = dither;
  quant->ns = ns;
  quant->flags = flags;
  quant->format = format;
  if (flags & GST_AUDIO_QUANTIZE_FLAG_NON_INTERLEAVED) {
    quant->stride = 1;
    quant->blocks = channels;
  } else {
    quant->stride = channels;
    quant->blocks = 1;
  }
  quant->quantizer = quantizer;

  quant->shift = count_power (quantizer);
  if (quant->shift > 0)
    quant->bias = (1U << (quant->shift - 1));
  else
    quant->bias = 0;
  quant->mask = (1U << quant->shift) - 1;

  gst_audio_quantize_setup_dither (quant);
  gst_audio_quantize_setup_noise_shaping (quant);
  gst_audio_quantize_setup_quantize_func (quant);

  return quant;
}

/**
 * gst_audio_quantize_free:
 * @quant: a #GstAudioQuantize
 *
 * Free a #GstAudioQuantize.
 */
void
gst_audio_quantize_free (GstAudioQuantize * quant)
{
  g_return_if_fail (quant != NULL);

  g_free (quant->error_buf);
  g_free (quant->coeffs);
  g_free (quant->last_random);
  g_free (quant->dither_buf);

  g_slice_free (GstAudioQuantize, quant);
}

/**
 * gst_audio_quantize_reset:
 * @quant: a #GstAudioQuantize
 *
 * Reset @quant to the state is was when created, clearing any
 * history it might have.
 */
void
gst_audio_quantize_reset (GstAudioQuantize * quant)
{
  g_free (quant->error_buf);
  quant->error_buf = NULL;
  quant->error_size = 0;
}

/**
 * gst_audio_quantize_samples:
 * @quant: a #GstAudioQuantize
 * @in: input samples
 * @out: output samples
 * @samples: number of samples
 *
 * Perform quantization on @samples in @in and write the result to @out.
 *
 * In case the samples are interleaved, @in and @out must point to an
 * array with a single element pointing to a block of interleaved samples.
 *
 * If non-interleaved samples are used, @in and @out must point to an
 * array with pointers to memory blocks, one for each channel.
 *
 * @in and @out may point to the same memory location, in which case samples will be
 * modified in-place.
 */
void
gst_audio_quantize_samples (GstAudioQuantize * quant,
    const gpointer in[], gpointer out[], guint samples)
{
  guint i;

  g_return_if_fail (quant != NULL);
  g_return_if_fail (out != NULL || samples == 0);
  g_return_if_fail (in != NULL || samples == 0);

  for (i = 0; i < quant->blocks; i++)
    quant->quantize (quant, in[i], out[i], samples);
}
