// Copyright 2025 RnD Center "ELVEES", JSC

#include <elcore50-signal-lib/fft.h>

float maximum_norm(const float *dsp_res, const float *ref_res, const int size) {
  float diff = fabs(ref_res[0] - dsp_res[0]);
  for (size_t i = 1; i < size * 2; ++i) {
    float diff_i = fabs(ref_res[i] - dsp_res[i]);
    if (diff_i > diff) diff = diff_i;
  }

  return diff;
}

int main() {
  // Количество строк матрицы
  int16_t rows_number = 512;
  // Количество столбцов матрицы
  int16_t cols_number = 512;
  // Общее количество комплексных отсчетов
  const int32_t fft_size = rows_number * cols_number;

  // Выходной вектор для типа fractional16
  // Выравнивание нужно только при вычислениях во внешней памяти (DDR)
  int16_t *fdst = memalign(fft_size * 2 * sizeof(int16_t), fft_size * 2 * sizeof(int16_t));
  // Входной вектор для типа fractional16
  int16_t *fsrc = memalign(fft_size * 2 * sizeof(int16_t), fft_size * 2 * sizeof(int16_t));
  // Вектор точного решения
  float *fetalon = memalign(fft_size * 2 * sizeof(float), fft_size * 2 * sizeof(float));

  // Создание объекта преобразования
  fft_t fft_obj;
  fft_obj.src = fsrc;
  fft_obj.dst = fdst;
  fft_obj.int_mem = SL_XYRAM;
  fft_obj.dir = SL_IFFT;
  fft_obj.dim = SL_2D;
  int retval = fft_plan_2d(&fft_obj, rows_number, cols_number, SL_FRACTIONAL);
  if (retval) {
    print_error_message(retval);
    // Освобождение ресурсов
    fft_free_sources(&fft_obj);
    free(fdst);
    free(fsrc);
    free(fetalon);
    return 1;
  }

  // Частота входного сигнала
  const int16_t frequency1 = 2;
  const int16_t frequency2 = 1;
  // Входной сигнал
  memset(fsrc, 0.0, fft_size * 2 * sizeof(short));
  fsrc[2 * (frequency1 * cols_number + frequency2) + 1] = 0.99 * 0.99 * 0x8000;

  // Вычисление обратного БПФ
  retval = fft_execute(&fft_obj);
  if (retval) {
    print_error_message(retval);
    // Освобождение ресурсов
    fft_free_sources(&fft_obj);
    free(fdst);
    free(fsrc);
    free(fetalon);
    return 1;
  }

  // Подготовка для вычисления прямого БПФ
  float *temp = malloc(fft_size * 2 * sizeof(float));
  if (fft_obj.int_mem & fft_obj.big_size) {
    // Преобразование результата типа fract16 в тип float32
    fract16_to_float32(fsrc, temp, fft_size);
    // Транспонирование результата обратного БПФ
    float *temp_T = malloc(fft_size * 2 * sizeof(float));
    for (size_t i = 0; i < cols_number * 2; i += 2) {
      for (size_t j = 0; j < rows_number * 2; j += 2) {
        temp_T[j * cols_number + i] = temp[i * rows_number + j];
        temp_T[j * cols_number + i + 1] = temp[i * rows_number + j + 1];
      }
    }
    // Преобразование результата типа float32 обратно в тип fract16
    float32_to_fract16(temp_T, fsrc, fft_size);
    free(temp_T);
  }

  // Вычисление прямого БПФ
  fft_obj.src = fsrc;
  fft_obj.dst = fdst;
  fft_obj.dir = SL_FFT;
  retval = fft_execute(&fft_obj);
  if (retval) {
    print_error_message(retval);
    // Освобождение ресурсов
    fft_free_sources(&fft_obj);
    free(fdst);
    free(fsrc);
    free(fetalon);
    free(temp);
    return 1;
  }

  // Проверка результата
  // Преобразование результата типа fract16 в тип float32
  fract16_to_float32(fsrc, temp, fft_size);
  memset(fetalon, 0.0, fft_size * 2 * sizeof(float));
  if (fft_obj.int_mem && fft_obj.big_size)
    fetalon[2 * (frequency2 * rows_number + frequency1) + 1] = 0.99 * 0.99;
  else
    fetalon[2 * (frequency1 * cols_number + frequency2) + 1] = 0.99 * 0.99;
  float norm = maximum_norm(temp, fetalon, fft_size);
  printf("fract16 2d sample result:\n");
  printf("rows = %d, cols = %d, ", rows_number, cols_number);
  if (norm < 1e-2)
    printf("passed!\n");
  else
    printf("failed!\n");

  // Освобождение ресурсов
  fft_free_sources(&fft_obj);
  free(temp);
  free(fdst);
  free(fsrc);
  free(fetalon);
  return 0;
}
