GoerzelStream.h

#pragma once
 
#include <math.h>
#include "AudioStreams.h"
 
#ifndef M_PI
#  define M_PI (3.14159265358979323846f)
#endif
 
namespace audio_tools {
 
struct GoertzelConfig : public AudioInfo {
  float target_frequency = 440.0f;
  int block_size = 205;
  float threshold = 0.5f;
  float volume = 1.0f;
 
  GoertzelConfig() = default;
  GoertzelConfig(const AudioInfo& info) : AudioInfo(info) {
    target_frequency = 440.0f;
    block_size = 205;
    threshold = 0.5f;
  }
};
 
class GoertzelDetector {
 public:
  GoertzelDetector() = default;
 
  void begin(const GoertzelConfig& config) {
    this->config = config;
    this->sample_count = 0;
 
    // Calculate the coefficient k and omega
    float omega = (2.0f * M_PI * config.target_frequency) / config.sample_rate;
    coeff = 2.0f * cos(omega);
 
    // Reset state variables
    reset();
  }
 
  bool processSample(float sample) {
    // Goertzel algorithm core computation
    float s0 = sample + coeff * s1 - s2;
    s2 = s1;
    s1 = s0;
 
    sample_count++;
 
    // Check if we've processed a complete block
    if (sample_count >= config.block_size) {
      // Calculate magnitude squared
      float real = s1 - s2 * cos(2.0f * M_PI * config.target_frequency /
                                 config.sample_rate);
      float imag =
          s2 * sin(2.0f * M_PI * config.target_frequency / config.sample_rate);
      magnitude_squared = real * real + imag * imag;
      magnitude = sqrt(magnitude_squared);
 
      // Reset for next block
      reset();
      return true;
    }
 
    return false;
  }
 
  float getMagnitude() const { return magnitude; }
 
  float getMagnitudeSquared() const { return magnitude_squared; }
 
  bool isDetected(float threshold) const { return magnitude > threshold; }
 
  void reset() {
    s1 = 0.0f;
    s2 = 0.0f;
    sample_count = 0;
    magnitude = 0.0f;
    magnitude_squared = 0.0f;
  }
 
  float getTargetFrequency() const { return config.target_frequency; }
 
  int getBlockSize() const { return config.block_size; }
 
  const GoertzelConfig& getConfig() const { return config; }
 
 private:
  GoertzelConfig config;
  float coeff = 0.0f;
 
  // State variables
  float s1 = 0.0f;
  float s2 = 0.0f;
  int sample_count = 0;
 
  // Results
  float magnitude = 0.0f;
  float magnitude_squared = 0.0f;
};
 
class GoertzelStream : public AudioStream {
 public:
  GoertzelStream() = default;
 
  void setAudioInfo(AudioInfo info) override {
    AudioStream::setAudioInfo(info);
    single_config.copyFrom(info);
    // Initialize detectors for each channel
    detectors.resize(info.channels);
  }
 
  bool begin(const GoertzelConfig& config) {
    this->single_config = config;
    // Set up audio info from config
    setAudioInfo(config);
    return begin();
  }
 
  bool begin() {
    for (int i = 0; i < info.channels; i++) {
      detectors[i].begin(single_config);
    }
    return info.channels > 0;
  }
 
  void setStream(Stream& in) {
    p_stream = &in;
    p_print = &in;
  }
 
  void setOutput(Print& out) { p_print = &out; }
 
  void setChannelDetectionCallback(void (*callback)(
      int channel, float frequency, float magnitude, void* ref)) {
    channel_detection_callback = callback;
  }
 
  size_t write(const uint8_t* data, size_t len) override {
    if (p_print == nullptr) return 0;
 
    // Process samples for detection
    processSamples(data, len);
 
    if (p_print == nullptr) return len;
 
    // Pass data through unchanged
    return p_print->write(data, len);
  }
 
  size_t readBytes(uint8_t* data, size_t len) override {
    if (p_stream == nullptr) return 0;
 
    size_t result = p_stream->readBytes(data, len);
 
    // Process samples for detection
    processSamples(data, result);
 
    return result;
  }
 
  float getCurrentMagnitude(int channel = 0) {
    if (channel >= 0 && channel < detectors.size()) {
      return detectors[channel].getMagnitude();
    }
    return 0.0f;
  }
 
  bool isFrequencyDetected(int channel = 0) {
    if (channel >= 0 && channel < detectors.size()) {
      return detectors[channel].isDetected(single_config.threshold);
    }
    return false;
  }
 
  const GoertzelConfig& getConfig() const { return single_config; }
 
  void setReference(void* ref) { this->ref = ref; }
 
 protected:
  // Core detection components
  Vector<GoertzelDetector> detectors;  
  GoertzelConfig single_config;        
  
  // Stream I/O components
  Stream* p_stream = nullptr;  
  Print* p_print = nullptr;   
 
  // Callback system
  void (*channel_detection_callback)(int channel, float frequency,
                                     float magnitude, void* ref) = nullptr;  
  void* ref = this;  
 
  void checkDetection(int channel) {
    if (channel >= 0 && channel < detectors.size()) {
      float magnitude = detectors[channel].getMagnitude();
      if (magnitude > single_config.threshold) {
        float frequency = detectors[channel].getTargetFrequency();
        if (channel_detection_callback) {
          channel_detection_callback(channel, frequency, magnitude, ref);
        }
      }
    }
  }
 
  template <typename T>
  void processSamplesOfType(const uint8_t* data, size_t data_len,
                            int channels) {
    const T* samples = reinterpret_cast<const T*>(data);
    size_t num_samples = data_len / sizeof(T);
 
    for (size_t i = 0; i < num_samples; i++) {
      // Convert to normalized float and apply volume scaling
      float normalized = clip(NumberConverter::toFloatT<T>(samples[i]) * single_config.volume);
      // Distribute samples to channels in interleaved format
      int channel = i % channels;
      if (detectors[channel].processSample(normalized)) {
        checkDetection(channel);
      }
    }
  }
 
  float clip(float value) {
    // Clip the value to the range [-1.0, 1.0]
    if (value > 1.0f) return 1.0f;
    if (value < -1.0f) return -1.0f;
    return value;
  }
 
  void processSamples(const uint8_t* data, size_t data_len) {
    int channels = single_config.channels;
 
    switch (single_config.bits_per_sample) {
      case 8:
        processSamplesOfType<uint8_t>(data, data_len, channels);
        break;
      case 16:
        processSamplesOfType<int16_t>(data, data_len, channels);
        break;
      case 24:
        processSamplesOfType<int24_t>(data, data_len, channels);
        break;
      case 32:
        processSamplesOfType<int32_t>(data, data_len, channels);
        break;
      default:
        LOGE("Unsupported bits_per_sample: %d", single_config.bits_per_sample);
        break;
    }
  }
};
 
}  // namespace audio_tools