#include <Arduino.h>
#include <array>
#include <cmath>

// 1. Define a lightweight pre-trained Multi-Layer Perceptron (MLP) structure
// Weights representing a trained model embedded directly into Flash ROM
constexpr size_t INPUT_DIM = 1;
constexpr size_t HIDDEN_DIM = 4;
constexpr size_t OUTPUT_DIM = 1;

struct TinyModelWeights {
    float w1[INPUT_DIM][HIDDEN_DIM];
    float b1[HIDDEN_DIM];
    float w2[HIDDEN_DIM][OUTPUT_DIM];
    float b2[OUTPUT_DIM];
};

// Static pre-trained weights simulating an exported model
const TinyModelWeights model = {
    { { 2.5f, -1.2f, 0.8f, 1.5f } }, // w1
    { -0.1f, 0.2f, -0.3f, 0.1f },     // b1
    { { 1.5f }, { -2.0f }, { 0.5f }, { 1.0f } }, // w2
    { 0.05f }                        // b2
};

// 2. Pure C++ Activation Function
float relu(float x) {
    return x > 0.0f ? x : 0.0f;
}

// 3. On-chip inference execution function
float predict(float input_value) {
    std::array<float, HIDDEN_DIM> hidden_layer{};

    // Input to Hidden Layer (Dense + ReLU)
    for (size_t j = 0; j < HIDDEN_DIM; ++j) {
        float activation = model.b1[j];
        for (size_t i = 0; i < INPUT_DIM; ++i) {
            activation += input_value * model.w1[i][j];
        }
        hidden_layer[j] = relu(activation);
    }

    // Hidden to Output Layer (Dense Linear)
    float output = model.b2[0];
    for (size_t j = 0; j < HIDDEN_DIM; ++j) {
        output += hidden_layer[j] * model.w2[j][0];
    }

    return output;
}

void setup() {
    Serial.begin(115200);
    delay(1000);
    Serial.println("Native Embedded Neural Network Initialized!");
}

void loop() {
    // Dynamically sweep the input feature back and forth (simulating object tracking positions)
    static float sensor_input = 0.0f;
    static float direction = 0.1f;

    sensor_input += direction;
    if (sensor_input >= 1.0f || sensor_input <= 0.0f) {
        direction *= -1.0f; // Reverse at boundaries
    }
    
    // Run real-time inference on the ESP32 core
    float prediction = predict(sensor_input);

    Serial.print("Input Feature: ");
    Serial.print(sensor_input, 2);
    Serial.print(" | Native NN Inference Output: ");
    Serial.println(prediction, 4);

    delay(300);
}