#include <Arduino.h>
#include <pgmspace.h>

#include <TensorFlowLite_ESP32.h>
#include "tensorflow/lite/c/common.h"
#include "tensorflow/lite/micro/all_ops_resolver.h"
#include "tensorflow/lite/micro/micro_error_reporter.h"
#include "tensorflow/lite/micro/micro_interpreter.h"
#include "tensorflow/lite/schema/schema_generated.h"

#include "model_data.h"
#include "model_mean.h"
#include "model_std.h"
#include "body_replay_case.h"

#define MODEL_DATA     _content_drive_MyDrive_KFALL_RESULTS_kfall_real_160x6_esp32_kfall_160x6_int8_tflite
#define MODEL_DATA_LEN _content_drive_MyDrive_KFALL_RESULTS_kfall_real_160x6_esp32_kfall_160x6_int8_tflite_len

static const int PIN_DEPLOY = 2;
static const int PIN_BRAKE  = 4;
static const int PIN_ALERT  = 5;

static const int N = N_SAMPLES_CASE;
static const float FS = 80.0f;
static const uint32_t DT_US = (uint32_t)(1000000.0f / FS);

static const int SEG_LEN = 160;
static const int N_CH = 6;

static const float DEPLOY_THR = 0.60f;
static const float ALERT_THR  = 0.60f;

static const float ENC_ACC_EMA_ALPHA = 0.15f;
static const float BRAKE_SPEED_THR   = 1.00f;
static const float BRAKE_ACC_THR     = 17.0f;
static const float ENC_PEAK_DECAY    = 0.90f;
static const uint32_t BRAKE_HOLD_MS  = 500;

// ============================
// Encoder replay arrays
// ============================
const float ENC_SPEED[] PROGMEM = {
0.640586018, 0.513547322, 0.495718676, 0.529904447, 0.492370379, 0.625411230, 0.593995019, 0.581881656, 0.620493560, 0.519933143,
0.484293524, 0.567353553, 0.565464265, 0.507223887, 0.576792392, 0.623603752, 0.523334331, 0.559734043, 0.567212167, 0.574290934,
0.601145231, 0.579424065, 0.619982985, 0.509486965, 0.526394190, 0.590985826, 0.527878476, 0.522006927, 0.547725457, 0.567612494,
0.525116146, 0.579657794, 0.619378716, 0.460775252, 0.435662596, 0.587647630, 0.510183625, 0.622816936, 0.693470981, 0.542912217,
0.530819625, 0.549652216, 0.564078168, 0.547731104, 0.504810311, 0.570713482, 0.449043792, 0.569666963, 0.555024637, 0.615412074,
0.558167475, 0.574268310, 0.502682550, 0.558723827, 0.588849651, 0.601050963, 0.555936023, 0.629269716, 0.613773962, 0.517398945,
0.426615626, 0.495121953, 0.553417353, 0.557966032, 0.535324455, 0.528995677, 0.582460864, 0.520589313, 0.443632052, 0.534482730,
0.600902684, 0.507457224, 0.581983188, 0.578762407, 0.549185670, 0.570440592, 0.504123721, 0.553014572, 0.625012341, 0.546194964,
0.494883201, 0.623188469, 0.481928875, 0.632606082, 0.554759267, 0.580903335, 0.567693061, 0.551018368, 0.603921489, 0.546145522,
0.509129922, 0.552241943, 0.462758506, 0.542842705, 0.461416964, 0.551212533, 0.586473131, 0.582318884, 0.543155062, 0.615314580,
0.761157773, 0.883455402, 0.823284970, 0.916183462, 0.890654112, 1.037548355, 0.887971125, 1.001897994, 1.043357156, 1.122475820,
1.133093352, 1.174061168, 1.086239473, 1.334456432, 1.279609159, 1.272552622, 1.408134049, 1.317620451, 1.425524781, 1.492845953,
1.516159604, 1.524452812, 1.542449583, 1.531161294, 1.632997967, 1.586684103, 0.603948874, 0.584716942, 0.534453635, 0.433751896,
0.546726525, 0.537214594, 0.450112136, 0.472039893, 0.488278068, 0.503293089, 0.448808598, 0.421501361, 0.426024941, 0.399473727,
0.498226865, 0.490494249, 0.335552997, 0.388691688, 0.450109158, 0.434525226, 0.648893051, 0.613521057, 0.573900054, 0.566061798,
0.620364368, 0.573231470, 0.468737765, 0.558350776, 0.462054475, 0.591915451, 0.567513758, 0.543949781, 0.564400708, 0.464671740
};

const float ENC_ACC_RAW[] PROGMEM = {
0.000000000, -10.163095650, -1.426291708, 2.734861729, -3.002725465, 10.643268070, -2.513296891, -0.969069057, 3.088952344, -8.044833362,
-2.851169541, 6.644802352, -0.151143024, -4.659230290, 5.565480433, 3.744908789, -8.021553666, 2.911976988, 0.598249907, 0.566301307,
2.148343822, -1.737693311, 3.244713569, -8.839681523, 1.352577944, 5.167330901, -5.048587966, -0.469723961, 2.057482439, 1.590962902,
-3.399707830, 4.363331881, 3.177673706, -12.688277110, -2.009012428, 12.158802680, -6.197120372, 9.010664824, 5.652323631, -12.044701090,
-0.967407392, 1.506607319, 1.154076109, -1.307765120, -3.433663463, 5.272253732, -9.733575197, 9.649853682, -1.171386070, 4.830994970,
-4.579567960, 1.288066793, -5.726860815, 4.483302224, 2.410065906, 0.976104957, -3.609195237, 5.866695462, -1.239660311, -7.710001405,
-7.262665497, 5.480506185, 4.663631998, 0.363894287, -1.811326158, -0.506302254, 4.277214996, -4.949724094, -6.156580898, 7.268054268,
5.313596297, -7.475636750, 5.962077111, -0.257662485, -2.366138970, 1.700393766, -5.305349679, 3.911268044, 5.759821575, -6.305390166,
-4.104941081, 10.264421450, -11.300767540, 12.054176580, -6.227745206, 2.091525447, -1.056821883, -1.333975456, 4.232249644, -4.622077368,
-2.961247945, 3.448961684, -7.158675007, 6.406735922, -6.514059266, 7.183645517, 2.820847867, -0.332339794, -3.133105721, 5.772761389,
11.667455430, 9.783810362, -4.813634608, 7.431879374, -2.042347962, 11.751539430, -11.966178390, 9.114149466, 3.316732944, 6.329493141,
0.849402562, 3.277425297, -7.025735603, 19.857356710, -4.387781796, -0.564523005, 10.846514210, -7.241087877, 8.632346433, 5.385693760,
1.865092045, 0.663456628, 1.439741709, -0.903063101, 8.146933783, -3.705109120, -78.618818310, -1.538554577, -4.021064560, -8.056139076,
9.037970305, -0.760954490, -6.968196665, 1.754220615, 1.299054010, 1.201201659, -4.358759250, -2.184578995, 0.361886430, -2.124097176,
7.900251100, -0.618609343, -12.395300170, 4.251095299, 4.913397600, -1.246714567, 17.149425990, -2.829759486, -3.169680239, -0.627060462,
4.344205545, -3.770631817, -8.359496416, 7.169040896, -7.703704106, 10.388878100, -1.952135409, -1.885118176, 1.636074186, -7.978317475
};

// Smaller arena to avoid static RAM overflow
constexpr int kTensorArenaSize = 48 * 1024;

// Dynamic allocation instead of big global static arrays
uint8_t* tensor_arena = nullptr;
float* imu_buf = nullptr;

tflite::MicroErrorReporter micro_error_reporter;
tflite::ErrorReporter* error_reporter = &micro_error_reporter;
const tflite::Model* tfl_model = nullptr;
tflite::MicroInterpreter* interpreter = nullptr;
TfLiteTensor* input_tensor = nullptr;
TfLiteTensor* output_tensor = nullptr;
tflite::AllOpsResolver resolver;

static inline float rd(const float* arr, int i) {
  return pgm_read_float(&arr[i]);
}

void setLED(int pin, bool on) {
  digitalWrite(pin, on ? HIGH : LOW);
}

void setupBuffers() {
  tensor_arena = (uint8_t*)malloc(kTensorArenaSize);
  imu_buf = (float*)malloc(SEG_LEN * N_CH * sizeof(float));

  if (!tensor_arena) {
    Serial.println("tensor_arena malloc failed!");
    while (true) delay(1000);
  }

  if (!imu_buf) {
    Serial.println("imu_buf malloc failed!");
    while (true) delay(1000);
  }

  Serial.println("Dynamic buffers allocated.");
}

void setupModel() {
  tfl_model = tflite::GetModel(MODEL_DATA);

  static tflite::MicroInterpreter* static_interpreter = nullptr;
  static_interpreter = new tflite::MicroInterpreter(
    tfl_model,
    resolver,
    tensor_arena,
    kTensorArenaSize,
    error_reporter
  );
  interpreter = static_interpreter;

  TfLiteStatus alloc_status = interpreter->AllocateTensors();
  if (alloc_status != kTfLiteOk) {
    error_reporter->Report("AllocateTensors failed!");
    while (true) delay(1000);
  }

  input_tensor = interpreter->input(0);
  output_tensor = interpreter->output(0);

  Serial.println("Model ready.");
  Serial.print("Input bytes: ");
  Serial.println(input_tensor->bytes);
  Serial.print("Output bytes: ");
  Serial.println(output_tensor->bytes);
}

void runModel(float &p_adl, float &p_deploy, float &p_alert) {
  float in_scale = input_tensor->params.scale;
  int in_zero = input_tensor->params.zero_point;

  for (int t = 0; t < SEG_LEN; t++) {
    for (int c = 0; c < N_CH; c++) {
      float x = imu_buf[t * N_CH + c];
      float xn = (x - MODEL_MEAN[c]) / MODEL_STD[c];
      int q = (int)roundf(xn / in_scale + in_zero);
      if (q < -128) q = -128;
      if (q > 127)  q = 127;
      input_tensor->data.int8[t * N_CH + c] = (int8_t)q;
    }
  }

  if (interpreter->Invoke() != kTfLiteOk) {
    Serial.println("Invoke failed!");
    p_adl = p_deploy = p_alert = 0.0f;
    return;
  }

  float out_scale = output_tensor->params.scale;
  int out_zero = output_tensor->params.zero_point;

  p_adl    = (output_tensor->data.int8[0] - out_zero) * out_scale;
  p_deploy = (output_tensor->data.int8[1] - out_zero) * out_scale;
  p_alert  = (output_tensor->data.int8[2] - out_zero) * out_scale;
}

void loadReplayToBuffer() {
  for (int i = 0; i < SEG_LEN; i++) {
    imu_buf[i * N_CH + 0] = rd(ACCX_CASE, i);
    imu_buf[i * N_CH + 1] = rd(ACCY_CASE, i);
    imu_buf[i * N_CH + 2] = rd(ACCZ_CASE, i);
    imu_buf[i * N_CH + 3] = rd(GYRX_CASE, i);
    imu_buf[i * N_CH + 4] = rd(GYRY_CASE, i);
    imu_buf[i * N_CH + 5] = rd(GYRZ_CASE, i);
  }
}

void printReplayPreview() {
  Serial.println("First 10 replay samples:");
  for (int i = 0; i < 10; i++) {
    Serial.print("i="); Serial.print(i);
    Serial.print(" ax="); Serial.print(imu_buf[i * N_CH + 0], 3);
    Serial.print(" ay="); Serial.print(imu_buf[i * N_CH + 1], 3);
    Serial.print(" az="); Serial.print(imu_buf[i * N_CH + 2], 3);
    Serial.print(" gx="); Serial.print(imu_buf[i * N_CH + 3], 3);
    Serial.print(" gy="); Serial.print(imu_buf[i * N_CH + 4], 3);
    Serial.print(" gz="); Serial.println(imu_buf[i * N_CH + 5], 3);
  }
}

void setup() {
  Serial.begin(115200);
  delay(1000);

  pinMode(PIN_DEPLOY, OUTPUT);
  pinMode(PIN_BRAKE, OUTPUT);
  pinMode(PIN_ALERT, OUTPUT);

  setLED(PIN_DEPLOY, false);
  setLED(PIN_BRAKE, false);
  setLED(PIN_ALERT, false);

  setupBuffers();
  setupModel();
  loadReplayToBuffer();
  printReplayPreview();

  Serial.println("=== ESP32 REAL MODEL + ENCODER REPLAY ===");
  Serial.print("Replay length = ");
  Serial.println(N);
}

void loop() {
  uint32_t t0 = micros();

  float encAccEma  = 0.0f;
  float encAccPeak = 0.0f;
  bool brake = false;
  bool brake_seen = false;
  uint32_t brake_t = 0;
  uint32_t brake_hold_until_ms = 0;

  for (int i = 0; i < N; i++) {
    uint32_t step_start = micros();

    float encSpeed = rd(ENC_SPEED, i);
    float encAcc   = rd(ENC_ACC_RAW, i);

    encAccEma  = ENC_ACC_EMA_ALPHA * encAcc + (1.0f - ENC_ACC_EMA_ALPHA) * encAccEma;
    encAccPeak = max(fabs(encAccEma), encAccPeak * ENC_PEAK_DECAY);

    bool brake_event = (encSpeed > BRAKE_SPEED_THR) || (encAccPeak > BRAKE_ACC_THR);
    if (brake_event) {
      brake_hold_until_ms = millis() + BRAKE_HOLD_MS;
    }
    brake = (millis() < brake_hold_until_ms);

    setLED(PIN_BRAKE, brake);

    if (brake && !brake_seen) {
      brake_seen = true;
      brake_t = micros() - t0;
    }

    uint32_t elapsed = micros() - step_start;
    if (elapsed < DT_US) delayMicroseconds(DT_US - elapsed);
  }

  float p_adl = 0.0f, p_deploy = 0.0f, p_alert = 0.0f;
  runModel(p_adl, p_deploy, p_alert);

  bool deploy = false;
  bool alert = false;

  if (p_alert >= ALERT_THR) {
    alert = true;
    deploy = false;
  } else if (p_deploy >= DEPLOY_THR) {
    alert = false;
    deploy = true;
  } else {
    alert = false;
    deploy = false;
  }

  setLED(PIN_DEPLOY, deploy);
  setLED(PIN_ALERT, alert);

  Serial.println();
  Serial.println("=== Model Output ===");
  Serial.print("P(ADL)    = "); Serial.println(p_adl, 4);
  Serial.print("P(DEPLOY) = "); Serial.println(p_deploy, 4);
  Serial.print("P(ALERT)  = "); Serial.println(p_alert, 4);

  Serial.println();
  Serial.println("=== Final Decision ===");
  Serial.print("Deploy legs? "); Serial.println(deploy ? "YES" : "NO");
  Serial.print("Emergency alert? "); Serial.println(alert ? "YES" : "NO");
  Serial.print("Brake seen? "); Serial.println(brake_seen ? "YES" : "NO");
  if (brake_seen) {
    Serial.print("Brake first trigger (ms): ");
    Serial.println(brake_t / 1000.0f, 2);
  }

  Serial.println();
  Serial.println("Done.");
  while (true) delay(1000);
}