/*
 * array_utils.hpp
 *
 * This file contains utility functions for working with arrays. These are
 * mostly convenience functions for setting up and inspecting data for testing.
 * They are not optimized and are not intended for use in performance-critical
 * code.
 *
 */

#ifndef ARRAY_UTILS_H
#define ARRAY_UTILS_H

#include <algorithm> // std::max_element
#include <array>
#include <cmath>
#include <random>
#include <stdexcept>
#include <string>
#include <utility>

#include "utils/logging.hpp"
#include "numeric_types/half.hpp"

namespace gpu {

static constexpr int kShowMaxRows = 8;
static constexpr int kShowMaxCols = 8;

/**
 * @brief Show a 2D array as a string, base implementation.
 *
 * @param a The array to show.
 * @param rows The number of rows in the array.
 * @param cols The number of columns in the array.
 * @param name The name of the array to show.
 * @return std::string The string representation of the array.
 * @code
 *   std::array<float, 4> a = {1.0, 2.0, 3.0, 4.0};
 *   printf("%s", show<float>(a.data(), 2, 2, "a").c_str());
 * @endcode
 */
template <typename numtype>
std::string show(const numtype *a, size_t rows, size_t cols,
                 const std::string &name = "") {
  std::string output = "\n";
  if (name != "") {
    output += "\n" + name + " (" + std::to_string(rows) + ", " +
              std::to_string(cols) + ")\n\n";
  } else {
    output +=
        "\n(" + std::to_string(rows) + ", " + std::to_string(cols) + ")\n\n";
  }
  // spacing as log10 of max value
  int spacing = 1;
  if constexpr (std::is_same<numtype, int>::value) {
    int max = *std::max_element(a, a + rows * cols);
    spacing = std::max(0, (int)log10(max + .01)) + 2;
  } else if constexpr (std::is_same<numtype, float>::value) {
    // spacing = std::max(0, (int)log10(max + .01)) + 1;
    spacing = 8; // scientific notation
  } else if constexpr (std::is_same<numtype, half>::value) {
    spacing = 8;
  } else {
    throw std::runtime_error("Unsupported number type for show()");
  }
  // print to stdout line break for each row
  for (size_t i = 0; i < rows; i++) {
    if (i == kShowMaxRows / 2 && rows > kShowMaxRows) {
      output += "...\n";
      i = rows - kShowMaxRows / 2;
    }
    for (size_t j = 0; j < cols; j++) {
      if (j == kShowMaxCols / 2 && cols > kShowMaxCols) {
        output += " ..";
        j = cols - kShowMaxCols / 2;
      }
      char buffer[50];
      if constexpr (std::is_same<numtype, int>::value) {
        snprintf(buffer, spacing, "%*d", spacing, a[i * cols + j]);
      } else if constexpr (std::is_same<numtype, float>::value) {
        if (std::abs(a[i * cols + j]) < 1000 &&
                std::abs(a[i * cols + j]) > 0.01 ||
            a[i * cols + j] == 0.0) {
          snprintf(buffer, 16, "%9.2f", a[i * cols + j]);
        } else
          snprintf(buffer, 16, "%10.2e", a[i * cols + j]);
      } else if constexpr (std::is_same<numtype, half>::value) {
	float tmp = halfToFloat(a[i * cols + j]);
        if (std::abs(tmp) < 1000 &&
                std::abs(tmp) > 0.01 ||
            tmp == 0.0) {
          snprintf(buffer, 16, "%9.2f", tmp);
        } else
          snprintf(buffer, 16, "%10.2e", tmp);
      } else {
        throw std::runtime_error("Unsupported number type for show()");
      }
      output += buffer;
    }
    output += "\n";
  }
  output += "\n";
  return output;
}

/**
 * @brief Overload of `show()` for std::array.
 *
 * @param a The array to show.
 * @param name The name of the array to show.
 * @return std::string The string representation of the array.
 * @code
 *  std::array<float, 4> a = {1.0, 2.0, 3.0, 4.0};
 *  printf("%s", show<float>(a, "a").c_str());
 * @endcode
 */
template <typename numtype, size_t rows, size_t cols>
std::string show(const std::array<numtype, rows * cols> &a,
                 const std::string &name = "") {
  return show<numtype>(a.data(), rows, cols, name);
}

/**
 * @brief Overload of `show()` for float std::array.
 * @param a The array to show.
 * @param name The name of the array to show.
 * @return std::string The string representation of the array.
 *
 * @code
 * std::array<float, 4> a = {1.0, 2.0, 3.0, 4.0};
 * printf("%s", show(a, "a").c_str());
 * @endcode
 @  
 */
template <size_t rows, size_t cols>
std::string show(const std::array<float, rows * cols> &a,
                 const std::string &name = "") {
  return show<float, rows, cols>(a, name);
}

/**
 * @brief Populate the array with a range of values. This is mostly for testing
 * purposes.
 * @param input The array to populate.
 * @param N The number of elements in the array.
 * @param start The starting value.
 * @param step The step size.
 */
void range(float *input, size_t N, float start = 0.0, float step = 1.0) {
  // TODO(avh): currently unused - check
  float curr = start;
  for (size_t i = 0; i < N; i++) {
    input[i] = curr;
    curr += step;
  }
}

/**
 * @brief Overload of `range()` for std::array.
 * @param input The array to populate.
 * @param start The starting value.
 * @param step The step size.
 */
template <size_t N>
void range(std::array<float, N> &input, float start = 0.0, float step = 1.0) {
  float curr = start;
  for (size_t i = start; i < N; i++) {
    input[i] = curr;
    curr += step;
  }
}

/**
 * @brief Populate the array with random integers.
 * @param a The array to populate.
 * @param N The number of elements in the array.
 * @param gen The random number generator.
 * @param min The minimum value for the random integers.
 * @param max The maximum value for the random integers.
 */
void randint(float *a, size_t N, std::mt19937 &gen, int min = -1, int max = 1) {
  std::uniform_int_distribution<> dist(min, max);
  for (int i = 0; i < N; i++) {
    a[i] = static_cast<float>(dist(gen));
  }
}

/**
 * @brief Overload of `randint()` for std::array.
 * @param a The array to populate.
 * @param gen The random number generator.
 * @param min The minimum value for the random integers.
 * @param max The maximum value for the random integers.
 */
template <typename numtype, size_t size>
void randint(std::array<numtype, size> &a, std::mt19937 &gen, int min = -1,
             int max = 1) {
  std::uniform_int_distribution<> dist(min, max);
  for (int i = 0; i < size; i++) {
    a[i] = static_cast<numtype>(dist(gen));
  }
}

/**
 * @brief Populate the array with random floats, generated from a Gaussian distribution.
 * @param a The array to populate.
 * @param N The number of elements in the array.
 * @param gen The random number generator.
 * @param mean The mean of the Gaussian distribution.
 * @param std The standard deviation of the Gaussian distribution.
 */
inline void randn(float *a, size_t N, std::mt19937 &gen, float mean = 0.0,
                  float std = 1.0) {
  std::normal_distribution<float> dist(mean, std);
  for (int i = 0; i < N; i++) {
    a[i] = static_cast<float>(dist(gen));
  }
}

inline void randn(half *a, size_t N, std::mt19937 &gen, float mean = 0.0,
                  float std = 1.0) {
  std::normal_distribution<float> dist(mean, std);
  for (int i = 0; i < N; i++) {
    a[i] = halfFromFloat(dist(gen));
  }
}

/**
 * @brief Overload of `randn()` for std::array.
 * @param a The array to populate.
 * @param gen The random number generator.
 * @param mean The mean of the Gaussian distribution.
 * @param std The standard deviation of the Gaussian distribution.
 */
template <size_t size>
void randn(std::array<float, size> &a, std::mt19937 &gen, float mean = 0.0,
           float std = 1.0) {
  std::normal_distribution<float> dist(mean, std);
  for (int i = 0; i < size; i++) {
    a[i] = static_cast<float>(dist(gen));
  }
}

/**
 * @brief Populate a square matrix with the identity matrix.
 * @param a The array to populate.
 * @param N The number of rows and columns in the square matrix.
 */
inline void eye(float *a, size_t N) {
  for (size_t i = 0; i < N; i++) {
    for (size_t j = 0; j < N; j++) {
      a[i * N + j] = (i == j) ? 1.0 : 0.0;
    }
  }
}

// Note transformation operations here are purely for testing - they are not
// optimized to be used in hot paths.

/**
 * @brief Transpose a matrix.
 * @param input The input matrix.
 * @param output The output matrix.
 * @param M The number of rows in the input matrix.
 * @param N The number of columns in the input matrix.
 */
inline void transpose(float *input, float *output, size_t M, size_t N) {
  for (size_t i = 0; i < M; i++) {
    for (size_t j = 0; j < N; j++) {
      output[j * M + i] = input[i * N + j];
    }
  }
}

inline void transpose(half *input, half *output, size_t M, size_t N) {
  for (size_t i = 0; i < M; i++) {
    for (size_t j = 0; j < N; j++) {
      output[j * M + i] = input[i * N + j];
    }
  }
}

/**
 * @brief Flip a matrix horizontally or vertically.
 * @param a The matrix to flip.
 * @param R The number of rows in the matrix.
 * @param C The number of columns in the matrix.
 * @param horizontal Whether to flip horizontally (true) or vertically (false).
 */
inline void flip(float *a, size_t R, size_t C, bool horizontal = true) {
  if (horizontal) {
    for (size_t i = 0; i < R; i++) {
      for (size_t j = 0; j < C / 2; j++) {
        std::swap(a[i * C + j], a[i * C + C - j - 1]);
      }
    }
  } else {
    for (size_t i = 0; i < R / 2; i++) {
      for (size_t j = 0; j < C; j++) {
        std::swap(a[i * C + j], a[(R - i - 1) * C + j]);
      }
    }
  }
}

/**
 * @brief Determine if the values of two arrays are close to each other.
 * @param a The first array.
 * @param b The second array.
 * @param n The number of elements in the arrays.
 * @param tol The tolerance for closeness.
 * @return bool True if the arrays are close, false otherwise.
 */
inline bool isclose(float *a, float *b, size_t n, float tol = 1e-3) {
  for (size_t i = 0; i < n; i++) {
    if (std::abs(a[i] - b[i]) > tol || std::isnan(a[i]) || std::isnan(b[i])) {
      LOG(kDefLog, kError, "Mismatch at index %d: %f != %f", i, a[i], b[i]);
      return false;
    }
  }
  return true;
}

inline bool isclose(half *a, half *b, size_t n, float tol = 1) {
  for (size_t i = 0; i < n; i++) {
    float ai = halfToFloat(a[i]);
    float bi = halfToFloat(b[i]);
    if (std::abs(ai - bi) > tol || std::isnan(ai) || std::isnan(bi)) {
      LOG(kDefLog, kError, "Mismatch at index %d: %f != %f", i, ai, bi);
      return false;
    }
  }
  return true;
}

} // namespace gpu

#endif