#include <Arduino.h>
#include <Chirale_TensorFlowLite.h> 

#include "tensorflow/lite/micro/micro_mutable_op_resolver.h"
#include "tensorflow/lite/micro/micro_interpreter.h"
#include "tensorflow/lite/schema/schema_generated.h"
#include "model.h"
#include <math.h>

// ===== TensorFlow Framework Objects =====
const tflite::Model* model = nullptr;
tflite::MicroInterpreter* interpreter = nullptr;
TfLiteTensor* input = nullptr;
TfLiteTensor* output = nullptr;

// ===== Strict Alignment Memory Arena =====
constexpr int kTensorArenaSize = 4 * 1024; // 4KB matching simulation requirements
alignas(16) uint8_t tensor_arena[kTensorArenaSize];

void setup() {
  Serial.begin(115200);
  delay(1500); 
  Serial.println("\n==================================================");
  Serial.println("=== TinyML Cosine Predictor – Hardware Boot   ===");
  Serial.println("==================================================");

  // 1. Initialize Model Reference
  model = tflite::GetModel(g_model);
  Serial.printf("Model size        : %d bytes\n", g_model_len);
  Serial.printf("Arena allocated   : %d bytes\n", kTensorArenaSize);

  // 2. Hardware Resource Op Allocator
  static tflite::MicroMutableOpResolver<2> resolver;
  resolver.AddFullyConnected();
  resolver.AddRelu();

  // 3. Engine Initialization
  static tflite::MicroInterpreter static_interpreter(
      model, resolver, tensor_arena, kTensorArenaSize
  );
  interpreter = &static_interpreter;

  // 4. Memory Allocations
  if (interpreter->AllocateTensors() != kTfLiteOk) {
    Serial.println("System allocation error! Infinite Loop Halted.");
    while (true) delay(1000);
  }

  // 5. Extraction of Hardware Tensor Pointers
  input = interpreter->input(0);
  output = interpreter->output(0);

  // Print system parameters mirroring Python extraction metrics
  Serial.printf("Input  scale      : %.6f\n", input->params.scale);
  Serial.printf("Input  zero_point : %d\n", input->params.zero_point);
  Serial.printf("Output scale      : %.6f\n", output->params.scale);
  Serial.printf("Output zero_point : %d\n", output->params.zero_point);
  Serial.println("==================================================");
  
  Serial.println("predicted,actual");
}

void loop() {
  static int step = 0;
  const int steps_per_period = 80;

  // ── a) Compute float x matching simulation indexing geometry ──────────────
  int i = step % steps_per_period;
  float x = ((float)i / (float)(steps_per_period - 1)) * 2.0 * PI;

  // ── b) Hardware Quantization Math: float → int8 ───────────────────────────
  float raw_quant = (x / input->params.scale);
  // Manual rounding implementation to reflect Python's round() exactly
  int32_t rounded_quant = (raw_quant >= 0) ? (int32_t)(raw_quant + 0.5f) : (int32_t)(raw_quant - 0.5f);
  int32_t q_value = rounded_quant + input->params.zero_point;
  
  // Hard clipping inside int8 borders
  if (q_value < -128) q_value = -128;
  if (q_value > 127) q_value = 127;

  // ── c) Load data buffer and prompt inference execution ────────────────────
  input->data.int8[0] = (int8_t)q_value;
  
  if (interpreter->Invoke() != kTfLiteOk) {
    Serial.println("Inference Pipeline Fault");
    return;
  }

  // ── d) Read output data matrix and perform Dequantization ─────────────────
  int8_t y_q = output->data.int8[0];
  float y_predicted = ((float)y_q - (float)output->params.zero_point) * input->params.scale;

  // ── e) Mathematical Ground Truth ──────────────────────────────────────────
  float y_actual = cos(x);

  // Emit data to Wokwi virtual serial log plotter window
  Serial.print(y_predicted, 4);
  Serial.print(",");
  Serial.println(y_actual, 4);

  // Increment step counter indefinitely
  step++;

  delay(50);
}