#include "phash.h"

#include <cmath>
#include <cstring>
#include <vector>

bool decode_image_bytes(const uint8_t* data, int len, std::vector<uint8_t>& rgba, int& width,
                        int& height);
void resize_rgba_to_rgb24x24(const uint8_t* rgba, int width, int height, uint8_t* out_rgb);

namespace {

constexpr int HASH_SIZE = PHASH_SIZE;
constexpr int FREQ_SIZE = HASH_SIZE / 4;

void compute_dct(const double* pixels, double* out) {
  const int n = HASH_SIZE;
  for (int u = 0; u < n; u++) {
    for (int v = 0; v < n; v++) {
      double sum = 0.0;
      for (int y = 0; y < n; y++) {
        for (int x = 0; x < n; x++) {
          sum += pixels[y * n + x] *
                 std::cos((2.0 * x + 1.0) * u * M_PI / (2.0 * n)) *
                 std::cos((2.0 * y + 1.0) * v * M_PI / (2.0 * n));
        }
      }
      const double cu = (u == 0) ? 1.0 / std::sqrt(2.0) : 1.0;
      const double cv = (v == 0) ? 1.0 / std::sqrt(2.0) : 1.0;
      out[u * n + v] = (2.0 * cu * cv / n) * sum;
    }
  }
}

int hash_from_rgb24x24_internal(const uint8_t* rgb, char* out_hash, int out_cap) {
  if (!rgb || !out_hash || out_cap < PHASH_BITS + 1) {
    return 1;
  }

  double r[HASH_SIZE * HASH_SIZE];
  double g[HASH_SIZE * HASH_SIZE];
  double b[HASH_SIZE * HASH_SIZE];

  for (int i = 0; i < HASH_SIZE * HASH_SIZE; i++) {
    r[i] = rgb[i * 3 + 0];
    g[i] = rgb[i * 3 + 1];
    b[i] = rgb[i * 3 + 2];
  }

  double dct_r[HASH_SIZE * HASH_SIZE];
  double dct_g[HASH_SIZE * HASH_SIZE];
  double dct_b[HASH_SIZE * HASH_SIZE];
  compute_dct(r, dct_r);
  compute_dct(g, dct_g);
  compute_dct(b, dct_b);

  double sum_r = 0, sum_g = 0, sum_b = 0;
  int count = 0;
  for (int y = 0; y < FREQ_SIZE; y++) {
    for (int x = 0; x < FREQ_SIZE; x++) {
      if (x == 0 && y == 0) continue;
      const int pos = y * HASH_SIZE + x;
      sum_r += dct_r[pos];
      sum_g += dct_g[pos];
      sum_b += dct_b[pos];
      count++;
    }
  }
  const double avg_r = sum_r / count;
  const double avg_g = sum_g / count;
  const double avg_b = sum_b / count;

  int idx = 0;
  for (int y = 0; y < FREQ_SIZE; y++) {
    for (int x = 0; x < FREQ_SIZE; x++) {
      if (x == 0 && y == 0) continue;
      const int pos = y * HASH_SIZE + x;
      out_hash[idx++] = dct_r[pos] > avg_r ? '1' : '0';
      out_hash[idx++] = dct_g[pos] > avg_g ? '1' : '0';
      out_hash[idx++] = dct_b[pos] > avg_b ? '1' : '0';
    }
  }
  out_hash[PHASH_BITS] = '\0';
  return 0;
}

}  // namespace

int phash_from_rgb24x24(const uint8_t* rgb, char* out_hash, int out_cap) {
  return hash_from_rgb24x24_internal(rgb, out_hash, out_cap);
}

int phash_from_rgba(const uint8_t* rgba, int width, int height, char* out_hash, int out_cap) {
  if (!rgba || width <= 0 || height <= 0) {
    return 2;
  }
  uint8_t rgb[HASH_SIZE * HASH_SIZE * 3];
  resize_rgba_to_rgb24x24(rgba, width, height, rgb);
  return hash_from_rgb24x24_internal(rgb, out_hash, out_cap);
}

int phash_from_bytes(const uint8_t* data, int len, char* out_hash, int out_cap) {
  if (!data || len <= 0) {
    return 2;
  }
  std::vector<uint8_t> rgba;
  int width = 0, height = 0;
  if (!decode_image_bytes(data, len, rgba, width, height)) {
    return 3;
  }
  return phash_from_rgba(rgba.data(), width, height, out_hash, out_cap);
}
