// ESP32 Firmware: LD2410 + TFLite Inference + SSD1306 Display
// Tác giả: Auto-generated from train_ld2410_presence.ipynb
// Chức năng: Read LD2410 sensor → Extract features → Run TFLite inference → Display on SSD1306

#include <Arduino.h>
#include <HardwareSerial.h>
#include <Wire.h>
#include <Adafruit_GFX.h>
#include <Adafruit_SSD1306.h>
#include "model_int8.h"

#include <tensorflow/lite/micro/micro_interpreter.h>
#include <tensorflow/lite/micro/micro_log.h>
#include <tensorflow/lite/micro/micro_mutable_op_resolver.h>
#include <tensorflow/lite/micro/system_setup.h>
#include <tensorflow/lite/schema/schema_generated.h>

// ==================== SSD1306 Display Setup ====================
#define SCREEN_WIDTH 128
#define SCREEN_HEIGHT 64
#define OLED_RESET -1
#define SCREEN_ADDRESS 0x3C  // 0x3D for some displays

Adafruit_SSD1306 display(SCREEN_WIDTH, SCREEN_HEIGHT, &Wire, OLED_RESET);

// ==================== Test Data (20 samples) ====================
const float TEST_DATA[][4] = {
    // No person detected (10 samples)
    {5.2f, 0.5f, 25.0f, 15.0f},
    {4.8f, 0.3f, 22.0f, 12.0f},
    {5.1f, 0.4f, 26.0f, 14.0f},
    {4.9f, 0.2f, 24.0f, 13.0f},
    {5.0f, 0.6f, 25.0f, 14.0f},
    {5.3f, 0.5f, 27.0f, 15.0f},
    {4.7f, 0.3f, 23.0f, 12.0f},
    {5.2f, 0.4f, 25.0f, 13.0f},
    {4.9f, 0.5f, 26.0f, 14.0f},
    {5.1f, 0.2f, 24.0f, 12.0f},
    // Person detected (10 samples)
    {125.7f, 45.3f, 92.0f, 95.0f},
    {128.2f, 48.1f, 94.0f, 96.0f},
    {130.5f, 46.8f, 91.0f, 94.0f},
    {127.3f, 44.6f, 93.0f, 95.0f},
    {132.1f, 50.2f, 95.0f, 97.0f},
    {126.8f, 47.3f, 92.0f, 95.0f},
    {129.4f, 49.1f, 94.0f, 96.0f},
    {128.7f, 45.9f, 93.0f, 94.0f},
    {131.2f, 48.5f, 95.0f, 96.0f},
    {127.5f, 46.2f, 92.0f, 95.0f}
};
const int TEST_DATA_SIZE = 20;
int test_data_index = 0;
bool use_test_data = true;  // Set to false to use real LD2410

// ==================== Configuration ====================
#define LD2410_RX_PIN 16
#define LD2410_TX_PIN 17
#define LD2410_BAUD 115200

#define INFERENCE_SERIAL Serial  // Output inference results via USB

// Use sizes from model_data.h - don't redefine
const float THRESHOLD_PROB = 0.5f;  // Probability threshold

// ==================== TFLite Setup ====================
namespace {
const tflite::Model* model = nullptr;
tflite::MicroInterpreter* interpreter = nullptr;
TfLiteTensor* input = nullptr;
TfLiteTensor* output = nullptr;

// Tensor arena for model inference (increase if needed)
constexpr int kTensorArenaSize = 150 * 1024;  // 150KB
uint8_t tensor_arena[kTensorArenaSize];
}  // namespace

// ==================== LD2410 Sensor Data ====================
struct LD2410Sample {
    float distance;    // f1: distance reading
    float velocity;    // f2: velocity reading
    float signal;      // f3: signal strength
    float confidence;  // f4: confidence level
};

// Circular buffer for windowing
class SensorBuffer {
   private:
    LD2410Sample buffer[WINDOW_SIZE];
    int head = 0;
    int count = 0;

   public:
    void add(LD2410Sample sample) {
        buffer[head] = sample;
        head = (head + 1) % WINDOW_SIZE;
        if (count < WINDOW_SIZE) count++;
    }

    bool is_full() const { return count == WINDOW_SIZE; }

    int get_count() const { return count; }

    void get_window(float out[WINDOW_SIZE][4]) {
        for (int i = 0; i < WINDOW_SIZE; i++) {
            int idx = (head - WINDOW_SIZE + i + WINDOW_SIZE) % WINDOW_SIZE;
            out[i][0] = buffer[idx].distance;
            out[i][1] = buffer[idx].velocity;
            out[i][2] = buffer[idx].signal;
            out[i][3] = buffer[idx].confidence;
        }
    }
};

SensorBuffer sensor_buf;

// ==================== Display Functions ====================
void init_display() {
    if (!display.begin(SSD1306_SWITCHCAPVCC, SCREEN_ADDRESS)) {
        INFERENCE_SERIAL.println(F("SSD1306 allocation failed"));
        for (;;) delay(100);
    }
    display.clearDisplay();
    display.setTextSize(1);
    display.setTextColor(SSD1306_WHITE);
    display.setCursor(0, 0);
    display.println("LD2410 + TFLite");
    display.println("Initializing...");
    display.display();
}

void display_debug(float distance, float velocity, float prob, int prediction) {
    display.clearDisplay();
    display.setTextSize(1);
    display.setTextColor(SSD1306_WHITE);
    display.setCursor(0, 0);
    
    display.println("=== LD2410 Inference ===");
    
    display.print("Distance: ");
    display.print(distance, 1);
    display.println("m");
    
    display.print("Velocity: ");
    display.print(velocity, 1);
    display.println("m/s");
    
    display.print("Probability: ");
    display.print(prob, 2);
    display.println("");
    
    display.print("Status: ");
    display.println(prediction ? "DETECTED" : "NOT DETECTED");
    
    display.print("Sample: ");
    display.print(sensor_buf.get_count());
    display.print("/");
    display.println(WINDOW_SIZE);
    
    display.display();
}

void display_status(const char* msg) {
    display.clearDisplay();
    display.setTextSize(2);
    display.setTextColor(SSD1306_WHITE);
    display.setCursor(0, 20);
    display.println(msg);
    display.display();
}

// ==================== Feature Extraction ====================
void compute_window_features(float window[WINDOW_SIZE][4], float features[NUM_FEATURES]) {
    // For each of 4 channels, compute 6 statistics
    // stats: mean, std, min, max, range, delta

    for (int ch = 0; ch < 4; ch++) {
        // Extract channel values
        float values[WINDOW_SIZE];
        for (int i = 0; i < WINDOW_SIZE; i++) {
            values[i] = window[i][ch];
        }

        // Mean
        float mean = 0;
        for (int i = 0; i < WINDOW_SIZE; i++) {
            mean += values[i];
        }
        mean /= WINDOW_SIZE;

        // Std
        float var = 0;
        for (int i = 0; i < WINDOW_SIZE; i++) {
            var += (values[i] - mean) * (values[i] - mean);
        }
        float std = sqrt(var / WINDOW_SIZE);

        // Min, Max
        float min_val = values[0];
        float max_val = values[0];
        for (int i = 1; i < WINDOW_SIZE; i++) {
            if (values[i] < min_val) min_val = values[i];
            if (values[i] > max_val) max_val = values[i];
        }

        // Range
        float range = max_val - min_val;

        // Delta (last - first)
        float delta = values[WINDOW_SIZE - 1] - values[0];

        // Store as features
        int base = ch * 6;
        features[base + 0] = mean;
        features[base + 1] = std;
        features[base + 2] = min_val;
        features[base + 3] = max_val;
        features[base + 4] = range;
        features[base + 5] = delta;
    }
}

// ==================== Normalization ====================
void normalize_features(float features[NUM_FEATURES]) {
    for (int i = 0; i < NUM_FEATURES; i++) {
        features[i] = (features[i] - SCALER_MEAN[i]) / SCALER_SCALE[i];
    }
}

// ==================== LD2410 Serial Parser ====================
void inject_test_data() {
    if (use_test_data && test_data_index < TEST_DATA_SIZE) {
        LD2410Sample sample;
        sample.distance = TEST_DATA[test_data_index][0];
        sample.velocity = TEST_DATA[test_data_index][1];
        sample.signal = TEST_DATA[test_data_index][2];
        sample.confidence = TEST_DATA[test_data_index][3];
        sensor_buf.add(sample);
        test_data_index++;
        
        // Loop test data continuously
        if (test_data_index >= TEST_DATA_SIZE) {
            test_data_index = 0;
        }
    }
}

bool parse_ld2410_line(const char* line, LD2410Sample& sample) {
    // Expected format: "f1,f2,f3,f4"
    // Example: "120.5,25.3,85,95"

    int n = sscanf(line, "%f,%f,%f,%f", &sample.distance, &sample.velocity,
                    &sample.signal, &sample.confidence);
    return n == 4;
}

void read_ld2410_sensor() {
    static String buffer;

    while (Serial1.available()) {
        char c = Serial1.read();
        if (c == '\n') {
            buffer.trim();
            if (buffer.length() > 0) {
                LD2410Sample sample;
                if (parse_ld2410_line(buffer.c_str(), sample)) {
                    sensor_buf.add(sample);
                }
            }
            buffer = "";
        } else {
            buffer += c;
        }
    }
}

// ==================== TFLite Inference ====================
void run_inference() {
    if (!sensor_buf.is_full()) {
        return;  // Wait for window to fill
    }

    // Get window data
    float window[WINDOW_SIZE][4];
    sensor_buf.get_window(window);

    // Extract features
    float features[NUM_FEATURES];
    compute_window_features(window, features);

    // Normalize using scaler from model_data.h
    for (int i = 0; i < NUM_FEATURES; i++) {
        features[i] = (features[i] - SCALER_MEAN[i]) / SCALER_SCALE[i];
    }

    // Run inference
    if (interpreter != nullptr) {
        for (int i = 0; i < NUM_FEATURES; i++) {
            input->data.f[i] = features[i];
        }

        TfLiteStatus invoke_status = interpreter->Invoke();
        if (invoke_status != kTfLiteOk) {
            INFERENCE_SERIAL.println("ERROR: TFLite invoke failed");
            return;
        }

        // Get output (probability of person detected)
        float prob = output->data.f[0];
        int prediction = (prob >= THRESHOLD_PROB) ? 1 : 0;

        // Serial output format: timestamp,distance,velocity,prob,decision
        INFERENCE_SERIAL.print(millis());
        INFERENCE_SERIAL.print(",");
        INFERENCE_SERIAL.print(window[WINDOW_SIZE - 1][0], 2);
        INFERENCE_SERIAL.print(",");
        INFERENCE_SERIAL.print(window[WINDOW_SIZE - 1][1], 2);
        INFERENCE_SERIAL.print(",");
        INFERENCE_SERIAL.print(prob, 4);
        INFERENCE_SERIAL.print(",");
        INFERENCE_SERIAL.println(prediction ? "DETECTED" : "NOT_DETECTED");
    } else {
        // No model - test mode: classify based on distance threshold
        float dist = window[WINDOW_SIZE - 1][0];
        float vel = window[WINDOW_SIZE - 1][1];
        
        // Simple heuristic: distance > 100m → person detected
        int test_prediction = (dist > 100.0f) ? 1 : 0;
        float test_prob = dist / 200.0f;  // Scale distance to [0,1]
        test_prob = (test_prob > 1.0f) ? 1.0f : test_prob;
        
        INFERENCE_SERIAL.print(millis());
        INFERENCE_SERIAL.print(",");
        INFERENCE_SERIAL.print(dist, 2);
        INFERENCE_SERIAL.print(",");
        INFERENCE_SERIAL.print(vel, 2);
        INFERENCE_SERIAL.print(",");
        INFERENCE_SERIAL.print(test_prob, 4);
        INFERENCE_SERIAL.print(",");
        INFERENCE_SERIAL.println(test_prediction ? "DETECTED" : "NOT_DETECTED");
    }
}

// ==================== Setup & Loop ====================
void setup() {
    // USB Serial for logging/inference output
    INFERENCE_SERIAL.begin(115200);
    delay(1000);
    INFERENCE_SERIAL.println("\n=== LD2410 TFLite Inference ===");
    INFERENCE_SERIAL.println("Initializing...");

    // Load TFLite model
    INFERENCE_SERIAL.println("Loading TFLite model...");
    model = tflite::GetModel(MODEL_DATA);
    if (model->version() != TFLITE_SCHEMA_VERSION) {
        INFERENCE_SERIAL.println("ERROR: Model schema version mismatch");
        return;
    }

    // Create interpreter
    static tflite::MicroMutableOpResolver<4> resolver;
    resolver.AddFullyConnected();
    resolver.AddRelu();
    resolver.AddLogistic();
    resolver.AddQuantize();

    static tflite::MicroInterpreter static_interpreter(
        model, resolver, tensor_arena, kTensorArenaSize, nullptr);
    interpreter = &static_interpreter;

    // Allocate tensors
    TfLiteStatus allocate_status = interpreter->AllocateTensors();
    if (allocate_status != kTfLiteOk) {
        INFERENCE_SERIAL.println("ERROR: AllocateTensors failed");
        interpreter = nullptr;
        return;
    }

    // Get input/output tensors
    input = interpreter->input(0);
    output = interpreter->output(0);

    INFERENCE_SERIAL.println("Model loaded successfully!");
    INFERENCE_SERIAL.println("Ready. Processing sensor data...");
}

void loop() {
    // Read incoming sensor data or inject test data
    if (use_test_data) {
        inject_test_data();
    } else {
        read_ld2410_sensor();
    }

    // Run inference only once when window is full
    static bool done = false;
    if (sensor_buf.is_full() && !done) {
        run_inference();
        done = true;
        INFERENCE_SERIAL.println("\n=== DONE ===");
        while(1) delay(1000);  // Halt
    }

    delay(10);
}