#include "arduinoFFT.h"

#define NUM_SAMPLES 4096 // Number of samples; must be a power of 2
#define SAMPLE_RATE 48000 // Sampling rate in Hz

// LED Pins
#define RED_LED_PIN 2   // GPIO2
#define GREEN_LED_PIN 4 // GPIO4

// Detection parameters
const float low_range = 5000.0f;
const float high_range = 10000.0f;
const float peak_threshold_std = 1.0f;
const float peaks_magnitude_rel_threshold = 0.5f;

// Global variables
static float vReal[NUM_SAMPLES];
static float vImag[NUM_SAMPLES];
ArduinoFFT<float> FFT = ArduinoFFT<float>(vReal, vImag, NUM_SAMPLES, SAMPLE_RATE);

// ------ Brake sound detector function ------
bool brake_sound_detector(float in[], uint16_t num_samples, float samplingFrequency) {
    // Center and normalize
    float mean = 0.0f;
    for (int i = 0; i < num_samples; i++) {
        mean += in[i];
    }
    mean /= num_samples;
    
    for (int i = 0; i < num_samples; i++) {
        in[i] -= mean;
    }
    
    float max_val = 0.0f;
    for (int i = 0; i < num_samples; i++) {
        if (abs(in[i]) > max_val) {
            max_val = abs(in[i]);
        }
    }
    
    for (int i = 0; i < num_samples; i++) {
        in[i] /= max_val;
    }

    // Prepare input data for FFT
    for (int i = 0; i < num_samples; i++) {
        vReal[i] = in[i];  // Load real part
        // vImag[i] is already 0
    }

    // Perform FFT
    FFT.windowing(FFTWindow::Hamming, FFTDirection::Forward);
    FFT.compute(FFTDirection::Forward);
    FFT.complexToMagnitude();

    // Analyze magnitudes to detect brake sound
    float sum_magnitudes = 0, mean_magnitude = 0, std_dev = 0;
    int count = 0;

    // Calculate mean of magnitudes in the target range
    for (uint16_t i = 0; i < num_samples / 2; i++) {
        float freq = (i * samplingFrequency / num_samples);
        if (freq >= low_range && freq <= high_range) {
            sum_magnitudes += vReal[i];
            count++;
        }
    }

    if (count > 0) {
        mean_magnitude = sum_magnitudes / count;

        // Calculate standard deviation
        for (uint16_t i = 0; i < num_samples / 2; i++) {
            float freq = (i * samplingFrequency / num_samples);
            if (freq >= low_range && freq <= high_range) {
                std_dev += pow(vReal[i] - mean_magnitude, 2);
            }
        }
        std_dev = sqrt(std_dev / count);
        float peak_threshold = mean_magnitude + peak_threshold_std * std_dev;

        // Identify peaks above the threshold
        float peak_sum = 0;
        for (uint16_t i = 0; i < num_samples / 2; i++) {
            float freq = (i * samplingFrequency / num_samples);
            if (freq >= low_range && freq <= high_range && vReal[i] > peak_threshold) {
                peak_sum += vReal[i];
            }
        }

        return peak_sum > sum_magnitudes * peaks_magnitude_rel_threshold;
    }

    return false;
}

void generate_signal(const char* mode) {
    const float amplitude = 50.0f;  // Amplitude of the sine wave
    const float target_frequency_1 = 6000.0f;  // Brake sound frequency 1 in Hz
    const float target_frequency_2 = 7500.0f;  // Brake sound frequency 2 in Hz
    const float target_frequency_3 = 9000.0f;  // Brake sound frequency 3 in Hz

    for (int i = 0; i < NUM_SAMPLES; i++) {
        if (strcmp(mode, "normal") == 0) {
            vReal[i] = (rand() % 65536) - 32768;  // Simulate 16-bit audio noise
            vImag[i] = 0.0f;
        } else if (strcmp(mode, "brake_sound") == 0) {
            float time = i / float(SAMPLE_RATE);
            float sine_wave_1 = amplitude * sin(2 * M_PI * target_frequency_1 * time);
            float sine_wave_2 = amplitude * sin(2 * M_PI * target_frequency_2 * time);
            float sine_wave_3 = amplitude * sin(2 * M_PI * target_frequency_3 * time);
            float noise = (rand() % 65536) - 32768;
            vReal[i] = noise / 32768.0f + sine_wave_1 + sine_wave_2 + sine_wave_3;
            vImag[i] = 0.0f;
        }
    }
}

void setup() {
  Serial.begin(115200);
  pinMode(RED_LED_PIN, OUTPUT);
  pinMode(GREEN_LED_PIN, OUTPUT);
  while(!Serial);
  Serial.println("Setup Complete - Ready to simulate");
}

void loop() {
  generate_signal("normal");
  if (brake_sound_detector(vReal, NUM_SAMPLES, SAMPLE_RATE)) {
    digitalWrite(RED_LED_PIN, HIGH);
    digitalWrite(GREEN_LED_PIN, LOW);
    Serial.println("Analysis output: Brake sound detected");
  } else {
    digitalWrite(RED_LED_PIN, LOW);
    digitalWrite(GREEN_LED_PIN, HIGH);
    Serial.println("Analysis output: Normal");
  }
  delay(1000); // Wait a second

  generate_signal("brake_sound");
  if (brake_sound_detector(vReal, NUM_SAMPLES, SAMPLE_RATE)) {
    digitalWrite(RED_LED_PIN, HIGH);
    digitalWrite(GREEN_LED_PIN, LOW);
    Serial.println("Analysis output: Brake sound detected");
  } else {
    digitalWrite(RED_LED_PIN, LOW);
    digitalWrite(GREEN_LED_PIN, HIGH);
    Serial.println("Analysis output: Normal");
  }
  delay(1000); // Delay before the next cycle
}