#include "defines.h"

#include "CTfLiteClass.h"
#include "ClassLogFile.h"
#include "Helper.h"
#include "psram.h"
#include "esp_log.h"

#include <sys/stat.h>

static const char *TAG = "TFLITE";

CTfLiteClass::CTfLiteClass()
{
  model = nullptr;
  modelfile = NULL;
  interpreter = nullptr;
  input = nullptr;
  output = nullptr;
  kTensorArenaSize = TENSOR_ARENA_SIZE;
  tensor_arena = (uint8_t *)psram_get_shared_tensor_arena_memory();
}

CTfLiteClass::~CTfLiteClass()
{
  delete interpreter;
  psram_free_shared_tensor_arena_and_model_memory();
}

bool CTfLiteClass::MakeStaticResolver(void)
{
  if (resolver.AddFullyConnected() != kTfLiteOk)
  {
    ESP_LOGE(TAG, "load AddFullyConnected() failed");
    return false;
  }
  if (resolver.AddReshape() != kTfLiteOk)
  {
    ESP_LOGE(TAG, "load AddReshape() failed");
    return false;
  }
  if (resolver.AddSoftmax() != kTfLiteOk)
  {
    ESP_LOGE(TAG, "load AddSoftmax() failed");
    return false;
  }
  if (resolver.AddConv2D() != kTfLiteOk)
  {
    ESP_LOGE(TAG, "load AddConv2D() failed");
    return false;
  }
  if (resolver.AddMaxPool2D() != kTfLiteOk)
  {
    ESP_LOGE(TAG, "load AddMaxPool2D() failed");
    return false;
  }
  if (resolver.AddQuantize() != kTfLiteOk)
  {
    ESP_LOGE(TAG, "load AddQuantize() failed");
    return false;
  }
  if (resolver.AddMul() != kTfLiteOk)
  {
    ESP_LOGE(TAG, "load AddMul() failed");
    return false;
  }
  if (resolver.AddAdd() != kTfLiteOk)
  {
    ESP_LOGE(TAG, "load AddAdd() failed");
    return false;
  }
  if (resolver.AddLeakyRelu() != kTfLiteOk)
  {
    ESP_LOGE(TAG, "load AddLeakyRelu() failed");
    return false;
  }
  if (resolver.AddDequantize() != kTfLiteOk)
  {
    ESP_LOGE(TAG, "load AddDequantize() failed");
    return false;
  }

  return true;
}

float CTfLiteClass::GetOutputValue(int nr)
{
  TfLiteTensor *output2 = interpreter->output(0);

  int numer_output = output2->dims->data[1];
  if ((nr + 1) > numer_output)
  {
    return -1000;
  }

  return output2->data.f[nr];
}

int CTfLiteClass::GetClassFromImageBasis(CImageBasis *rs)
{
  if (!LoadInputImageBasis(rs))
  {
    return -1000;
  }

  Invoke();

  return GetOutClassification();
}

int CTfLiteClass::GetOutClassification(int _von, int _bis)
{
  TfLiteTensor *output2 = interpreter->output(0);

  float zw_max;
  float zw;
  int zw_class;

  if (output2 == NULL)
  {
    return -1;
  }

  int numeroutput = output2->dims->data[1];
  // ESP_LOGD(TAG, "number output neurons: %d", numeroutput);

  if (_bis == -1)
  {
    _bis = numeroutput - 1;
  }

  if (_von == -1)
  {
    _von = 0;
  }

  if (_bis >= numeroutput)
  {
    ESP_LOGD(TAG, "NUMBER OF OUTPUT NEURONS does not match required classification!");
    return -1;
  }

  zw_max = output2->data.f[_von];
  zw_class = _von;
  for (int i = _von + 1; i <= _bis; ++i)
  {
    zw = output2->data.f[i];
    if (zw > zw_max)
    {
      zw_max = zw;
      zw_class = i;
    }
  }
  return (zw_class - _von);
}

void CTfLiteClass::GetInputDimension(bool silent = false)
{
  TfLiteTensor *input2 = interpreter->input(0);

  int numdim = input2->dims->size;
  if (!silent)
  {
    ESP_LOGD(TAG, "NumDimension: %d", numdim);
  }

  int sizeofdim;
  for (int j = 0; j < numdim; ++j)
  {
    sizeofdim = input2->dims->data[j];
    if (!silent)
    {
      ESP_LOGD(TAG, "SizeOfDimension %d: %d", j, sizeofdim);
    }
    if (j == 1)
    {
      im_height = sizeofdim;
    }
    if (j == 2)
    {
      im_width = sizeofdim;
    }
    if (j == 3)
    {
      im_channel = sizeofdim;
    }
  }
}

int CTfLiteClass::ReadInputDimenstion(int _dim)
{
  if (_dim == 0)
  {
    return im_width;
  }
  if (_dim == 1)
  {
    return im_height;
  }
  if (_dim == 2)
  {
    return im_channel;
  }

  return -1;
}

int CTfLiteClass::GetAnzOutPut(bool silent)
{
  TfLiteTensor *output2 = interpreter->output(0);

  int numdim = output2->dims->size;
  if (!silent)
  {
    ESP_LOGD(TAG, "NumDimension: %d", numdim);
  }

  int sizeofdim;
  for (int j = 0; j < numdim; ++j)
  {
    sizeofdim = output2->dims->data[j];
    if (!silent)
    {
      ESP_LOGD(TAG, "SizeOfDimension %d: %d", j, sizeofdim);
    }
  }

  float fo;

  // Process the inference results.
  int numeroutput = output2->dims->data[1];
  for (int i = 0; i < numeroutput; ++i)
  {
    fo = output2->data.f[i];
    if (!silent)
    {
      ESP_LOGD(TAG, "Result %d: %f", i, fo);
    }
  }
  return numeroutput;
}

void CTfLiteClass::Invoke()
{
  if (interpreter != nullptr)
  {
    interpreter->Invoke();
  }
}

bool CTfLiteClass::LoadInputImageBasis(CImageBasis *rs)
{
  unsigned int w = rs->width;
  unsigned int h = rs->height;
  unsigned char red, green, blue;
  // ESP_LOGD(TAG, "Image: %s size: %d x %d\n", _fn.c_str(), w, h);

  input_i = 0;
  float *input_data_ptr = (interpreter->input(0))->data.f;

  for (int y = 0; y < h; ++y)
  {
    for (int x = 0; x < w; ++x)
    {
      red = rs->GetPixelColor(x, y, 0);
      green = rs->GetPixelColor(x, y, 1);
      blue = rs->GetPixelColor(x, y, 2);

      *(input_data_ptr) = (float)red;
      input_data_ptr++;

      *(input_data_ptr) = (float)green;
      input_data_ptr++;

      *(input_data_ptr) = (float)blue;
      input_data_ptr++;
    }
  }

  return true;
}

bool CTfLiteClass::MakeAllocate(void)
{
  LogFile.WriteToFile(ESP_LOG_DEBUG, TAG, "CTfLiteClass::MakeAllocate");

  if (!MakeStaticResolver())
  {
    LogFile.WriteToFile(ESP_LOG_ERROR, TAG, "CTfLiteClass::MakeAllocate - resolver could not be loaded!");
    return false;
  }

  if (!model)
  {
    LogFile.WriteToFile(ESP_LOG_ERROR, TAG, "CTfLiteClass::MakeAllocate - no model loaded!");
    return false;
  }

  if (model->version() != TFLITE_SCHEMA_VERSION)
  {
    LogFile.WriteToFile(ESP_LOG_ERROR, TAG, "The selected model does not match the tflite schema version!");
    return false;
  }

  if (!tensor_arena)
  {
    LogFile.WriteToFile(ESP_LOG_ERROR, TAG, "CTfLiteClass::MakeAllocate - tensor_arena not allocate");
    return false;
  }

  interpreter = new tflite::MicroInterpreter(model, resolver, tensor_arena, kTensorArenaSize);
  // LogFile.WriteToFile(ESP_LOG_INFO, TAG, "Trying to load the model. If it crashes here, it ist most likely due to a corrupted model!");

  if (interpreter)
  {
    TfLiteStatus allocate_status = interpreter->AllocateTensors();
    if (allocate_status != kTfLiteOk)
    {
      LogFile.WriteToFile(ESP_LOG_ERROR, TAG, "AllocateTensors() failed");
      GetInputDimension();
      return false;
    }
  }
  else
  {
    LogFile.WriteToFile(ESP_LOG_ERROR, TAG, "new tflite::MicroInterpreter failed");
    LogFile.WriteHeapInfo("CTfLiteClass::MakeAllocate-new tflite::MicroInterpreter failed");
    return false;
  }

  return true;
}

long CTfLiteClass::GetFileSize(std::string filename)
{
  struct stat stat_buf;
  long rc = -1;

  FILE *pFile = fopen(filename.c_str(), "rb"); // previously only "rb

  if (pFile != NULL)
  {
    rc = stat(filename.c_str(), &stat_buf);
    fclose(pFile);
  }

  return (rc == 0 ? stat_buf.st_size : -1);
}

bool CTfLiteClass::ReadFileToModel(std::string filename)
{
  LogFile.WriteToFile(ESP_LOG_DEBUG, TAG, "CTfLiteClass::ReadFileToModel: " + filename);

  long size = GetFileSize(filename);
  if (size == -1)
  {
    LogFile.WriteToFile(ESP_LOG_ERROR, TAG, "Model file doesn't exist: " + filename + "!");
    return false;
  }
  else if (size > MAX_MODEL_SIZE)
  {
    LogFile.WriteToFile(ESP_LOG_ERROR, TAG, "Unable to load model '" + filename + "'! It does not fit in the reserved shared memory in PSRAM!");
    return false;
  }

  LogFile.WriteToFile(ESP_LOG_DEBUG, TAG, "Loading Model " + filename + " /size: " + std::to_string(size) + " bytes...");

  modelfile = (unsigned char *)psram_get_shared_model_memory();
  if (modelfile != NULL)
  {
    FILE *pFile = fopen(filename.c_str(), "rb"); // previously only "rb
    if (pFile != NULL)
    {
      fread(modelfile, 1, size, pFile);
      fclose(pFile);

      return true;
    }
    else
    {
      LogFile.WriteToFile(ESP_LOG_ERROR, TAG, "CTfLiteClass::ReadFileToModel: Model does not exist");
      return false;
    }
  }
  else
  {
    LogFile.WriteToFile(ESP_LOG_ERROR, TAG, "CTfLiteClass::ReadFileToModel: Can't allocate enough memory: " + std::to_string(size));
    LogFile.WriteHeapInfo("CTfLiteClass::ReadFileToModel");

    return false;
  }
}

bool CTfLiteClass::LoadModel(std::string filename)
{
  LogFile.WriteToFile(ESP_LOG_DEBUG, TAG, "CTfLiteClass::LoadModel");

  if (!ReadFileToModel(filename.c_str()))
  {
    return false;
  }

  model = tflite::GetModel(modelfile);

  if (model == nullptr)
  {
    return false;
  }

  return true;
}