/*
libspecbleach - A spectral processing library

Copyright 2022 Luciano Dato <lucianodato@gmail.com>

This library 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; either
version 2.1 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
Lesser General Public License for more details.

You should have received a copy of the GNU Lesser General Public
License along with this library; if not, write to the Free Software
Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301  USA
*/

#include "spectral_smoother.h"
#include "transient_detector.h"
#include <math.h>
#include <stdlib.h>
#include <string.h>

static void spectrum_time_smoothing(SpectralSmoother* self, float smoothing);
static void spectrum_transient_aware_time_smoothing(SpectralSmoother* self,
                                                    float smoothing,
                                                    float* spectrum);

struct SpectralSmoother {
  uint32_t fft_size;
  uint32_t real_spectrum_size;
  float adaptive_coefficient;
  float previous_adaptive_coefficient;
  TimeSmoothingType type;

  float* noise_spectrum;
  float* smoothed_spectrum;
  float* smoothed_spectrum_previous;

  TransientDetector* transient_detection;
};

SpectralSmoother* spectral_smoothing_initialize(const uint32_t fft_size,
                                                TimeSmoothingType type) {
  SpectralSmoother* self =
      (SpectralSmoother*)calloc(1U, sizeof(SpectralSmoother));

  if (!self) {
    return NULL;
  }

  self->fft_size = fft_size;
  self->real_spectrum_size = (self->fft_size / 2U) + 1U;
  self->type = type;
  self->previous_adaptive_coefficient = 0.F;
  self->adaptive_coefficient = 0.F;

  self->noise_spectrum =
      (float*)calloc(self->real_spectrum_size, sizeof(float));
  self->smoothed_spectrum =
      (float*)calloc(self->real_spectrum_size, sizeof(float));
  self->smoothed_spectrum_previous =
      (float*)calloc(self->real_spectrum_size, sizeof(float));

  self->transient_detection = transient_detector_initialize(self->fft_size);

  if (!self->noise_spectrum || !self->smoothed_spectrum ||
      !self->smoothed_spectrum_previous || !self->transient_detection) {
    spectral_smoothing_free(self);
    return NULL;
  }

  return self;
}

void spectral_smoothing_free(SpectralSmoother* self) {
  if (!self) {
    return;
  }
  transient_detector_free(self->transient_detection);

  free(self->noise_spectrum);
  free(self->smoothed_spectrum);
  free(self->smoothed_spectrum_previous);

  free(self);
}

bool spectral_smoothing_run(SpectralSmoother* self,
                            TimeSmoothingParameters parameters,
                            float* signal_spectrum) {
  if (!self || !signal_spectrum) {
    return false;
  }

  memcpy(self->smoothed_spectrum, signal_spectrum,
         sizeof(float) * self->real_spectrum_size);

  switch (self->type) {
    case FIXED:
      spectrum_time_smoothing(self, parameters.smoothing);
      break;
    case TRANSIENT_AWARE:
      spectrum_transient_aware_time_smoothing(self, parameters.smoothing,
                                              signal_spectrum);
      break;
    default:
      break;
  }

  memcpy(self->smoothed_spectrum_previous, self->smoothed_spectrum,
         sizeof(float) * self->real_spectrum_size);
  memcpy(signal_spectrum, self->smoothed_spectrum,
         sizeof(float) * self->real_spectrum_size);

  return true;
}

static void spectrum_transient_aware_time_smoothing(SpectralSmoother* self,
                                                    const float smoothing,
                                                    float* spectrum) {

  if (!transient_detector_run(self->transient_detection, spectrum)) {
    for (uint32_t k = 0U; k < self->real_spectrum_size; k++) {
      if (self->smoothed_spectrum[k] > self->smoothed_spectrum_previous[k]) {
        self->smoothed_spectrum[k] =
            (smoothing * self->smoothed_spectrum_previous[k]) +
            ((1.F - smoothing) * self->smoothed_spectrum[k]);
      }
    }
  }
}

static void spectrum_time_smoothing(SpectralSmoother* self,
                                    const float smoothing) {
  for (uint32_t k = 0U; k < self->real_spectrum_size; k++) {
    if (self->smoothed_spectrum[k] > self->smoothed_spectrum_previous[k]) {
      self->smoothed_spectrum[k] =
          (smoothing * self->smoothed_spectrum_previous[k]) +
          ((1.F - smoothing) * self->smoothed_spectrum[k]);
    }
  }
}
