#include <TensorFlowLite_ESP32.h>
#include "sine_model.h" // The header file containing your converted model

// Include necessary TFLite Micro components
#include <tensorflow/lite/micro/all_ops_resolver.h>
#include <tensorflow/lite/micro/micro_interpreter.h>
#include <tensorflow/lite/micro/micro_log.h>
#include <tensorflow/lite/micro/system_setup.h>
#include <tensorflow/lite/schema/schema_generated.h>

// Define the size of the tensor arena (memory buffer for tensors)
// 40KB is a common starting point for small models on ESP32-S3.
constexpr int kTensorArenaSize = 20 * 1024; 
uint8_t tensor_arena[kTensorArenaSize];

// Pointers for the interpreter and tensors
tflite::MicroInterpreter* interpreter = nullptr;
TfLiteTensor* input = nullptr;
TfLiteTensor* output = nullptr;

void setup() {
  Serial.begin(115200);
  delay(1000); 
  Serial.println("--- Starting TFLite Micro Setup ---");

  // 1. Get Model Data
  const tflite::Model* model = tflite::GetModel(sine_model_tflite);
  if (model->version() != TFLITE_SCHEMA_VERSION) {
    Serial.println("Model schema version mismatch!");
    while(1);
  }

  // 2. Instantiate Operation Resolver
  // This provides all necessary math operations (Ops) for the model.
  static tflite::AllOpsResolver resolver;

  // 3. Instantiate Micro Interpreter
  static tflite::MicroInterpreter static_interpreter(
    model, resolver, tensor_arena, kTensorArenaSize);
  
  interpreter = &static_interpreter;

  // 4. Allocate Tensors
  // This assigns segments of the tensor_arena to the model's tensors.
  TfLiteStatus allocate_status = interpreter->AllocateTensors();
  if (allocate_status != kTfLiteOk) {
    Serial.println("FATAL: AllocateTensors() failed!");
    while(1);
  }

  // 5. Get Input and Output Pointers
  // These pointers are used to feed data into and read data from the model.
  input = interpreter->input(0);
  output = interpreter->output(0);
  
  Serial.print("Input Shape: ");
  Serial.println(input->dims->data[1]);
  Serial.print("Output Shape: ");
  Serial.println(output->dims->data[1]);
  Serial.println("✅ TFLite Model loaded successfully.");
}

void loop() {
  // Wait 3 seconds between inferences
  delay(3000); 

  // 1. Generate a random input angle (x-value) between -PI and +PI
  // The 'random' function returns long, we scale it to fit the range.
  float random_angle = ((float)random(10000) / 10000.0) * (2 * PI) - PI;
  float true_sine = sin(random_angle);

  // 2. Feed the input into the model's input tensor
  // TFLite uses an array for data, even for a single value: input->data.f[0]
  input->data.f[0] = random_angle;

  // 3. Run the model (Inference)
  TfLiteStatus invoke_status = interpreter->Invoke();
  if (invoke_status != kTfLiteOk) {
    Serial.println("❌ Invoke failed!");
    return;
  }

  // 4. Read the prediction from the model's output tensor
  float predicted_sine = output->data.f[0];

  // 5. Print results to the Serial Monitor
  Serial.println("-------------------------------------");
  Serial.print("Input Angle (rad): ");
  Serial.print(random_angle, 4);
  Serial.print(" | True sin(): ");
  Serial.print(true_sine, 4);
  Serial.print(" | NN Predict: ");
  Serial.println(predicted_sine, 4);
}