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 = 0.0f;
  int block_size = 205;
  float threshold = 0.5f;
  float volume = 1.0f;
  uint8_t channel = 0;
 
  GoertzelConfig() = default;
  GoertzelConfig(const AudioInfo& info) : AudioInfo(info) {
    target_frequency = 440.0f;
    block_size = 205;
    threshold = 0.5f;
  }
};
 
class GoertzelDetector {
 public:
  GoertzelDetector() = default;
 
  bool begin(const GoertzelConfig& config) {
    this->config = config;
    this->sample_count = 0;
 
    if (config.target_frequency == 0.0f) {
      return false;
    }
 
    // 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();
    return true;
  }
 
  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; }
 
  bool isDetected() const { return isDetected(config.threshold); }
 
  void reset() {
    s1 = 0.0f;
    s2 = 0.0f;
    sample_count = 0;
    magnitude_squared = 0.0f;
  }
 
  float getTargetFrequency() const { return config.target_frequency; }
 
  int getBlockSize() const { return config.block_size; }
 
  const GoertzelConfig& getConfig() const { return config; }
 
  void setReference(void* ref) { this->reference = ref; }
 
  void* getReference() { return reference; }
 
 protected:
  GoertzelConfig config;
  float coeff = 0.0f;
  void* reference = nullptr;
 
  // 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:
  // Default Constructor with no output or input
  GoertzelStream() = default;
  GoertzelStream(Print& out) { setOutput(out); }
  GoertzelStream(AudioOutput& out) { setOutput(out); }
  GoertzelStream(Stream& io) { setStream(io); };
  GoertzelStream(AudioStream& io) { setStream(io); };
 
  void setAudioInfo(AudioInfo info) override {
    AudioStream::setAudioInfo(info);
    default_config.copyFrom(info);
    // Initialize detectors for each frequency
    begin();
  }
 
  GoertzelConfig defaultConfig() {
    GoertzelConfig result;
    return result;
  }
 
  bool begin(GoertzelConfig config) {
    config.target_frequency = 0;
    if (config.channel >= config.channels) {
      LOGE("GoertzelStream: channel %d out of range (max %d)", config.channel,
           config.channels);
      return false;
    }
    this->default_config = config;
 
    // Set up audio info from config
    setAudioInfo(config);
    return begin();
  }
 
  bool begin() {
    int i = 0;
    detectors.clear();
    for (float freq : frequencies) {
      GoertzelConfig cfg = default_config;
      cfg.target_frequency = freq;
      GoertzelDetector detector;
      if (i < references.size()) {
        detector.setReference(references[i]);
      }
      detector.begin(cfg);
      detectors.push_back(detector);
      i++;
    }
    sample_no = 0;
    return true;
  }
 
  void end() {
    detectors.clear();
    AudioStream::end();
  }
 
  void setStream(Stream& in) {
    p_stream = &in;
    p_print = &in;
  }
 
  void setStream(AudioStream& io) {
    setStream((Stream&)io);
    addNotifyAudioChange(io);
  }
 
  void setOutput(Print& out) { p_print = &out; }
 
  void setOutput(AudioOutput& out) {
    setOutput((Print&)out);
    addNotifyAudioChange(out);
  }
 
  void setFrequencyDetectionCallback(void (*callback)(float frequency,
                                                      float magnitude,
                                                      void* ref)) {
    frequency_detection_callback = callback;
  }
 
  size_t write(const uint8_t* data, size_t len) override {
    // 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;
  }
 
  const GoertzelConfig& getConfig() const { return default_config; }
 
  void setReference(void* ref) { this->ref = ref; }
 
  GoertzelDetector& getDetector(int no) { return detectors[no]; }
 
  void addFrequency(float freq) { frequencies.push_back(freq); }
 
  void addFrequency(float freq, void* ref) {
    frequencies.push_back(freq);
    references.push_back(ref);
  }
 
 protected:
  // Core detection components
  Vector<GoertzelDetector>
      detectors;                  
  Vector<float> frequencies;      
  Vector<void*> references;       
  GoertzelConfig default_config;  
  // Stream I/O components
  Stream* p_stream = nullptr;  
  Print* p_print = nullptr;    
 
  // Callback system
  void (*frequency_detection_callback)(float frequency, float magnitude,
                                       void* ref) =
      nullptr;       
  void* ref = this;  
 
  size_t sample_no = 0;  
 
  void checkDetection(GoertzelDetector& detector) {
    float magnitude = detector.getMagnitude();
    if (magnitude > 0.0f)
      LOGD("frequency: %f / magnitude: %f / threshold: %f", detector.getTargetFrequency(), magnitude, default_config.threshold);
 
    if (magnitude > default_config.threshold) {
      float frequency = detector.getTargetFrequency();
      void* reference = detector.getReference();
      if (reference == nullptr) reference = ref;
      if (frequency_detection_callback != nullptr) {
        frequency_detection_callback(frequency, magnitude, reference);
      }
    }
  }
 
  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++) {
      // select only samples for the configured channel
      if (sample_no % channels == default_config.channel) {
        float normalized = clip(NumberConverter::toFloatT<T>(samples[i]) *
                                default_config.volume);
        LOGD("sample: %f", normalized);                        
        // process all frequencies
        for (auto& detector : detectors) {
          if (detector.processSample(normalized)) {
            checkDetection(detector);
          }
        }
      }
      sample_no++;
    }
  }
 
  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) {
    // return if there is nothing to detect
    if (detectors.empty()) return;
    int channels = default_config.channels;
 
    switch (default_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", default_config.bits_per_sample);
        break;
    }
  }
};
 
}  // namespace audio_tools