#include <Wire.h>
#include <MPU6050.h>
#include <WiFi.h>
#include <PubSubClient.h>
#include <FFT.h>

#define SDA_PIN 6
#define SCL_PIN 7
#define SAMPLE_COUNT 64   // FFT sample size (Power of 2)
#define SAMPLE_RATE 500   // Hz sampling rate
#define FILTER_SIZE 5     // Moving average filter size

// WiFi Credentials
const char* ssid = "Wokwi-GUEST";
const char* password = "";
const char* mqtt_server = "broker.hivemq.com";

WiFiClient espClient;
PubSubClient client(espClient);
MPU6050 mpu;
FFT fft;
fft.computeFFT(vReal, vImag, SAMPLE_COUNT);



int16_t axBuffer[FILTER_SIZE], ayBuffer[FILTER_SIZE], azBuffer[FILTER_SIZE];
int filterIndex = 0;
double vReal[SAMPLE_COUNT];  // Store FFT sensor values
double vImag[SAMPLE_COUNT];  // Imaginary FFT part

// Moving Average Function
int16_t movingAverage(int16_t *buffer, int16_t newVal) {
    buffer[filterIndex] = newVal;
    int32_t sum = 0;
    for (int i = 0; i < FILTER_SIZE; i++) sum += buffer[i];
    return sum / FILTER_SIZE;
}

// Normalize Data Function
float normalize(float value, float minVal, float maxVal) {
    return (value - minVal) / (maxVal - minVal);
}

// MQTT Setup
void setup_wifi() {
    Serial.print("Connecting to WiFi...");
    WiFi.begin(ssid, password);
    while (WiFi.status() != WL_CONNECTED) {
        delay(500);
        Serial.print(".");
    }
    Serial.println("WiFi connected!");
}

void reconnect_mqtt() {
    while (!client.connected()) {
        Serial.print("Connecting to MQTT...");
        if (client.connect("ESP32C6_Client")) Serial.println("Connected!");
        else {
            Serial.print("Failed, rc=");
            Serial.print(client.state());
            Serial.println(" retrying...");
            delay(5000);
        }
    }
}

void setup() {
    Serial.begin(115200);
    Wire.begin(SDA_PIN, SCL_PIN);
    mpu.initialize();

    setup_wifi();
    client.setServer(mqtt_server, 1883);
}

void loop() {
    // Ensure MQTT Connection
    if (!client.connected()) reconnect_mqtt();
    client.loop();

    int16_t ax, ay, az, gx, gy, gz, temp;
    mpu.getAcceleration(&ax, &ay, &az);
    mpu.getRotation(&gx, &gy, &gz);
    temp = mpu.getTemperature();
    float tempC = (temp / 340.00) + 36.53;

    // Apply Moving Average Filtering
    ax = movingAverage(axBuffer, ax);
    ay = movingAverage(ayBuffer, ay);
    az = movingAverage(azBuffer, az);

    // Compute RMS Vibration
    float rmsAccel = sqrt((ax * ax + ay * ay + az * az) / 3.0);

    // Normalize Sensor Data
    float axNormalized = normalize(ax, -16000, 16000);
    float ayNormalized = normalize(ay, -16000, 16000);
    float azNormalized = normalize(az, -16000, 16000);

    // FFT Analysis
    for (int i = 0; i < SAMPLE_COUNT; i++) {
        vReal[i] = sqrt(ax * ax + ay * ay + az * az);
        vImag[i] = 0;
        delayMicroseconds(1000000 / SAMPLE_RATE);  // Sampling rate
    }

    FFT.Windowing(vReal, SAMPLE_COUNT, FFT_WIN_TYP_HAMMING, FFT_FORWARD);
    FFT.Compute(vReal, vImag, SAMPLE_COUNT, FFT_FORWARD);
    FFT.ComplexToMagnitude(vReal, vImag, SAMPLE_COUNT);

    // Prepare MQTT Payload
    String payload = "{\"raw_data\":{\"temp\":" + String(tempC) + 
                     ",\"rms\":" + String(rmsAccel) +
                     ",\"ax\":" + String(axNormalized) + ",\"ay\":" + String(ayNormalized) + ",\"az\":" + String(azNormalized) +
                     "},\"fft\":[";

    for (int i = 0; i < (SAMPLE_COUNT / 2); i++) {
        double frequency = i * SAMPLE_RATE / SAMPLE_COUNT;
        payload += "{\"freq\":" + String(frequency) + ",\"mag\":" + String(vReal[i]) + "},";
    }

    payload += "]}";
    client.publish("esp32c6/mpu6050_full", payload.c_str());

    Serial.println("Data sent via MQTT: " + payload);
    delay(500);
}