// Copyright (c) the JPEG XL Project Authors. All rights reserved.
//
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

#include "lib/jpegli/color_quantize.h"

#include <algorithm>
#include <cmath>
#include <cstddef>
#include <cstdint>
#include <cstring>
#include <limits>
#include <memory>
#include <unordered_map>
#include <vector>

#include "lib/jpegli/common.h"
#include "lib/jpegli/common_internal.h"
#include "lib/jpegli/decode_internal.h"
#include "lib/jpegli/error.h"
#include "lib/jpegli/memory_manager.h"

namespace jpegli {

namespace {

constexpr int kNumColorCellBits[kMaxComponents] = {3, 4, 3, 3};
constexpr int kCompW[kMaxComponents] = {2, 3, 1, 1};

int Pow(int a, int b) {
  int r = 1;
  for (int i = 0; i < b; ++i) {
    r *= a;
  }
  return r;
}

int ComponentOrder(j_decompress_ptr cinfo, int i) {
  if (cinfo->out_color_components == 3) {
    return i < 2 ? 1 - i : i;
  }
  return i;
}

int GetColorComponent(int i, int N) {
  return (i * 255 + (N - 1) / 2) / (N - 1);
}

}  // namespace

void ChooseColorMap1Pass(j_decompress_ptr cinfo) {
  jpeg_decomp_master* m = cinfo->master;
  int components = cinfo->out_color_components;
  int desired = std::min(cinfo->desired_number_of_colors, 256);
  int num = 1;
  while (Pow(num + 1, components) <= desired) {
    ++num;
  }
  if (num == 1) {
    JPEGLI_ERROR("Too few colors (%d) in requested colormap", desired);
  }
  int actual = Pow(num, components);
  for (int i = 0; i < components; ++i) {
    m->num_colors_[i] = num;
  }
  while (actual < desired) {
    int total = actual;
    for (int i = 0; i < components; ++i) {
      int c = ComponentOrder(cinfo, i);
      int new_total = (actual / m->num_colors_[c]) * (m->num_colors_[c] + 1);
      if (new_total <= desired) {
        ++m->num_colors_[c];
        actual = new_total;
      }
    }
    if (actual == total) {
      break;
    }
  }
  cinfo->actual_number_of_colors = actual;
  cinfo->colormap = (*cinfo->mem->alloc_sarray)(
      reinterpret_cast<j_common_ptr>(cinfo), JPOOL_IMAGE, actual, components);
  int next_color[kMaxComponents] = {0};
  for (int i = 0; i < actual; ++i) {
    for (int c = 0; c < components; ++c) {
      cinfo->colormap[c][i] =
          GetColorComponent(next_color[c], m->num_colors_[c]);
    }
    int c = components - 1;
    while (c > 0 && next_color[c] + 1 == m->num_colors_[c]) {
      next_color[c--] = 0;
    }
    ++next_color[c];
  }
  if (!m->colormap_lut_) {
    m->colormap_lut_ = Allocate<uint8_t>(cinfo, components * 256, JPOOL_IMAGE);
  }
  int stride = actual;
  for (int c = 0; c < components; ++c) {
    int N = m->num_colors_[c];
    stride /= N;
    for (int i = 0; i < 256; ++i) {
      int index = ((2 * i - 1) * (N - 1) + 254) / 510;
      m->colormap_lut_[c * 256 + i] = index * stride;
    }
  }
}

namespace {

// 2^13 priority levels for the PQ seems to be a good compromise between
// accuracy, running time and stack space usage.
const int kMaxPriority = 1 << 13;
const int kMaxLevel = 3;

// This function is used in the multi-resolution grid to be able to compute
// the keys for the different resolutions by just shifting the first key.
inline int InterlaceBitsRGB(uint8_t r, uint8_t g, uint8_t b) {
  int z = 0;
  for (int i = 0; i < 7; ++i) {
    z += (r >> 5) & 4;
    z += (g >> 6) & 2;
    z += (b >> 7);
    z <<= 3;
    r <<= 1;
    g <<= 1;
    b <<= 1;
  }
  z += (r >> 5) & 4;
  z += (g >> 6) & 2;
  z += (b >> 7);
  return z;
}

// This function will compute the actual priorities of the colors based on
// the current distance from the palette, the population count and the signals
// from the multi-resolution grid.
inline int Priority(int d, int n, const int* density, const int* radius) {
  int p = d * n;
  for (int level = 0; level < kMaxLevel; ++level) {
    if (d > radius[level]) {
      p += density[level] * (d - radius[level]);
    }
  }
  return std::min(kMaxPriority - 1, p >> 4);
}

inline int ColorIntQuadDistanceRGB(uint8_t r1, uint8_t g1, uint8_t b1,
                                   uint8_t r2, uint8_t g2, uint8_t b2) {
  // weights for the intensity calculation
  static constexpr int ired = 2;
  static constexpr int igreen = 5;
  static constexpr int iblue = 1;
  // normalization factor for the intensity calculation (2^ishift)
  static constexpr int ishift = 3;
  const int rd = r1 - r2;
  const int gd = g1 - g2;
  const int bd = b1 - b2;
  const int id = ired * rd + igreen * gd + iblue * bd;
  return rd * rd + gd * gd + bd * bd + ((id * id) >> (2 * ishift));
}

inline int ScaleQuadDistanceRGB(int d) {
  return static_cast<int>(std::lround(sqrt(d * 0.25)));
}

// The function updates the minimal distances, the clustering and the
// quantization error after the insertion of the new color into the palette.
void AddToRGBPalette(const uint8_t* red, const uint8_t* green,
                     const uint8_t* blue,
                     const int* count,  // histogram of colors
                     const int index,   // index of color to be added
                     const int k,       // size of current palette
                     const int n,       // number of colors
                     int* dist,         // array of distances from palette
                     int* cluster,      // mapping of color indices to palette
                     int* center,       // the inverse mapping
                     int64_t* error) {  // measure of the quantization error
  center[k] = index;
  cluster[index] = k;
  *error -=
      static_cast<int64_t>(dist[index]) * static_cast<int64_t>(count[index]);
  dist[index] = 0;
  for (int j = 0; j < n; ++j) {
    if (dist[j] > 0) {
      const int d = ColorIntQuadDistanceRGB(
          red[index], green[index], blue[index], red[j], green[j], blue[j]);
      if (d < dist[j]) {
        *error += static_cast<int64_t>((d - dist[j])) *
                  static_cast<int64_t>(count[j]);
        dist[j] = d;
        cluster[j] = k;
      }
    }
  }
}

struct RGBPixelHasher {
  // A quick but good-enough hash to get 24 bits of RGB into the lower 12 bits.
  size_t operator()(uint32_t a) const { return (a ^ (a >> 12)) * 0x9e3779b9; }
};

struct WangHasher {
  // Thomas Wang's Hash.  Nearly perfect and still quite fast.  Above (for
  // pixels) we use a simpler hash because the number of hash calls is
  // proportional to the number of pixels and that hash dominates; we want the
  // cost to be minimal and we start with a large table.  We can use a better
  // hash for the histogram since the number of hash calls is proportional to
  // the number of unique colors in the image, which is hopefully much smaller.
  // Note that the difference is slight; e.g. replacing RGBPixelHasher with
  // WangHasher only slows things down by 5% on an Opteron.
  size_t operator()(uint32_t a) const {
    a = (a ^ 61) ^ (a >> 16);
    a = a + (a << 3);
    a = a ^ (a >> 4);
    a = a * 0x27d4eb2d;
    a = a ^ (a >> 15);
    return a;
  }
};

// Build an index of all the different colors in the input
// image. To do this we map the 24 bit RGB representation of the colors
// to a unique integer index assigned to the different colors in order of
// appearance in the image.  Return the number of unique colors found.
// The colors are pre-quantized to 3 * 6 bits precision.
int BuildRGBColorIndex(const uint8_t* const image, int const num_pixels,
                       int* const count, uint8_t* const red,
                       uint8_t* const green, uint8_t* const blue) {
  // Impossible because rgb are in the low 24 bits, and the upper 8 bits is 0.
  const uint32_t impossible_pixel_value = 0x10000000;
  std::unordered_map<uint32_t, int, RGBPixelHasher> index_map(1 << 12);
  std::unordered_map<uint32_t, int, RGBPixelHasher>::iterator index_map_lookup;
  const uint8_t* imagep = &image[0];
  uint32_t prev_pixel = impossible_pixel_value;
  int index = 0;
  int n = 0;
  for (int i = 0; i < num_pixels; ++i) {
    uint8_t r = ((*imagep++) & 0xfc) + 2;
    uint8_t g = ((*imagep++) & 0xfc) + 2;
    uint8_t b = ((*imagep++) & 0xfc) + 2;
    uint32_t pixel = (b << 16) | (g << 8) | r;
    if (pixel != prev_pixel) {
      prev_pixel = pixel;
      index_map_lookup = index_map.find(pixel);
      if (index_map_lookup != index_map.end()) {
        index = index_map_lookup->second;
      } else {
        index_map[pixel] = index = n++;
        red[index] = r;
        green[index] = g;
        blue[index] = b;
      }
    }
    ++count[index];
  }
  return n;
}

}  // namespace

void ChooseColorMap2Pass(j_decompress_ptr cinfo) {
  if (cinfo->out_color_space != JCS_RGB) {
    JPEGLI_ERROR("Two-pass quantizer must use RGB output color space.");
  }
  jpeg_decomp_master* m = cinfo->master;
  const size_t num_pixels = cinfo->output_width * cinfo->output_height;
  const int max_color_count = std::max<size_t>(num_pixels, 1u << 18);
  const int max_palette_size = cinfo->desired_number_of_colors;
  std::unique_ptr<uint8_t[]> red(new uint8_t[max_color_count]);
  std::unique_ptr<uint8_t[]> green(new uint8_t[max_color_count]);
  std::unique_ptr<uint8_t[]> blue(new uint8_t[max_color_count]);
  std::vector<int> count(max_color_count, 0);
  // number of colors
  int n = BuildRGBColorIndex(m->pixels_, num_pixels, count.data(), &red[0],
                             &green[0], &blue[0]);

  std::vector<int> dist(n, std::numeric_limits<int>::max());
  std::vector<int> cluster(n);
  std::vector<bool> in_palette(n, false);
  int center[256];
  int k = 0;  // palette size
  const int count_threshold = (num_pixels * 4) / max_palette_size;
  static constexpr int kAveragePixelErrorThreshold = 1;
  const int64_t error_threshold = num_pixels * kAveragePixelErrorThreshold;
  int64_t error = 0;  // quantization error

  int max_count = 0;
  int winner = 0;
  for (int i = 0; i < n; ++i) {
    if (count[i] > max_count) {
      max_count = count[i];
      winner = i;
    }
    if (!in_palette[i] && count[i] > count_threshold) {
      AddToRGBPalette(&red[0], &green[0], &blue[0], count.data(), i, k++, n,
                      dist.data(), cluster.data(), &center[0], &error);
      in_palette[i] = true;
    }
  }
  if (k == 0) {
    AddToRGBPalette(&red[0], &green[0], &blue[0], count.data(), winner, k++, n,
                    dist.data(), cluster.data(), &center[0], &error);
    in_palette[winner] = true;
  }

  // Calculation of the multi-resolution density grid.
  std::vector<int> density(n * kMaxLevel);
  std::vector<int> radius(n * kMaxLevel);
  std::unordered_map<uint32_t, int, WangHasher> histogram[kMaxLevel];
  for (int level = 0; level < kMaxLevel; ++level) {
    // This value is never used because key = InterlaceBitsRGB(...) >> 6
  }

  for (int i = 0; i < n; ++i) {
    if (!in_palette[i]) {
      const int key = InterlaceBitsRGB(red[i], green[i], blue[i]) >> 6;
      for (int level = 0; level < kMaxLevel; ++level) {
        histogram[level][key >> (3 * level)] += count[i];
      }
    }
  }
  for (int i = 0; i < n; ++i) {
    if (!in_palette[i]) {
      for (int level = 0; level < kMaxLevel; ++level) {
        const int mask = (4 << level) - 1;
        const int rd = std::max(red[i] & mask, mask - (red[i] & mask));
        const int gd = std::max(green[i] & mask, mask - (green[i] & mask));
        const int bd = std::max(blue[i] & mask, mask - (blue[i] & mask));
        radius[i * kMaxLevel + level] =
            ScaleQuadDistanceRGB(ColorIntQuadDistanceRGB(0, 0, 0, rd, gd, bd));
      }
      const int key = InterlaceBitsRGB(red[i], green[i], blue[i]) >> 6;
      if (kMaxLevel > 0) {
        density[i * kMaxLevel] = histogram[0][key] - count[i];
      }
      for (int level = 1; level < kMaxLevel; ++level) {
        density[i * kMaxLevel + level] =
            (histogram[level][key >> (3 * level)] -
             histogram[level - 1][key >> (3 * level - 3)]);
      }
    }
  }

  // Calculate the initial error now that the palette has been initialized.
  error = 0;
  for (int i = 0; i < n; ++i) {
    error += static_cast<int64_t>(dist[i]) * static_cast<int64_t>(count[i]);
  }

  std::unique_ptr<std::vector<int>[]> bucket_array(
      new std::vector<int>[kMaxPriority]);
  int top_priority = -1;
  for (int i = 0; i < n; ++i) {
    if (!in_palette[i]) {
      int priority = Priority(ScaleQuadDistanceRGB(dist[i]), count[i],
                              &density[i * kMaxLevel], &radius[i * kMaxLevel]);
      bucket_array[priority].push_back(i);
      top_priority = std::max(priority, top_priority);
    }
  }
  double error_accum = 0;
  while (top_priority >= 0 && k < max_palette_size) {
    if (error < error_threshold) {
      error_accum += std::min(error_threshold, error_threshold - error);
      if (error_accum >= 10 * error_threshold) {
        break;
      }
    }
    int i = bucket_array[top_priority].back();
    int priority = Priority(ScaleQuadDistanceRGB(dist[i]), count[i],
                            &density[i * kMaxLevel], &radius[i * kMaxLevel]);
    if (priority < top_priority) {
      bucket_array[priority].push_back(i);
    } else {
      AddToRGBPalette(&red[0], &green[0], &blue[0], count.data(), i, k++, n,
                      dist.data(), cluster.data(), &center[0], &error);
    }
    bucket_array[top_priority].pop_back();
    while (top_priority >= 0 && bucket_array[top_priority].empty()) {
      --top_priority;
    }
  }

  cinfo->actual_number_of_colors = k;
  cinfo->colormap = (*cinfo->mem->alloc_sarray)(
      reinterpret_cast<j_common_ptr>(cinfo), JPOOL_IMAGE, k, 3);
  for (int i = 0; i < k; ++i) {
    int index = center[i];
    cinfo->colormap[0][i] = red[index];
    cinfo->colormap[1][i] = green[index];
    cinfo->colormap[2][i] = blue[index];
  }
}

namespace {

void FindCandidatesForCell(j_decompress_ptr cinfo, int ncomp, const int cell[],
                           std::vector<uint8_t>* candidates) {
  int cell_min[kMaxComponents];
  int cell_max[kMaxComponents];
  int cell_center[kMaxComponents];
  for (int c = 0; c < ncomp; ++c) {
    cell_min[c] = cell[c] << (8 - kNumColorCellBits[c]);
    cell_max[c] = cell_min[c] + (1 << (8 - kNumColorCellBits[c])) - 1;
    cell_center[c] = (cell_min[c] + cell_max[c]) >> 1;
  }
  int min_maxdist = std::numeric_limits<int>::max();
  int mindist[256];
  for (int i = 0; i < cinfo->actual_number_of_colors; ++i) {
    int dmin = 0;
    int dmax = 0;
    for (int c = 0; c < ncomp; ++c) {
      int palette_c = cinfo->colormap[c][i];
      int dminc = 0;
      int dmaxc;
      if (palette_c < cell_min[c]) {
        dminc = cell_min[c] - palette_c;
        dmaxc = cell_max[c] - palette_c;
      } else if (palette_c > cell_max[c]) {
        dminc = palette_c - cell_max[c];
        dmaxc = palette_c - cell_min[c];
      } else if (palette_c > cell_center[c]) {
        dmaxc = palette_c - cell_min[c];
      } else {
        dmaxc = cell_max[c] - palette_c;
      }
      dminc *= kCompW[c];
      dmaxc *= kCompW[c];
      dmin += dminc * dminc;
      dmax += dmaxc * dmaxc;
    }
    mindist[i] = dmin;
    min_maxdist = std::min(dmax, min_maxdist);
  }
  for (int i = 0; i < cinfo->actual_number_of_colors; ++i) {
    if (mindist[i] < min_maxdist) {
      candidates->push_back(i);
    }
  }
}

}  // namespace

void CreateInverseColorMap(j_decompress_ptr cinfo) {
  jpeg_decomp_master* m = cinfo->master;
  int ncomp = cinfo->out_color_components;
  JPEGLI_CHECK(ncomp > 0);
  JPEGLI_CHECK(ncomp <= kMaxComponents);
  int num_cells = 1;
  for (int c = 0; c < ncomp; ++c) {
    num_cells *= (1 << kNumColorCellBits[c]);
  }
  m->candidate_lists_.resize(num_cells);

  int next_cell[kMaxComponents] = {0};
  for (int i = 0; i < num_cells; ++i) {
    m->candidate_lists_[i].clear();
    FindCandidatesForCell(cinfo, ncomp, next_cell, &m->candidate_lists_[i]);
    int c = ncomp - 1;
    while (c > 0 && next_cell[c] + 1 == (1 << kNumColorCellBits[c])) {
      next_cell[c--] = 0;
    }
    ++next_cell[c];
  }
  m->regenerate_inverse_colormap_ = false;
}

int LookupColorIndex(j_decompress_ptr cinfo, const JSAMPLE* pixel) {
  jpeg_decomp_master* m = cinfo->master;
  int num_channels = cinfo->out_color_components;
  int index = 0;
  if (m->quant_mode_ == 1) {
    for (int c = 0; c < num_channels; ++c) {
      index += m->colormap_lut_[c * 256 + pixel[c]];
    }
  } else {
    size_t cell_idx = 0;
    size_t stride = 1;
    for (int c = num_channels - 1; c >= 0; --c) {
      cell_idx += (pixel[c] >> (8 - kNumColorCellBits[c])) * stride;
      stride <<= kNumColorCellBits[c];
    }
    JPEGLI_CHECK(cell_idx < m->candidate_lists_.size());
    int mindist = std::numeric_limits<int>::max();
    const auto& candidates = m->candidate_lists_[cell_idx];
    for (uint8_t i : candidates) {
      int dist = 0;
      for (int c = 0; c < num_channels; ++c) {
        int d = (cinfo->colormap[c][i] - pixel[c]) * kCompW[c];
        dist += d * d;
      }
      if (dist < mindist) {
        mindist = dist;
        index = i;
      }
    }
  }
  JPEGLI_CHECK(index < cinfo->actual_number_of_colors);
  return index;
}

void CreateOrderedDitherTables(j_decompress_ptr cinfo) {
  jpeg_decomp_master* m = cinfo->master;
  static constexpr size_t kDitherSize = 4;
  static constexpr size_t kDitherMask = kDitherSize - 1;
  static constexpr float kBaseDitherMatrix[] = {
      0,  8,  2,  10,  //
      12, 4,  14, 6,   //
      3,  11, 1,  9,   //
      15, 7,  13, 5,   //
  };
  m->dither_size_ = kDitherSize;
  m->dither_mask_ = kDitherMask;
  size_t ncells = m->dither_size_ * m->dither_size_;
  for (int c = 0; c < cinfo->out_color_components; ++c) {
    float spread = 1.0f / (m->num_colors_[c] - 1);
    float mul = spread / ncells;
    float offset = 0.5f * spread;
    if (m->dither_[c] == nullptr) {
      m->dither_[c] = Allocate<float>(cinfo, ncells, JPOOL_IMAGE_ALIGNED);
    }
    for (size_t idx = 0; idx < ncells; ++idx) {
      m->dither_[c][idx] = kBaseDitherMatrix[idx] * mul - offset;
    }
  }
}

void InitFSDitherState(j_decompress_ptr cinfo) {
  jpeg_decomp_master* m = cinfo->master;
  for (int c = 0; c < cinfo->out_color_components; ++c) {
    if (m->error_row_[c] == nullptr) {
      m->error_row_[c] =
          Allocate<float>(cinfo, cinfo->output_width, JPOOL_IMAGE_ALIGNED);
      m->error_row_[c + kMaxComponents] =
          Allocate<float>(cinfo, cinfo->output_width, JPOOL_IMAGE_ALIGNED);
    }
    memset(m->error_row_[c], 0.0, cinfo->output_width * sizeof(float));
    memset(m->error_row_[c + kMaxComponents], 0.0,
           cinfo->output_width * sizeof(float));
  }
}

}  // namespace jpegli