/*

 * Copyright (c) 2019, Alliance for Open Media. All rights reserved

 *

 * This source code is subject to the terms of the BSD 2 Clause License and

 * the Alliance for Open Media Patent License 1.0. If the BSD 2 Clause License

 * was not distributed with this source code in the LICENSE file, you can

 * obtain it at www.aomedia.org/license/software. 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 www.aomedia.org/license/patent.

 */

#include <limits.h>

#include <math.h>

#include <stdbool.h>

#include <stdio.h>

#include "config/aom_config.h"

#include "config/aom_dsp_rtcd.h"

#include "config/av1_rtcd.h"

#include "aom_dsp/aom_dsp_common.h"

#include "aom_dsp/binary_codes_writer.h"

#include "aom_ports/mem.h"

#include "aom_ports/aom_timer.h"

#include "av1/common/reconinter.h"

#include "av1/common/blockd.h"

#include "av1/encoder/encodeframe.h"

#include "av1/encoder/var_based_part.h"

#include "av1/encoder/reconinter_enc.h"

extern const uint8_t AV1_VAR_OFFS[];

typedef struct {

  VPVariance *part_variances;

  VPartVar *split[4];

} variance_node;

static AOM_INLINE void tree_to_node(void *data, BLOCK_SIZE bsize,

                                    variance_node *node) {

  int i;

  node->part_variances = NULL;

  switch (bsize) {

    case BLOCK_128X128: {

      VP128x128 *vt = (VP128x128 *)data;

      node->part_variances = &vt->part_variances;

      for (i = 0; i < 4; i++)

        node->split[i] = &vt->split[i].part_variances.none;

      break;

    }

    case BLOCK_64X64: {

      VP64x64 *vt = (VP64x64 *)data;

      node->part_variances = &vt->part_variances;

      for (i = 0; i < 4; i++)

        node->split[i] = &vt->split[i].part_variances.none;

      break;

    }

    case BLOCK_32X32: {

      VP32x32 *vt = (VP32x32 *)data;

      node->part_variances = &vt->part_variances;

      for (i = 0; i < 4; i++)

        node->split[i] = &vt->split[i].part_variances.none;

      break;

    }

    case BLOCK_16X16: {

      VP16x16 *vt = (VP16x16 *)data;

      node->part_variances = &vt->part_variances;

      for (i = 0; i < 4; i++)

        node->split[i] = &vt->split[i].part_variances.none;

      break;

    }

    case BLOCK_8X8: {

      VP8x8 *vt = (VP8x8 *)data;

      node->part_variances = &vt->part_variances;

      for (i = 0; i < 4; i++)

        node->split[i] = &vt->split[i].part_variances.none;

      break;

    }

    default: {

      VP4x4 *vt = (VP4x4 *)data;

      assert(bsize == BLOCK_4X4);

      node->part_variances = &vt->part_variances;

      for (i = 0; i < 4; i++) node->split[i] = &vt->split[i];

      break;

    }

  }

}

// Set variance values given sum square error, sum error, count.

static AOM_INLINE void fill_variance(uint32_t s2, int32_t s, int c,

                                     VPartVar *v) {

  v->sum_square_error = s2;

  v->sum_error = s;

  v->log2_count = c;

}

static AOM_INLINE void get_variance(VPartVar *v) {

  v->variance =

      (int)(256 * (v->sum_square_error -

                   (uint32_t)(((int64_t)v->sum_error * v->sum_error) >>

                              v->log2_count)) >>

            v->log2_count);

}

static AOM_INLINE void sum_2_variances(const VPartVar *a, const VPartVar *b,

                                       VPartVar *r) {

  assert(a->log2_count == b->log2_count);

  fill_variance(a->sum_square_error + b->sum_square_error,

                a->sum_error + b->sum_error, a->log2_count + 1, r);

}

static AOM_INLINE void fill_variance_tree(void *data, BLOCK_SIZE bsize) {

  variance_node node;

  memset(&node, 0, sizeof(node));

  tree_to_node(data, bsize, &node);

  sum_2_variances(node.split[0], node.split[1], &node.part_variances->horz[0]);

  sum_2_variances(node.split[2], node.split[3], &node.part_variances->horz[1]);

  sum_2_variances(node.split[0], node.split[2], &node.part_variances->vert[0]);

  sum_2_variances(node.split[1], node.split[3], &node.part_variances->vert[1]);

  sum_2_variances(&node.part_variances->vert[0], &node.part_variances->vert[1],

                  &node.part_variances->none);

}

static AOM_INLINE void set_block_size(AV1_COMP *const cpi, MACROBLOCK *const x,

                                      MACROBLOCKD *const xd, int mi_row,

                                      int mi_col, BLOCK_SIZE bsize) {

  if (cpi->common.mi_params.mi_cols > mi_col &&

      cpi->common.mi_params.mi_rows > mi_row) {

    set_mode_info_offsets(&cpi->common.mi_params, &cpi->mbmi_ext_info, x, xd,

                          mi_row, mi_col);

    xd->mi[0]->bsize = bsize;

  }

}

static int set_vt_partitioning(AV1_COMP *cpi, MACROBLOCK *const x,

                               MACROBLOCKD *const xd,

                               const TileInfo *const tile, void *data,

                               BLOCK_SIZE bsize, int mi_row, int mi_col,

                               int64_t threshold, BLOCK_SIZE bsize_min,

                               int force_split) {

  AV1_COMMON *const cm = &cpi->common;

  variance_node vt;

  const int block_width = mi_size_wide[bsize];

  const int block_height = mi_size_high[bsize];

  int bs_width_check = block_width;

  int bs_height_check = block_height;

  int bs_width_vert_check = block_width >> 1;

  int bs_height_horiz_check = block_height >> 1;

  // On the right and bottom boundary we only need to check

  // if half the bsize fits, because boundary is extended

  // up to 64. So do this check only for sb_size = 64X64.

  if (cm->seq_params->sb_size == BLOCK_64X64) {

    if (tile->mi_col_end == cm->mi_params.mi_cols) {

      bs_width_check = (block_width >> 1) + 1;

      bs_width_vert_check = (block_width >> 2) + 1;

    }

    if (tile->mi_row_end == cm->mi_params.mi_rows) {

      bs_height_check = (block_height >> 1) + 1;

      bs_height_horiz_check = (block_height >> 2) + 1;

    }

  }

  assert(block_height == block_width);

  tree_to_node(data, bsize, &vt);

  if (force_split == 1) return 0;

  // For bsize=bsize_min (16x16/8x8 for 8x8/4x4 downsampling), select if

  // variance is below threshold, otherwise split will be selected.

  // No check for vert/horiz split as too few samples for variance.

  if (bsize == bsize_min) {

    // Variance already computed to set the force_split.

    if (frame_is_intra_only(cm)) get_variance(&vt.part_variances->none);

    if (mi_col + bs_width_check <= tile->mi_col_end &&

        mi_row + bs_height_check <= tile->mi_row_end &&

        vt.part_variances->none.variance < threshold) {

      set_block_size(cpi, x, xd, mi_row, mi_col, bsize);

      return 1;

    }

    return 0;

  } else if (bsize > bsize_min) {

    // Variance already computed to set the force_split.

    if (frame_is_intra_only(cm)) get_variance(&vt.part_variances->none);

    // For key frame: take split for bsize above 32X32 or very high variance.

    if (frame_is_intra_only(cm) &&

        (bsize > BLOCK_32X32 ||

         vt.part_variances->none.variance > (threshold << 4))) {

      return 0;

    }

    // If variance is low, take the bsize (no split).

    if (mi_col + bs_width_check <= tile->mi_col_end &&

        mi_row + bs_height_check <= tile->mi_row_end &&

        vt.part_variances->none.variance < threshold) {

      set_block_size(cpi, x, xd, mi_row, mi_col, bsize);

      return 1;

    }

    // Check vertical split.

    if (mi_row + bs_height_check <= tile->mi_row_end &&

        mi_col + bs_width_vert_check <= tile->mi_col_end) {

      BLOCK_SIZE subsize = get_partition_subsize(bsize, PARTITION_VERT);

      get_variance(&vt.part_variances->vert[0]);

      get_variance(&vt.part_variances->vert[1]);

      if (vt.part_variances->vert[0].variance < threshold &&

          vt.part_variances->vert[1].variance < threshold &&

          get_plane_block_size(subsize, xd->plane[1].subsampling_x,

                               xd->plane[1].subsampling_y) < BLOCK_INVALID) {

        set_block_size(cpi, x, xd, mi_row, mi_col, subsize);

        set_block_size(cpi, x, xd, mi_row, mi_col + block_width / 2, subsize);

        return 1;

      }

    }

    // Check horizontal split.

    if (mi_col + bs_width_check <= tile->mi_col_end &&

        mi_row + bs_height_horiz_check <= tile->mi_row_end) {

      BLOCK_SIZE subsize = get_partition_subsize(bsize, PARTITION_HORZ);

      get_variance(&vt.part_variances->horz[0]);

      get_variance(&vt.part_variances->horz[1]);

      if (vt.part_variances->horz[0].variance < threshold &&

          vt.part_variances->horz[1].variance < threshold &&

          get_plane_block_size(subsize, xd->plane[1].subsampling_x,

                               xd->plane[1].subsampling_y) < BLOCK_INVALID) {

        set_block_size(cpi, x, xd, mi_row, mi_col, subsize);

        set_block_size(cpi, x, xd, mi_row + block_height / 2, mi_col, subsize);

        return 1;

      }

    }

    return 0;

  }

  return 0;

}

static AOM_INLINE void fill_variance_8x8avg(const uint8_t *s, int sp,

                                            const uint8_t *d, int dp,

                                            int x16_idx, int y16_idx,

                                            VP16x16 *vst,

#if CONFIG_AV1_HIGHBITDEPTH

                                            int highbd_flag,

#endif

                                            int pixels_wide, int pixels_high,

                                            int is_key_frame) {

  int k;

  for (k = 0; k < 4; k++) {

    int x8_idx = x16_idx + ((k & 1) << 3);

    int y8_idx = y16_idx + ((k >> 1) << 3);

    unsigned int sse = 0;

    int sum = 0;

    if (x8_idx < pixels_wide && y8_idx < pixels_high) {

      int s_avg;

      int d_avg = 128;

#if CONFIG_AV1_HIGHBITDEPTH

      if (highbd_flag & YV12_FLAG_HIGHBITDEPTH) {

        s_avg = aom_highbd_avg_8x8(s + y8_idx * sp + x8_idx, sp);

        if (!is_key_frame)

          d_avg = aom_highbd_avg_8x8(d + y8_idx * dp + x8_idx, dp);

      } else {

        s_avg = aom_avg_8x8(s + y8_idx * sp + x8_idx, sp);

        if (!is_key_frame) d_avg = aom_avg_8x8(d + y8_idx * dp + x8_idx, dp);

      }

#else

      s_avg = aom_avg_8x8(s + y8_idx * sp + x8_idx, sp);

      if (!is_key_frame) d_avg = aom_avg_8x8(d + y8_idx * dp + x8_idx, dp);

#endif

      sum = s_avg - d_avg;

      sse = sum * sum;

    }

    fill_variance(sse, sum, 0, &vst->split[k].part_variances.none);

  }

}

static int compute_minmax_8x8(const uint8_t *s, int sp, const uint8_t *d,

                              int dp, int x16_idx, int y16_idx,

#if CONFIG_AV1_HIGHBITDEPTH

                              int highbd_flag,

#endif

                              int pixels_wide, int pixels_high) {

  int k;

  int minmax_max = 0;

  int minmax_min = 255;

  // Loop over the 4 8x8 subblocks.

  for (k = 0; k < 4; k++) {

    int x8_idx = x16_idx + ((k & 1) << 3);

    int y8_idx = y16_idx + ((k >> 1) << 3);

    int min = 0;

    int max = 0;

    if (x8_idx < pixels_wide && y8_idx < pixels_high) {

#if CONFIG_AV1_HIGHBITDEPTH

      if (highbd_flag & YV12_FLAG_HIGHBITDEPTH) {

        aom_highbd_minmax_8x8(s + y8_idx * sp + x8_idx, sp,

                              d + y8_idx * dp + x8_idx, dp, &min, &max);

      } else {

        aom_minmax_8x8(s + y8_idx * sp + x8_idx, sp, d + y8_idx * dp + x8_idx,

                       dp, &min, &max);

      }

#else

      aom_minmax_8x8(s + y8_idx * sp + x8_idx, sp, d + y8_idx * dp + x8_idx, dp,

                     &min, &max);

#endif

      if ((max - min) > minmax_max) minmax_max = (max - min);

      if ((max - min) < minmax_min) minmax_min = (max - min);

    }

  }

  return (minmax_max - minmax_min);

}

static AOM_INLINE void fill_variance_4x4avg(const uint8_t *s, int sp,

                                            const uint8_t *d, int dp,

                                            int x8_idx, int y8_idx, VP8x8 *vst,

#if CONFIG_AV1_HIGHBITDEPTH

                                            int highbd_flag,

#endif

                                            int pixels_wide, int pixels_high,

                                            int is_key_frame) {

  int k;

  for (k = 0; k < 4; k++) {

    int x4_idx = x8_idx + ((k & 1) << 2);

    int y4_idx = y8_idx + ((k >> 1) << 2);

    unsigned int sse = 0;

    int sum = 0;

    if (x4_idx < pixels_wide && y4_idx < pixels_high) {

      int s_avg;

      int d_avg = 128;

#if CONFIG_AV1_HIGHBITDEPTH

      if (highbd_flag & YV12_FLAG_HIGHBITDEPTH) {

        s_avg = aom_highbd_avg_4x4(s + y4_idx * sp + x4_idx, sp);

        if (!is_key_frame)

          d_avg = aom_highbd_avg_4x4(d + y4_idx * dp + x4_idx, dp);

      } else {

        s_avg = aom_avg_4x4(s + y4_idx * sp + x4_idx, sp);

        if (!is_key_frame) d_avg = aom_avg_4x4(d + y4_idx * dp + x4_idx, dp);

      }

#else

      s_avg = aom_avg_4x4(s + y4_idx * sp + x4_idx, sp);

      if (!is_key_frame) d_avg = aom_avg_4x4(d + y4_idx * dp + x4_idx, dp);

#endif

      sum = s_avg - d_avg;

      sse = sum * sum;

    }

    fill_variance(sse, sum, 0, &vst->split[k].part_variances.none);

  }

}

// TODO(kyslov) Bring back threshold adjustment based on content state

static int64_t scale_part_thresh_content(int64_t threshold_base, int speed,

                                         int width, int height,

                                         int non_reference_frame) {

  (void)width;

  (void)height;

  int64_t threshold = threshold_base;

  if (non_reference_frame) threshold = (3 * threshold) >> 1;

  if (speed >= 8) {

    return (5 * threshold) >> 2;

  }

  return threshold;

}

static AOM_INLINE void set_vbp_thresholds(AV1_COMP *cpi, int64_t thresholds[],

                                          int q, int content_lowsumdiff,

                                          int source_sad, int segment_id) {

  AV1_COMMON *const cm = &cpi->common;

  const int is_key_frame = frame_is_intra_only(cm);

  const int threshold_multiplier = is_key_frame ? 120 : 1;

  const int ac_q = av1_ac_quant_QTX(q, 0, cm->seq_params->bit_depth);

  int64_t threshold_base = (int64_t)(threshold_multiplier * ac_q);

  const int current_qindex = cm->quant_params.base_qindex;

  const int threshold_left_shift = cpi->sf.rt_sf.var_part_split_threshold_shift;

  if (is_key_frame) {

    if (cpi->sf.rt_sf.force_large_partition_blocks_intra) {

      const int shift_steps =

          threshold_left_shift - (cpi->oxcf.mode == ALLINTRA ? 7 : 8);

      assert(shift_steps >= 0);

      threshold_base <<= shift_steps;

    }

    thresholds[0] = threshold_base;

    thresholds[1] = threshold_base;

    if (cm->width * cm->height < 1280 * 720) {

      thresholds[2] = threshold_base / 3;

      thresholds[3] = threshold_base >> 1;

    } else {

      int shift_val = 2;

      if (cpi->sf.rt_sf.force_large_partition_blocks_intra) {

        shift_val = 0;

      }

      thresholds[2] = threshold_base >> shift_val;

      thresholds[3] = threshold_base >> shift_val;

    }

    thresholds[4] = threshold_base << 2;

    return;

  }

  // Increase partition thresholds for noisy content. Apply it only for

  // superblocks where sumdiff is low, as we assume the sumdiff of superblock

  // whose only change is due to noise will be low (i.e, noise will average

  // out over large block).

  if (cpi->noise_estimate.enabled && content_lowsumdiff &&

      (cm->width * cm->height > 640 * 480) &&

      cm->current_frame.frame_number > 60) {

    NOISE_LEVEL noise_level =

        av1_noise_estimate_extract_level(&cpi->noise_estimate);

    if (noise_level == kHigh)

      threshold_base = (5 * threshold_base) >> 1;

    else if (noise_level == kMedium &&

             !cpi->sf.rt_sf.force_large_partition_blocks)

      threshold_base = (5 * threshold_base) >> 2;

  }

  // TODO(kyslov) Enable var based partition adjusment on temporal denoising

#if 0  // CONFIG_AV1_TEMPORAL_DENOISING

  if (cpi->oxcf.noise_sensitivity > 0 && denoise_svc(cpi) &&

      cpi->oxcf.speed > 5 && cpi->denoiser.denoising_level >= kDenLow)

      threshold_base =

          av1_scale_part_thresh(threshold_base, cpi->denoiser.denoising_level,

                                content_state, cpi->svc.temporal_layer_id);

  else

    threshold_base =

        scale_part_thresh_content(threshold_base, cpi->oxcf.speed, cm->width,

                                  cm->height, cpi->svc.non_reference_frame);

#else

  // Increase base variance threshold based on content_state/sum_diff level.

  threshold_base =

      scale_part_thresh_content(threshold_base, cpi->oxcf.speed, cm->width,

                                cm->height, cpi->svc.non_reference_frame);

#endif

  thresholds[0] = threshold_base >> 1;

  thresholds[1] = threshold_base;

  thresholds[3] = threshold_base << threshold_left_shift;

  if (cm->width >= 1280 && cm->height >= 720)

    thresholds[3] = thresholds[3] << 1;

  if (cm->width * cm->height <= 352 * 288) {

    if (current_qindex >= QINDEX_HIGH_THR) {

      threshold_base = (5 * threshold_base) >> 1;

      thresholds[1] = threshold_base >> 3;

      thresholds[2] = threshold_base << 2;

      thresholds[3] = threshold_base << 5;

    } else if (current_qindex < QINDEX_LOW_THR) {

      thresholds[1] = threshold_base >> 3;

      thresholds[2] = threshold_base >> 1;

      thresholds[3] = threshold_base << 3;

    } else {

      int64_t qi_diff_low = current_qindex - QINDEX_LOW_THR;

      int64_t qi_diff_high = QINDEX_HIGH_THR - current_qindex;

      int64_t threshold_diff = QINDEX_HIGH_THR - QINDEX_LOW_THR;

      int64_t threshold_base_high = (5 * threshold_base) >> 1;

      threshold_diff = threshold_diff > 0 ? threshold_diff : 1;

      threshold_base =

          (qi_diff_low * threshold_base_high + qi_diff_high * threshold_base) /

          threshold_diff;

      thresholds[1] = threshold_base >> 3;

      thresholds[2] = ((qi_diff_low * threshold_base) +

                       qi_diff_high * (threshold_base >> 1)) /

                      threshold_diff;

      thresholds[3] = ((qi_diff_low * (threshold_base << 5)) +

                       qi_diff_high * (threshold_base << 3)) /

                      threshold_diff;

    }

  } else if (cm->width < 1280 && cm->height < 720) {

    thresholds[2] = (5 * threshold_base) >> 2;

  } else if (cm->width < 1920 && cm->height < 1080) {

    thresholds[2] = threshold_base << 1;

  } else {

    thresholds[2] = (5 * threshold_base) >> 1;

  }

  if (cpi->sf.rt_sf.force_large_partition_blocks) {

    double weight;

    const int win = 20;

    if (current_qindex < QINDEX_LARGE_BLOCK_THR - win)

      weight = 1.0;

    else if (current_qindex > QINDEX_LARGE_BLOCK_THR + win)

      weight = 0.0;

    else

      weight =

          1.0 - (current_qindex - QINDEX_LARGE_BLOCK_THR + win) / (2 * win);

    if (cm->width * cm->height > 640 * 480) {

      for (int i = 0; i < 4; i++) {

        thresholds[i] <<= 1;

      }

    }

    if (cm->width * cm->height <= 352 * 288) {

      thresholds[3] = INT32_MAX;

      if (segment_id == 0) {

        thresholds[1] <<= 2;

        thresholds[2] <<= (source_sad == kLowSad) ? 5 : 4;

      } else {

        thresholds[1] <<= 1;

        thresholds[2] <<= 3;

      }

      // Condition the increase of partition thresholds on the segment

      // and the content. Avoid the increase for superblocks which have

      // high source sad, unless the whole frame has very high motion

      // (i.e, cpi->rc.avg_source_sad is very large, in which case all blocks

      // have high source sad).

    } else if (cm->width * cm->height > 640 * 480 && segment_id == 0 &&

               (source_sad != kHighSad || cpi->rc.avg_source_sad > 50000)) {

      thresholds[0] = (3 * thresholds[0]) >> 1;

      thresholds[3] = INT32_MAX;

      if (current_qindex > QINDEX_LARGE_BLOCK_THR) {

        thresholds[1] =

            (int)((1 - weight) * (thresholds[1] << 1) + weight * thresholds[1]);

        thresholds[2] =

            (int)((1 - weight) * (thresholds[2] << 1) + weight * thresholds[2]);

      }

    } else if (current_qindex > QINDEX_LARGE_BLOCK_THR && segment_id == 0 &&

               (source_sad != kHighSad || cpi->rc.avg_source_sad > 50000)) {

      thresholds[1] =

          (int)((1 - weight) * (thresholds[1] << 2) + weight * thresholds[1]);

      thresholds[2] =

          (int)((1 - weight) * (thresholds[2] << 4) + weight * thresholds[2]);

      thresholds[3] = INT32_MAX;

    }

  }

}

// Set temporal variance low flag for superblock 64x64.

// Only first 25 in the array are used in this case.

static AOM_INLINE void set_low_temp_var_flag_64x64(

    CommonModeInfoParams *mi_params, PartitionSearchInfo *part_info,

    MACROBLOCKD *xd, VP64x64 *vt, const int64_t thresholds[], int mi_col,

    int mi_row) {

  if (xd->mi[0]->bsize == BLOCK_64X64) {

    if ((vt->part_variances).none.variance < (thresholds[0] >> 1))

      part_info->variance_low[0] = 1;

  } else if (xd->mi[0]->bsize == BLOCK_64X32) {

    for (int i = 0; i < 2; i++) {

      if (vt->part_variances.horz[i].variance < (thresholds[0] >> 2))

        part_info->variance_low[i + 1] = 1;

    }

  } else if (xd->mi[0]->bsize == BLOCK_32X64) {

    for (int i = 0; i < 2; i++) {

      if (vt->part_variances.vert[i].variance < (thresholds[0] >> 2))

        part_info->variance_low[i + 3] = 1;

    }

  } else {

    static const int idx[4][2] = { { 0, 0 }, { 0, 8 }, { 8, 0 }, { 8, 8 } };

    for (int i = 0; i < 4; i++) {

      const int idx_str =

          mi_params->mi_stride * (mi_row + idx[i][0]) + mi_col + idx[i][1];

      MB_MODE_INFO **this_mi = mi_params->mi_grid_base + idx_str;

      if (mi_params->mi_cols <= mi_col + idx[i][1] ||

          mi_params->mi_rows <= mi_row + idx[i][0])

        continue;

      if (*this_mi == NULL) continue;

      if ((*this_mi)->bsize == BLOCK_32X32) {

        int64_t threshold_32x32 = (5 * thresholds[1]) >> 3;

        if (vt->split[i].part_variances.none.variance < threshold_32x32)

          part_info->variance_low[i + 5] = 1;

      } else {

        // For 32x16 and 16x32 blocks, the flag is set on each 16x16 block

        // inside.

        if ((*this_mi)->bsize == BLOCK_16X16 ||

            (*this_mi)->bsize == BLOCK_32X16 ||

            (*this_mi)->bsize == BLOCK_16X32) {

          for (int j = 0; j < 4; j++) {

            if (vt->split[i].split[j].part_variances.none.variance <

                (thresholds[2] >> 8))

              part_info->variance_low[(i << 2) + j + 9] = 1;

          }

        }

      }

    }

  }

}

static AOM_INLINE void set_low_temp_var_flag_128x128(

    CommonModeInfoParams *mi_params, PartitionSearchInfo *part_info,

    MACROBLOCKD *xd, VP128x128 *vt, const int64_t thresholds[], int mi_col,

    int mi_row) {

  if (xd->mi[0]->bsize == BLOCK_128X128) {

    if (vt->part_variances.none.variance < (thresholds[0] >> 1))

      part_info->variance_low[0] = 1;

  } else if (xd->mi[0]->bsize == BLOCK_128X64) {

    for (int i = 0; i < 2; i++) {

      if (vt->part_variances.horz[i].variance < (thresholds[0] >> 2))

        part_info->variance_low[i + 1] = 1;

    }

  } else if (xd->mi[0]->bsize == BLOCK_64X128) {

    for (int i = 0; i < 2; i++) {

      if (vt->part_variances.vert[i].variance < (thresholds[0] >> 2))

        part_info->variance_low[i + 3] = 1;

    }

  } else {

    static const int idx64[4][2] = {

      { 0, 0 }, { 0, 16 }, { 16, 0 }, { 16, 16 }

    };

    static const int idx32[4][2] = { { 0, 0 }, { 0, 8 }, { 8, 0 }, { 8, 8 } };

    for (int i = 0; i < 4; i++) {

      const int idx_str =

          mi_params->mi_stride * (mi_row + idx64[i][0]) + mi_col + idx64[i][1];

      MB_MODE_INFO **mi_64 = mi_params->mi_grid_base + idx_str;

      if (*mi_64 == NULL) continue;

      if (mi_params->mi_cols <= mi_col + idx64[i][1] ||

          mi_params->mi_rows <= mi_row + idx64[i][0])

        continue;

      const int64_t threshold_64x64 = (5 * thresholds[1]) >> 3;

      if ((*mi_64)->bsize == BLOCK_64X64) {

        if (vt->split[i].part_variances.none.variance < threshold_64x64)

          part_info->variance_low[5 + i] = 1;

      } else if ((*mi_64)->bsize == BLOCK_64X32) {

        for (int j = 0; j < 2; j++)

          if (vt->split[i].part_variances.horz[j].variance <

              (threshold_64x64 >> 1))

            part_info->variance_low[9 + (i << 1) + j] = 1;

      } else if ((*mi_64)->bsize == BLOCK_32X64) {

        for (int j = 0; j < 2; j++)

          if (vt->split[i].part_variances.vert[j].variance <

              (threshold_64x64 >> 1))

            part_info->variance_low[17 + (i << 1) + j] = 1;

      } else {

        for (int k = 0; k < 4; k++) {

          const int idx_str1 = mi_params->mi_stride * idx32[k][0] + idx32[k][1];

          MB_MODE_INFO **mi_32 = mi_params->mi_grid_base + idx_str + idx_str1;

          if (*mi_32 == NULL) continue;

          if (mi_params->mi_cols <= mi_col + idx64[i][1] + idx32[k][1] ||

              mi_params->mi_rows <= mi_row + idx64[i][0] + idx32[k][0])

            continue;

          const int64_t threshold_32x32 = (5 * thresholds[2]) >> 3;

          if ((*mi_32)->bsize == BLOCK_32X32) {

            if (vt->split[i].split[k].part_variances.none.variance <

                threshold_32x32)

              part_info->variance_low[25 + (i << 2) + k] = 1;

          } else {

            // For 32x16 and 16x32 blocks, the flag is set on each 16x16 block

            // inside.

            if ((*mi_32)->bsize == BLOCK_16X16 ||

                (*mi_32)->bsize == BLOCK_32X16 ||

                (*mi_32)->bsize == BLOCK_16X32) {

              for (int j = 0; j < 4; j++) {

                if (vt->split[i]

                        .split[k]

                        .split[j]

                        .part_variances.none.variance < (thresholds[3] >> 8))

                  part_info->variance_low[41 + (i << 4) + (k << 2) + j] = 1;

              }

            }

          }

        }

      }

    }

  }

}

static AOM_INLINE void set_low_temp_var_flag(

    AV1_COMP *cpi, PartitionSearchInfo *part_info, MACROBLOCKD *xd,

    VP128x128 *vt, int64_t thresholds[], MV_REFERENCE_FRAME ref_frame_partition,

    int mi_col, int mi_row) {

  AV1_COMMON *const cm = &cpi->common;

  // Check temporal variance for bsize >= 16x16, if LAST_FRAME was selected.

  // If the temporal variance is small set the flag

  // variance_low for the block. The variance threshold can be adjusted, the

  // higher the more aggressive.

  if (ref_frame_partition == LAST_FRAME) {

    const int is_small_sb = (cm->seq_params->sb_size == BLOCK_64X64);

    if (is_small_sb)

      set_low_temp_var_flag_64x64(&cm->mi_params, part_info, xd,

                                  &(vt->split[0]), thresholds, mi_col, mi_row);

    else

      set_low_temp_var_flag_128x128(&cm->mi_params, part_info, xd, vt,

                                    thresholds, mi_col, mi_row);

  }

}

static const int pos_shift_16x16[4][4] = {

  { 9, 10, 13, 14 }, { 11, 12, 15, 16 }, { 17, 18, 21, 22 }, { 19, 20, 23, 24 }

};

int av1_get_force_skip_low_temp_var_small_sb(const uint8_t *variance_low,

                                             int mi_row, int mi_col,

                                             BLOCK_SIZE bsize) {

  // Relative indices of MB inside the superblock.

  const int mi_x = mi_row & 0xF;

  const int mi_y = mi_col & 0xF;

  // Relative indices of 16x16 block inside the superblock.

  const int i = mi_x >> 2;

  const int j = mi_y >> 2;

  int force_skip_low_temp_var = 0;

  // Set force_skip_low_temp_var based on the block size and block offset.

  switch (bsize) {

    case BLOCK_64X64: force_skip_low_temp_var = variance_low[0]; break;

    case BLOCK_64X32:

      if (!mi_y && !mi_x) {

        force_skip_low_temp_var = variance_low[1];

      } else if (!mi_y && mi_x) {

        force_skip_low_temp_var = variance_low[2];

      }

      break;

    case BLOCK_32X64:

      if (!mi_y && !mi_x) {

        force_skip_low_temp_var = variance_low[3];

      } else if (mi_y && !mi_x) {

        force_skip_low_temp_var = variance_low[4];

      }

      break;

    case BLOCK_32X32:

      if (!mi_y && !mi_x) {

        force_skip_low_temp_var = variance_low[5];

      } else if (mi_y && !mi_x) {

        force_skip_low_temp_var = variance_low[6];

      } else if (!mi_y && mi_x) {

        force_skip_low_temp_var = variance_low[7];

      } else if (mi_y && mi_x) {

        force_skip_low_temp_var = variance_low[8];

      }

      break;

    case BLOCK_32X16:

    case BLOCK_16X32:

    case BLOCK_16X16:

      force_skip_low_temp_var = variance_low[pos_shift_16x16[i][j]];

      break;

    default: break;

  }

  return force_skip_low_temp_var;

}

int av1_get_force_skip_low_temp_var(const uint8_t *variance_low, int mi_row,

                                    int mi_col, BLOCK_SIZE bsize) {

  int force_skip_low_temp_var = 0;

  int x, y;

  x = (mi_col & 0x1F) >> 4;

  // y = (mi_row & 0x1F) >> 4;

  // const int idx64 = (y << 1) + x;

  y = (mi_row & 0x17) >> 3;

  const int idx64 = y + x;

  x = (mi_col & 0xF) >> 3;

  // y = (mi_row & 0xF) >> 3;

  // const int idx32 = (y << 1) + x;

  y = (mi_row & 0xB) >> 2;

  const int idx32 = y + x;

  x = (mi_col & 0x7) >> 2;

  // y = (mi_row & 0x7) >> 2;

  // const int idx16 = (y << 1) + x;

  y = (mi_row & 0x5) >> 1;

  const int idx16 = y + x;

  // Set force_skip_low_temp_var based on the block size and block offset.

  switch (bsize) {

    case BLOCK_128X128: force_skip_low_temp_var = variance_low[0]; break;

    case BLOCK_128X64:

      assert((mi_col & 0x1F) == 0);

      force_skip_low_temp_var = variance_low[1 + ((mi_row & 0x1F) != 0)];

      break;

    case BLOCK_64X128:

      assert((mi_row & 0x1F) == 0);

      force_skip_low_temp_var = variance_low[3 + ((mi_col & 0x1F) != 0)];

      break;

    case BLOCK_64X64:

      // Location of this 64x64 block inside the 128x128 superblock

      force_skip_low_temp_var = variance_low[5 + idx64];

      break;

    case BLOCK_64X32:

      x = (mi_col & 0x1F) >> 4;

      y = (mi_row & 0x1F) >> 3;

      /*

      .---------------.---------------.

      | x=0,y=0,idx=0 | x=0,y=0,idx=2 |

      :---------------+---------------:

      | x=0,y=1,idx=1 | x=1,y=1,idx=3 |

      :---------------+---------------:

      | x=0,y=2,idx=4 | x=1,y=2,idx=6 |

      :---------------+---------------:

      | x=0,y=3,idx=5 | x=1,y=3,idx=7 |

      '---------------'---------------'

      */

      const int idx64x32 = (x << 1) + (y % 2) + ((y >> 1) << 2);

      force_skip_low_temp_var = variance_low[9 + idx64x32];

      break;

    case BLOCK_32X64:

      x = (mi_col & 0x1F) >> 3;

      y = (mi_row & 0x1F) >> 4;

      const int idx32x64 = (y << 2) + x;

      force_skip_low_temp_var = variance_low[17 + idx32x64];

      break;

    case BLOCK_32X32:

      force_skip_low_temp_var = variance_low[25 + (idx64 << 2) + idx32];

      break;

    case BLOCK_32X16:

    case BLOCK_16X32:

    case BLOCK_16X16:

      force_skip_low_temp_var =

          variance_low[41 + (idx64 << 4) + (idx32 << 2) + idx16];

      break;

    default: break;

  }

  return force_skip_low_temp_var;

}

void av1_set_variance_partition_thresholds(AV1_COMP *cpi, int q,

                                           int content_lowsumdiff) {

  SPEED_FEATURES *const sf = &cpi->sf;

  if (sf->part_sf.partition_search_type != VAR_BASED_PARTITION) {

    return;

  } else {

    set_vbp_thresholds(cpi, cpi->vbp_info.thresholds, q, content_lowsumdiff, 0,

                       0);

    // The threshold below is not changed locally.

    cpi->vbp_info.threshold_minmax = 15 + (q >> 3);

  }

}

static AOM_INLINE void chroma_check(AV1_COMP *cpi, MACROBLOCK *x,

                                    BLOCK_SIZE bsize, unsigned int y_sad,

                                    int is_key_frame) {

  int i;

  MACROBLOCKD *xd = &x->e_mbd;

  if (is_key_frame || cpi->oxcf.tool_cfg.enable_monochrome) return;

  for (i = 1; i <= 2; ++i) {

    unsigned int uv_sad = UINT_MAX;

    struct macroblock_plane *p = &x->plane[i];

    struct macroblockd_plane *pd = &xd->plane[i];

    const BLOCK_SIZE bs =

        get_plane_block_size(bsize, pd->subsampling_x, pd->subsampling_y);

    if (bs != BLOCK_INVALID)

      uv_sad = cpi->ppi->fn_ptr[bs].sdf(p->src.buf, p->src.stride, pd->dst.buf,

                                        pd->dst.stride);

    if (uv_sad > (y_sad >> 1))

      x->color_sensitivity_sb[i - 1] = 1;

    else if (uv_sad < (y_sad >> 3))

      x->color_sensitivity_sb[i - 1] = 0;

    // Borderline case: to be refined at coding block level in nonrd_pickmode,

    // for coding block size < sb_size.

    else

      x->color_sensitivity_sb[i - 1] = 2;

  }

}

static void fill_variance_tree_leaves(

    AV1_COMP *cpi, MACROBLOCK *x, VP128x128 *vt, VP16x16 *vt2,

    unsigned char *force_split, int avg_16x16[][4], int maxvar_16x16[][4],

    int minvar_16x16[][4], int *variance4x4downsample, int64_t *thresholds,

    uint8_t *src, int src_stride, const uint8_t *dst, int dst_stride) {

  AV1_COMMON *cm = &cpi->common;

  MACROBLOCKD *xd = &x->e_mbd;

  const int is_key_frame = frame_is_intra_only(cm);

  const int is_small_sb = (cm->seq_params->sb_size == BLOCK_64X64);

  const int num_64x64_blocks = is_small_sb ? 1 : 4;

  // TODO(kyslov) Bring back compute_minmax_variance with content type detection

  const int compute_minmax_variance = 0;

  const int segment_id = xd->mi[0]->segment_id;

  int pixels_wide = 128, pixels_high = 128;

  if (is_small_sb) {

    pixels_wide = 64;