/*
 * Copyright (c) 2021, Alliance for Open Media. All rights reserved
 *
 * This source code is subject to the terms of the BSD 3-Clause Clear License
 * and the Alliance for Open Media Patent License 1.0. If the BSD 3-Clause Clear
 * License was not distributed with this source code in the LICENSE file, you
 * can obtain it at aomedia.org/license/software-license/bsd-3-c-c/.  If the
 * Alliance for Open Media Patent License 1.0 was not distributed with this
 * source code in the PATENTS file, you can obtain it at
 * aomedia.org/license/patent-license/.
 */

#include <math.h>
#include <stdlib.h>

#include "av2/common/pred_common.h"

#include "avm_ports/bitops.h"
#include "av2/encoder/block.h"
#include "av2/common/cost.h"
#include "av2/encoder/encoder.h"
#include "av2/encoder/intra_mode_search.h"
#include "av2/encoder/intra_mode_search_utils.h"
#include "av2/encoder/palette.h"
#include "av2/encoder/random.h"
#include "av2/encoder/rdopt_utils.h"
#include "av2/encoder/tx_search.h"

#define AV2_K_MEANS_DIM 1
#include "av2/encoder/k_means_template.h"
#undef AV2_K_MEANS_DIM
#define AV2_K_MEANS_DIM 2
#include "av2/encoder/k_means_template.h"
#undef AV2_K_MEANS_DIM

static int int_comparer(const void *a, const void *b) {
  return (*(int *)a - *(int *)b);
}

int av2_remove_duplicates(int *centroids, int num_centroids) {
  int num_unique;  // number of unique centroids
  int i;
  qsort(centroids, num_centroids, sizeof(*centroids), int_comparer);
  // Remove duplicates.
  num_unique = 1;
  for (i = 1; i < num_centroids; ++i) {
    if (centroids[i] != centroids[i - 1]) {  // found a new unique centroid
      centroids[num_unique++] = centroids[i];
    }
  }
  return num_unique;
}

static int delta_encode_cost(const int *colors, int num, int bit_depth,
                             int min_val) {
  if (num <= 0) return 0;
  int bits_cost = bit_depth;
  if (num == 1) return bits_cost;
  bits_cost += 2;
  int max_delta = 0;
  int deltas[PALETTE_MAX_SIZE];
  const int min_bits = bit_depth - 3;
  for (int i = 1; i < num; ++i) {
    const int delta = colors[i] - colors[i - 1];
    deltas[i - 1] = delta;
    assert(delta >= min_val);
    if (delta > max_delta) max_delta = delta;
  }
  int bits_per_delta = AVMMAX(avm_ceil_log2(max_delta + 1 - min_val), min_bits);
  assert(bits_per_delta <= bit_depth);
  int range = (1 << bit_depth) - colors[0] - min_val;
  for (int i = 0; i < num - 1; ++i) {
    bits_cost += bits_per_delta;
    range -= deltas[i];
    bits_per_delta = AVMMIN(bits_per_delta, avm_ceil_log2(range));
  }
  return bits_cost;
}

int av2_index_color_cache(const uint16_t *color_cache, int n_cache,
                          const uint16_t *colors, int n_colors,
                          uint8_t *cache_color_found, int *out_cache_colors) {
  if (n_cache <= 0) {
    for (int i = 0; i < n_colors; ++i) out_cache_colors[i] = colors[i];
    return n_colors;
  }
  memset(cache_color_found, 0, n_cache * sizeof(*cache_color_found));
#ifndef NDEBUG
  int n_in_cache = 0;
#endif  // NDEBUG
  int in_cache_flags[PALETTE_MAX_SIZE];
  memset(in_cache_flags, 0, sizeof(in_cache_flags));
  for (int i = 0; i < n_cache; ++i) {
    int duplicate = 0;
    for (int j = 0; j < i; j++) {
      if (color_cache[i] == color_cache[j]) duplicate = 1;
    }
    if (duplicate) continue;
    for (int j = 0; j < n_colors; ++j) {
      if (colors[j] == color_cache[i]) {
        in_cache_flags[j] = 1;
        cache_color_found[i] = 1;
#ifndef NDEBUG
        ++n_in_cache;
#endif  // NDEBUG
        break;
      }
    }
  }
  int j = 0;
  for (int i = 0; i < n_colors; ++i)
    if (!in_cache_flags[i]) out_cache_colors[j++] = colors[i];
  assert(j == n_colors - n_in_cache);
  return j;
}

int av2_palette_color_cost_y(const PALETTE_MODE_INFO *const pmi,
                             const uint16_t *color_cache, int n_cache,
                             int bit_depth) {
  const int n = pmi->palette_size[0];
  int out_cache_colors[PALETTE_MAX_SIZE];
  uint8_t cache_color_found[2 * PALETTE_MAX_SIZE];
  const int n_out_cache =
      av2_index_color_cache(color_cache, n_cache, pmi->palette_colors, n,
                            cache_color_found, out_cache_colors);
  const int total_bits =
      n_cache + delta_encode_cost(out_cache_colors, n_out_cache, bit_depth, 1);
  return av2_cost_literal(total_bits);
}

// Extends 'color_map' array from 'orig_width x orig_height' to 'new_width x
// new_height'. Extra rows and columns are filled in by copying last valid
// row/column.
static AVM_INLINE void extend_palette_color_map(uint8_t *const color_map,
                                                int orig_width, int orig_height,
                                                int new_width, int new_height) {
  int j;
  assert(new_width >= orig_width);
  assert(new_height >= orig_height);
  if (new_width == orig_width && new_height == orig_height) return;

  for (j = orig_height - 1; j >= 0; --j) {
    memmove(color_map + j * new_width, color_map + j * orig_width, orig_width);
    // Copy last column to extra columns.
    memset(color_map + j * new_width + orig_width,
           color_map[j * new_width + orig_width - 1], new_width - orig_width);
  }
  // Copy last row to extra rows.
  for (j = orig_height; j < new_height; ++j) {
    memcpy(color_map + j * new_width, color_map + (orig_height - 1) * new_width,
           new_width);
  }
}

// Bias toward using colors in the cache.
// TODO(huisu): Try other schemes to improve compression.
#define PALETTE_CACHE_BIAS_THRESH 4
static AVM_INLINE void optimize_palette_colors(uint16_t *color_cache,
                                               int n_cache, int n_colors,
                                               int stride, int *centroids,
                                               int bit_depth) {
  if (n_cache <= 0) return;
  for (int i = 0; i < n_colors * stride; i += stride) {
    int min_diff = abs(centroids[i] - (int)color_cache[0]);
    int idx = 0;
    for (int j = 1; j < n_cache; ++j) {
      const int this_diff = abs(centroids[i] - color_cache[j]);
      if (this_diff < min_diff) {
        min_diff = this_diff;
        idx = j;
      }
    }
    const int min_threshold = (PALETTE_CACHE_BIAS_THRESH) << (bit_depth - 8);
    if (min_diff <= min_threshold) centroids[i] = color_cache[idx];
  }
}

/*!\brief Calculate the luma palette cost from a given color palette
 *
 * \ingroup palette_mode_search
 * \callergraph
 * Given the base colors as specified in centroids[], calculate the RD cost
 * of palette mode.
 */
static AVM_INLINE void palette_rd_y(
    const AV2_COMP *const cpi, MACROBLOCK *x, MB_MODE_INFO *mbmi,
    BLOCK_SIZE bsize, int dc_mode_cost, const int *data, int *centroids, int n,
    uint16_t *color_cache, int n_cache, MB_MODE_INFO *best_mbmi,
    uint8_t *best_palette_color_map, int64_t *best_rd, int64_t *best_model_rd,
    int *rate, int *rate_tokenonly, int64_t *distortion, int *skippable,
    int *beat_best_rd, PICK_MODE_CONTEXT *ctx, uint8_t *blk_skip,
    TX_TYPE *tx_type_map, int *beat_best_palette_rd) {
  optimize_palette_colors(color_cache, n_cache, n, 1, centroids,
                          cpi->common.seq_params.bit_depth);
  const int num_unique_colors = av2_remove_duplicates(centroids, n);
  if (num_unique_colors < PALETTE_MIN_SIZE) {
    // Too few unique colors to create a palette. And DC_PRED will work
    // well for that case anyway. So skip.
    return;
  }
  PALETTE_MODE_INFO *const pmi = &mbmi->palette_mode_info;
  for (int i = 0; i < num_unique_colors; ++i) {
    pmi->palette_colors[i] =
        clip_pixel_highbd((int)centroids[i], cpi->common.seq_params.bit_depth);
  }
  pmi->palette_size[0] = num_unique_colors;
  MACROBLOCKD *const xd = &x->e_mbd;
  uint8_t *const color_map = xd->plane[0].color_index_map;
  int block_width, block_height, rows, cols;
  av2_get_block_dimensions(bsize, 0, xd, &block_width, &block_height, &rows,
                           &cols);
  av2_calc_indices(data, centroids, color_map, rows * cols, num_unique_colors,
                   1);
  extend_palette_color_map(color_map, cols, rows, block_width, block_height);

  const int palette_mode_cost =
      intra_mode_info_cost_y(cpi, x, mbmi, bsize, dc_mode_cost);
  mbmi->angle_delta[PLANE_TYPE_Y] = 0;

  mbmi->use_dpcm_y = 0;
  mbmi->dpcm_mode_y = 0;
  if (model_intra_yrd_and_prune(cpi, x, bsize, palette_mode_cost,
                                best_model_rd)) {
    return;
  }

  RD_STATS tokenonly_rd_stats;
  av2_pick_uniform_tx_size_type_yrd(cpi, x, &tokenonly_rd_stats, bsize,
                                    *best_rd);
  if (tokenonly_rd_stats.rate == INT_MAX) return;
  int this_rate = tokenonly_rd_stats.rate + palette_mode_cost;
  int64_t this_rd = RDCOST(x->rdmult, this_rate, tokenonly_rd_stats.dist);
  if (!xd->lossless[mbmi->segment_id] &&
      block_signals_txsize(mbmi->sb_type[PLANE_TYPE_Y])) {
    tokenonly_rd_stats.rate -= tx_size_cost(x, bsize, mbmi->tx_size);
  }
  // Collect mode stats for multiwinner mode processing
  const int txfm_search_done = 1;
  const MV_REFERENCE_FRAME refs[2] = { -1, -1 };
  store_winner_mode_stats(
      &cpi->common, x, mbmi, NULL, NULL, NULL, refs, DC_PRED, color_map, bsize,
      this_rd, cpi->sf.winner_mode_sf.multi_winner_mode_type, txfm_search_done);
  if (this_rd < *best_rd) {
    *best_rd = this_rd;
    // Setting beat_best_rd flag because current mode rd is better than best_rd.
    // This flag need to be updated only for palette evaluation in key frames
    if (beat_best_rd) *beat_best_rd = 1;
    memcpy(best_palette_color_map, color_map,
           block_width * block_height * sizeof(color_map[0]));
    *best_mbmi = *mbmi;
    memcpy(
        blk_skip, x->txfm_search_info.blk_skip[AVM_PLANE_Y],
        sizeof(*x->txfm_search_info.blk_skip[AVM_PLANE_Y]) * ctx->num_4x4_blk);
    av2_copy_array(tx_type_map, xd->tx_type_map, ctx->num_4x4_blk);
    if (rate) *rate = this_rate;
    if (rate_tokenonly) *rate_tokenonly = tokenonly_rd_stats.rate;
    if (distortion) *distortion = tokenonly_rd_stats.dist;
    if (skippable) *skippable = tokenonly_rd_stats.skip_txfm;
    if (beat_best_palette_rd) *beat_best_palette_rd = 1;
  }
}

static AVM_INLINE int is_iter_over(int curr_idx, int end_idx, int step_size) {
  assert(step_size != 0);
  return (step_size > 0) ? curr_idx >= end_idx : curr_idx <= end_idx;
}

// Performs count-based palette search with number of colors in interval
// [start_n, end_n) with step size step_size. If step_size < 0, then end_n can
// be less than start_n. Saves the last numbers searched in last_n_searched and
// returns the best number of colors found.
static AVM_INLINE int perform_top_color_palette_search(
    const AV2_COMP *const cpi, MACROBLOCK *x, MB_MODE_INFO *mbmi,
    BLOCK_SIZE bsize, int dc_mode_cost, const int *data, int *top_colors,
    int start_n, int end_n, int step_size, int *last_n_searched,
    uint16_t *color_cache, int n_cache, MB_MODE_INFO *best_mbmi,
    uint8_t *best_palette_color_map, int64_t *best_rd, int64_t *best_model_rd,
    int *rate, int *rate_tokenonly, int64_t *distortion, int *skippable,
    int *beat_best_rd, PICK_MODE_CONTEXT *ctx, uint8_t *best_blk_skip,
    TX_TYPE *tx_type_map) {
  int centroids[PALETTE_MAX_SIZE];
  int n = start_n;
  int top_color_winner = end_n;
  /* clang-format off */
  assert(IMPLIES(step_size < 0, start_n > end_n));
  /* clang-format on */
  assert(IMPLIES(step_size > 0, start_n < end_n));
  while (!is_iter_over(n, end_n, step_size)) {
    int beat_best_palette_rd = 0;
    memcpy(centroids, top_colors, n * sizeof(top_colors[0]));
    palette_rd_y(cpi, x, mbmi, bsize, dc_mode_cost, data, centroids, n,
                 color_cache, n_cache, best_mbmi, best_palette_color_map,
                 best_rd, best_model_rd, rate, rate_tokenonly, distortion,
                 skippable, beat_best_rd, ctx, best_blk_skip, tx_type_map,
                 &beat_best_palette_rd);
    *last_n_searched = n;
    if (beat_best_palette_rd) {
      top_color_winner = n;
    } else if (cpi->sf.intra_sf.prune_palette_search_level == 2) {
      // At search level 2, we return immediately if we don't see an improvement
      return top_color_winner;
    }
    n += step_size;
  }
  return top_color_winner;
}

// Performs k-means based palette search with number of colors in interval
// [start_n, end_n) with step size step_size. If step_size < 0, then end_n can
// be less than start_n. Saves the last numbers searched in last_n_searched and
// returns the best number of colors found.
static AVM_INLINE int perform_k_means_palette_search(
    const AV2_COMP *const cpi, MACROBLOCK *x, MB_MODE_INFO *mbmi,
    BLOCK_SIZE bsize, int dc_mode_cost, const int *data, int lb, int ub,
    int start_n, int end_n, int step_size, int *last_n_searched,
    uint16_t *color_cache, int n_cache, MB_MODE_INFO *best_mbmi,
    uint8_t *best_palette_color_map, int64_t *best_rd, int64_t *best_model_rd,
    int *rate, int *rate_tokenonly, int64_t *distortion, int *skippable,
    int *beat_best_rd, PICK_MODE_CONTEXT *ctx, uint8_t *best_blk_skip,
    TX_TYPE *tx_type_map, uint8_t *color_map, int data_points) {
  int centroids[PALETTE_MAX_SIZE];
  const int max_itr = 50;
  int n = start_n;
  int top_color_winner = end_n;
  /* clang-format off */
  assert(IMPLIES(step_size < 0, start_n > end_n));
  /* clang-format on */
  assert(IMPLIES(step_size > 0, start_n < end_n));
  while (!is_iter_over(n, end_n, step_size)) {
    int beat_best_palette_rd = 0;
    for (int i = 0; i < n; ++i) {
      centroids[i] = lb + (2 * i + 1) * (ub - lb) / n / 2;
    }
    av2_k_means(data, centroids, color_map, data_points, n, 1, max_itr);
    palette_rd_y(cpi, x, mbmi, bsize, dc_mode_cost, data, centroids, n,
                 color_cache, n_cache, best_mbmi, best_palette_color_map,
                 best_rd, best_model_rd, rate, rate_tokenonly, distortion,
                 skippable, beat_best_rd, ctx, best_blk_skip, tx_type_map,
                 &beat_best_palette_rd);
    *last_n_searched = n;
    if (beat_best_palette_rd) {
      top_color_winner = n;
    } else if (cpi->sf.intra_sf.prune_palette_search_level == 2) {
      // At search level 2, we return immediately if we don't see an improvement
      return top_color_winner;
    }
    n += step_size;
  }
  return top_color_winner;
}

// Sets the parameters to search the current number of colors +- 1
static AVM_INLINE void set_stage2_params(int *min_n, int *max_n, int *step_size,
                                         int winner, int end_n) {
  // Set min to winner - 1 unless we are already at the border, then we set it
  // to winner + 1
  *min_n = (winner == PALETTE_MIN_SIZE) ? (PALETTE_MIN_SIZE + 1)
                                        : AVMMAX(winner - 1, PALETTE_MIN_SIZE);
  // Set max to winner + 1 unless we are already at the border, then we set it
  // to winner - 1
  *max_n =
      (winner == end_n) ? (winner - 1) : AVMMIN(winner + 1, PALETTE_MAX_SIZE);

  // Set the step size to max_n - min_n so we only search those two values.
  // If max_n == min_n, then set step_size to 1 to avoid infinite loop later.
  *step_size = AVMMAX(1, *max_n - *min_n);
}

void av2_rd_pick_palette_intra_sby(
    const AV2_COMP *cpi, MACROBLOCK *x, BLOCK_SIZE bsize, int dc_mode_cost,
    MB_MODE_INFO *best_mbmi, uint8_t *best_palette_color_map, int64_t *best_rd,
    int64_t *best_model_rd, int *rate, int *rate_tokenonly, int64_t *distortion,
    int *skippable, int *beat_best_rd, PICK_MODE_CONTEXT *ctx,
    uint8_t *best_blk_skip, TX_TYPE *tx_type_map) {
  MACROBLOCKD *const xd = &x->e_mbd;
  MB_MODE_INFO *const mbmi = xd->mi[0];

  mbmi->mrl_index = 0;
  mbmi->multi_line_mrl = 0;
  for (int i = 0; i < MAX_TX_PARTITIONS; ++i) {
    mbmi->is_wide_angle[0][i] = 0;
    mbmi->mapped_intra_mode[0][i] = DC_PRED;
  }

  mbmi->fsc_mode[xd->tree_type == CHROMA_PART] = 0;
  assert(!is_inter_block(mbmi, xd->tree_type));
  assert(av2_allow_palette(
      PLANE_TYPE_Y, cpi->common.features.allow_screen_content_tools, bsize));
  assert(PALETTE_MAX_SIZE == 8);
  assert(PALETTE_MIN_SIZE == 2);
  mbmi->angle_delta[PLANE_TYPE_Y] = 0;
  mbmi->use_dpcm_y = 0;
  mbmi->dpcm_mode_y = 0;
  const int src_stride = x->plane[0].src.stride;
  uint16_t *src = x->plane[0].src.buf;
  int block_width, block_height, rows, cols;
  av2_get_block_dimensions(bsize, 0, xd, &block_width, &block_height, &rows,
                           &cols);
  const SequenceHeader *const seq_params = &cpi->common.seq_params;
  const int bit_depth = seq_params->bit_depth;
  int unused;

  int count_buf[1 << 12];      // Maximum (1 << 12) color levels.
  int count_buf_8bit[1 << 8];  // Maximum (1 << 8) bins for hbd path.
  int colors, colors_threshold = 0;
  av2_count_colors_highbd(src, src_stride, rows, cols, bit_depth, count_buf,
                          count_buf_8bit, &colors_threshold, &colors);

  uint8_t *const color_map = xd->plane[0].color_index_map;
  if (colors_threshold > 1 && colors_threshold <= 64) {
    int *const data = x->palette_buffer->kmeans_data_buf;
    int centroids[PALETTE_MAX_SIZE];
    int lb, ub;
    int *data_pt = data;
    lb = ub = src[0];
    for (int r = 0; r < rows; ++r) {
      for (int c = 0; c < cols; ++c) {
        const int val = src[c];
        data_pt[c] = val;
        lb = AVMMIN(lb, val);
        ub = AVMMAX(ub, val);
      }
      src += src_stride;
      data_pt += cols;
    }

    mbmi->mode = DC_PRED;
    mbmi->joint_y_mode_delta_angle = DC_PRED;
    mbmi->y_mode_idx = DC_PRED;
    mbmi->use_intra_dip = 0;

    uint16_t color_cache[2 * PALETTE_MAX_SIZE];
    const int n_cache = av2_get_palette_cache(xd, 0, color_cache);

    // Find the dominant colors, stored in top_colors[].
    int top_colors[PALETTE_MAX_SIZE] = { 0 };
    for (int i = 0; i < AVMMIN(colors, PALETTE_MAX_SIZE); ++i) {
      int max_count = 0;
      for (int j = 0; j < (1 << bit_depth); ++j) {
        if (count_buf[j] > max_count) {
          max_count = count_buf[j];
          top_colors[i] = j;
        }
      }
      assert(max_count > 0);
      count_buf[top_colors[i]] = 0;
    }

    // TODO(huisu@google.com): Try to avoid duplicate computation in cases
    // where the dominant colors and the k-means results are similar.
    if ((cpi->sf.intra_sf.prune_palette_search_level == 1) &&
        (colors > PALETTE_MIN_SIZE)) {
      // Start index and step size below are chosen to evaluate unique
      // candidates in neighbor search, in case a winner candidate is found in
      // coarse search. Example,
      // 1) 8 colors (end_n = 8): 2,3,4,5,6,7,8. start_n is chosen as 2 and step
      // size is chosen as 3. Therefore, coarse search will evaluate 2, 5 and 8.
      // If winner is found at 5, then 4 and 6 are evaluated. Similarly, for 2
      // (3) and 8 (7).
      // 2) 7 colors (end_n = 7): 2,3,4,5,6,7. If start_n is chosen as 2 (same
      // as for 8 colors) then step size should also be 2, to cover all
      // candidates. Coarse search will evaluate 2, 4 and 6. If winner is either
      // 2 or 4, 3 will be evaluated. Instead, if start_n=3 and step_size=3,
      // coarse search will evaluate 3 and 6. For the winner, unique neighbors
      // (3: 2,4 or 6: 5,7) would be evaluated.

      // Start index for coarse palette search for dominant colors and k-means
      const uint8_t start_n_lookup_table[PALETTE_MAX_SIZE + 1] = { 0, 0, 0,
                                                                   3, 3, 2,
                                                                   3, 3, 2 };
      // Step size for coarse palette search for dominant colors and k-means
      const uint8_t step_size_lookup_table[PALETTE_MAX_SIZE + 1] = { 0, 0, 0,
                                                                     3, 3, 3,
                                                                     3, 3, 3 };

      // Choose the start index and step size for coarse search based on number
      // of colors
      const int max_n = AVMMIN(colors, PALETTE_MAX_SIZE);
      const int min_n = start_n_lookup_table[max_n];
      const int step_size = step_size_lookup_table[max_n];
      assert(min_n >= PALETTE_MIN_SIZE);

      // Perform top color coarse palette search to find the winner candidate
      const int top_color_winner = perform_top_color_palette_search(
          cpi, x, mbmi, bsize, dc_mode_cost, data, top_colors, min_n, max_n + 1,
          step_size, &unused, color_cache, n_cache, best_mbmi,
          best_palette_color_map, best_rd, best_model_rd, rate, rate_tokenonly,
          distortion, skippable, beat_best_rd, ctx, best_blk_skip, tx_type_map);
      // Evaluate neighbors for the winner color (if winner is found) in the
      // above coarse search for dominant colors
      if (top_color_winner <= max_n) {
        int stage2_min_n, stage2_max_n, stage2_step_size;
        set_stage2_params(&stage2_min_n, &stage2_max_n, &stage2_step_size,
                          top_color_winner, max_n);
        // perform finer search for the winner candidate
        perform_top_color_palette_search(
            cpi, x, mbmi, bsize, dc_mode_cost, data, top_colors, stage2_min_n,
            stage2_max_n + 1, stage2_step_size, &unused, color_cache, n_cache,
            best_mbmi, best_palette_color_map, best_rd, best_model_rd, rate,
            rate_tokenonly, distortion, skippable, beat_best_rd, ctx,
            best_blk_skip, tx_type_map);
      }
      // K-means clustering.
      // Perform k-means coarse palette search to find the winner candidate
      const int k_means_winner = perform_k_means_palette_search(
          cpi, x, mbmi, bsize, dc_mode_cost, data, lb, ub, min_n, max_n + 1,
          step_size, &unused, color_cache, n_cache, best_mbmi,
          best_palette_color_map, best_rd, best_model_rd, rate, rate_tokenonly,
          distortion, skippable, beat_best_rd, ctx, best_blk_skip, tx_type_map,
          color_map, rows * cols);
      // Evaluate neighbors for the winner color (if winner is found) in the
      // above coarse search for k-means
      if (k_means_winner <= max_n) {
        int start_n_stage2, end_n_stage2, step_size_stage2;
        set_stage2_params(&start_n_stage2, &end_n_stage2, &step_size_stage2,
                          k_means_winner, max_n);
        // perform finer search for the winner candidate
        perform_k_means_palette_search(
            cpi, x, mbmi, bsize, dc_mode_cost, data, lb, ub, start_n_stage2,
            end_n_stage2 + 1, step_size_stage2, &unused, color_cache, n_cache,
            best_mbmi, best_palette_color_map, best_rd, best_model_rd, rate,
            rate_tokenonly, distortion, skippable, beat_best_rd, ctx,
            best_blk_skip, tx_type_map, color_map, rows * cols);
      }
    } else {
      const int max_n = AVMMIN(colors, PALETTE_MAX_SIZE),
                min_n = PALETTE_MIN_SIZE;
      // Perform top color palette search in descending order
      int last_n_searched = max_n;
      perform_top_color_palette_search(
          cpi, x, mbmi, bsize, dc_mode_cost, data, top_colors, max_n, min_n - 1,
          -1, &last_n_searched, color_cache, n_cache, best_mbmi,
          best_palette_color_map, best_rd, best_model_rd, rate, rate_tokenonly,
          distortion, skippable, beat_best_rd, ctx, best_blk_skip, tx_type_map);

      if (last_n_searched > min_n) {
        // Search in ascending order until we get to the previous best
        perform_top_color_palette_search(
            cpi, x, mbmi, bsize, dc_mode_cost, data, top_colors, min_n,
            last_n_searched, 1, &unused, color_cache, n_cache, best_mbmi,
            best_palette_color_map, best_rd, best_model_rd, rate,
            rate_tokenonly, distortion, skippable, beat_best_rd, ctx,
            best_blk_skip, tx_type_map);
      }
      // K-means clustering.
      if (colors == PALETTE_MIN_SIZE) {
        // Special case: These colors automatically become the centroids.
        assert(colors == 2);
        centroids[0] = lb;
        centroids[1] = ub;
        palette_rd_y(cpi, x, mbmi, bsize, dc_mode_cost, data, centroids, colors,
                     color_cache, n_cache, best_mbmi, best_palette_color_map,
                     best_rd, best_model_rd, rate, rate_tokenonly, distortion,
                     skippable, beat_best_rd, ctx, best_blk_skip, tx_type_map,
                     NULL);
      } else {
        // Perform k-means palette search in descending order
        last_n_searched = max_n;
        perform_k_means_palette_search(
            cpi, x, mbmi, bsize, dc_mode_cost, data, lb, ub, max_n, min_n - 1,
            -1, &last_n_searched, color_cache, n_cache, best_mbmi,
            best_palette_color_map, best_rd, best_model_rd, rate,
            rate_tokenonly, distortion, skippable, beat_best_rd, ctx,
            best_blk_skip, tx_type_map, color_map, rows * cols);
        if (last_n_searched > min_n) {
          // Search in ascending order until we get to the previous best
          perform_k_means_palette_search(
              cpi, x, mbmi, bsize, dc_mode_cost, data, lb, ub, min_n,
              last_n_searched, 1, &unused, color_cache, n_cache, best_mbmi,
              best_palette_color_map, best_rd, best_model_rd, rate,
              rate_tokenonly, distortion, skippable, beat_best_rd, ctx,
              best_blk_skip, tx_type_map, color_map, rows * cols);
        }
      }
    }
  }

  if (best_mbmi->palette_mode_info.palette_size[0] > 0) {
    memcpy(color_map, best_palette_color_map,
           block_width * block_height * sizeof(best_palette_color_map[0]));
    // Gather the stats to determine whether to use screen content tools in
    // function av2_determine_sc_tools_with_encoding().
    x->palette_pixels += (block_width * block_height);
  }
  *mbmi = *best_mbmi;
}