/*

 * Copyright (c) 2020, 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 "av1/common/common_data.h"

#include "av1/common/quant_common.h"

#include "av1/common/reconintra.h"

#include "av1/encoder/encoder.h"

#include "av1/encoder/encodeframe_utils.h"

#include "av1/encoder/partition_strategy.h"

#include "av1/encoder/rdopt.h"

#include "av1/encoder/aq_variance.h"

void av1_set_ssim_rdmult(const AV1_COMP *const cpi, int *errorperbit,

                         const BLOCK_SIZE bsize, const int mi_row,

                         const int mi_col, int *const rdmult) {

  const AV1_COMMON *const cm = &cpi->common;

  const int bsize_base = BLOCK_16X16;

  const int num_mi_w = mi_size_wide[bsize_base];

  const int num_mi_h = mi_size_high[bsize_base];

  const int num_cols = (cm->mi_params.mi_cols + num_mi_w - 1) / num_mi_w;

  const int num_rows = (cm->mi_params.mi_rows + num_mi_h - 1) / num_mi_h;

  const int num_bcols = (mi_size_wide[bsize] + num_mi_w - 1) / num_mi_w;

  const int num_brows = (mi_size_high[bsize] + num_mi_h - 1) / num_mi_h;

  int row, col;

  double num_of_mi = 0.0;

  double geom_mean_of_scale = 0.0;

  assert(cpi->oxcf.tune_cfg.tuning == AOM_TUNE_SSIM);

  for (row = mi_row / num_mi_w;

       row < num_rows && row < mi_row / num_mi_w + num_brows; ++row) {

    for (col = mi_col / num_mi_h;

         col < num_cols && col < mi_col / num_mi_h + num_bcols; ++col) {

      const int index = row * num_cols + col;

      geom_mean_of_scale += log(cpi->ssim_rdmult_scaling_factors[index]);

      num_of_mi += 1.0;

    }

  }

  geom_mean_of_scale = exp(geom_mean_of_scale / num_of_mi);

  *rdmult = (int)((double)(*rdmult) * geom_mean_of_scale + 0.5);

  *rdmult = AOMMAX(*rdmult, 0);

  av1_set_error_per_bit(errorperbit, *rdmult);

}

// TODO(angiebird): Move these function to tpl_model.c

#if !CONFIG_REALTIME_ONLY

static AOM_INLINE int set_deltaq_rdmult(const AV1_COMP *const cpi,

                                        const MACROBLOCK *const x) {

  const AV1_COMMON *const cm = &cpi->common;

  const CommonQuantParams *quant_params = &cm->quant_params;

  return av1_compute_rd_mult(cpi, quant_params->base_qindex + x->delta_qindex +

                                      quant_params->y_dc_delta_q);

}

// Return the end column for the current superblock, in unit of TPL blocks.

static int get_superblock_tpl_column_end(const AV1_COMMON *const cm, int mi_col,

                                         int num_mi_w) {

  // Find the start column of this superblock.

  const int sb_mi_col_start = (mi_col >> cm->seq_params->mib_size_log2)

                              << cm->seq_params->mib_size_log2;

  // Same but in superres upscaled dimension.

  const int sb_mi_col_start_sr =

      coded_to_superres_mi(sb_mi_col_start, cm->superres_scale_denominator);

  // Width of this superblock in mi units.

  const int sb_mi_width = mi_size_wide[cm->seq_params->sb_size];

  // Same but in superres upscaled dimension.

  const int sb_mi_width_sr =

      coded_to_superres_mi(sb_mi_width, cm->superres_scale_denominator);

  // Superblock end in mi units.

  const int sb_mi_end = sb_mi_col_start_sr + sb_mi_width_sr;

  // Superblock end in TPL units.

  return (sb_mi_end + num_mi_w - 1) / num_mi_w;

}

int av1_get_cb_rdmult(const AV1_COMP *const cpi, MACROBLOCK *const x,

                      const BLOCK_SIZE bsize, const int mi_row,

                      const int mi_col) {

  const AV1_COMMON *const cm = &cpi->common;

  assert(IMPLIES(cpi->ppi->gf_group.size > 0,

                 cpi->gf_frame_index < cpi->ppi->gf_group.size));

  const int tpl_idx = cpi->gf_frame_index;

  int deltaq_rdmult = set_deltaq_rdmult(cpi, x);

  if (!av1_tpl_stats_ready(&cpi->ppi->tpl_data, tpl_idx)) return deltaq_rdmult;

  if (cm->superres_scale_denominator != SCALE_NUMERATOR) return deltaq_rdmult;

  if (cpi->oxcf.q_cfg.aq_mode != NO_AQ) return deltaq_rdmult;

  if (x->rb == 0) return deltaq_rdmult;

  TplParams *const tpl_data = &cpi->ppi->tpl_data;

  TplDepFrame *tpl_frame = &tpl_data->tpl_frame[tpl_idx];

  TplDepStats *tpl_stats = tpl_frame->tpl_stats_ptr;

  const int mi_wide = mi_size_wide[bsize];

  const int mi_high = mi_size_high[bsize];

  int tpl_stride = tpl_frame->stride;

  double intra_cost_base = 0;

  double mc_dep_cost_base = 0;

  double cbcmp_base = 0;

  const int step = 1 << tpl_data->tpl_stats_block_mis_log2;

  for (int row = mi_row; row < mi_row + mi_high; row += step) {

    for (int col = mi_col; col < mi_col + mi_wide; col += step) {

      if (row >= cm->mi_params.mi_rows || col >= cm->mi_params.mi_cols)

        continue;

      TplDepStats *this_stats = &tpl_stats[av1_tpl_ptr_pos(

          row, col, tpl_stride, tpl_data->tpl_stats_block_mis_log2)];

      double cbcmp = (double)this_stats->srcrf_dist;

      int64_t mc_dep_delta =

          RDCOST(tpl_frame->base_rdmult, this_stats->mc_dep_rate,

                 this_stats->mc_dep_dist);

      double dist_scaled = (double)(this_stats->recrf_dist << RDDIV_BITS);

      intra_cost_base += log(dist_scaled) * cbcmp;

      mc_dep_cost_base += log(3 * dist_scaled + mc_dep_delta) * cbcmp;

      cbcmp_base += cbcmp;

    }

  }

  if (cbcmp_base == 0) return deltaq_rdmult;

  double rk = exp((intra_cost_base - mc_dep_cost_base) / cbcmp_base);

  deltaq_rdmult = (int)(deltaq_rdmult * (rk / x->rb));

  return AOMMAX(deltaq_rdmult, 1);

}

int av1_get_hier_tpl_rdmult(const AV1_COMP *const cpi, MACROBLOCK *const x,

                            const BLOCK_SIZE bsize, const int mi_row,

                            const int mi_col, int orig_rdmult) {

  const AV1_COMMON *const cm = &cpi->common;

  const GF_GROUP *const gf_group = &cpi->ppi->gf_group;

  assert(IMPLIES(cpi->ppi->gf_group.size > 0,

                 cpi->gf_frame_index < cpi->ppi->gf_group.size));

  const int tpl_idx = cpi->gf_frame_index;

  const int deltaq_rdmult = set_deltaq_rdmult(cpi, x);

  if (!av1_tpl_stats_ready(&cpi->ppi->tpl_data, tpl_idx)) return deltaq_rdmult;

  if (!is_frame_tpl_eligible(gf_group, cpi->gf_frame_index))

    return deltaq_rdmult;

  if (cpi->oxcf.q_cfg.aq_mode != NO_AQ) return deltaq_rdmult;

  const int mi_col_sr =

      coded_to_superres_mi(mi_col, cm->superres_scale_denominator);

  const int mi_cols_sr = av1_pixels_to_mi(cm->superres_upscaled_width);

  const int block_mi_width_sr =

      coded_to_superres_mi(mi_size_wide[bsize], cm->superres_scale_denominator);

  const int bsize_base = BLOCK_16X16;

  const int num_mi_w = mi_size_wide[bsize_base];

  const int num_mi_h = mi_size_high[bsize_base];

  const int num_cols = (mi_cols_sr + num_mi_w - 1) / num_mi_w;

  const int num_rows = (cm->mi_params.mi_rows + num_mi_h - 1) / num_mi_h;

  const int num_bcols = (block_mi_width_sr + num_mi_w - 1) / num_mi_w;

  const int num_brows = (mi_size_high[bsize] + num_mi_h - 1) / num_mi_h;

  // This is required because the end col of superblock may be off by 1 in case

  // of superres.

  const int sb_bcol_end = get_superblock_tpl_column_end(cm, mi_col, num_mi_w);

  int row, col;

  double base_block_count = 0.0;

  double geom_mean_of_scale = 0.0;

  for (row = mi_row / num_mi_w;

       row < num_rows && row < mi_row / num_mi_w + num_brows; ++row) {

    for (col = mi_col_sr / num_mi_h;

         col < num_cols && col < mi_col_sr / num_mi_h + num_bcols &&

         col < sb_bcol_end;

         ++col) {

      const int index = row * num_cols + col;

      geom_mean_of_scale += log(cpi->ppi->tpl_sb_rdmult_scaling_factors[index]);

      base_block_count += 1.0;

    }

  }

  geom_mean_of_scale = exp(geom_mean_of_scale / base_block_count);

  int rdmult = (int)((double)orig_rdmult * geom_mean_of_scale + 0.5);

  rdmult = AOMMAX(rdmult, 0);

  av1_set_error_per_bit(&x->errorperbit, rdmult);

#if !CONFIG_RD_COMMAND

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

    const int rdmult_sb = set_deltaq_rdmult(cpi, x);

    assert(rdmult_sb == rdmult);

    (void)rdmult_sb;

  }

#endif  // !CONFIG_RD_COMMAND

  return rdmult;

}

#endif  // !CONFIG_REALTIME_ONLY

static AOM_INLINE void update_filter_type_count(FRAME_COUNTS *counts,

                                                const MACROBLOCKD *xd,

                                                const MB_MODE_INFO *mbmi) {

  int dir;

  for (dir = 0; dir < 2; ++dir) {

    const int ctx = av1_get_pred_context_switchable_interp(xd, dir);

    InterpFilter filter = av1_extract_interp_filter(mbmi->interp_filters, dir);

    ++counts->switchable_interp[ctx][filter];

  }

}

static void reset_tx_size(MACROBLOCK *x, MB_MODE_INFO *mbmi,

                          const TX_MODE tx_mode) {

  MACROBLOCKD *const xd = &x->e_mbd;

  TxfmSearchInfo *txfm_info = &x->txfm_search_info;

  if (xd->lossless[mbmi->segment_id]) {

    mbmi->tx_size = TX_4X4;

  } else if (tx_mode != TX_MODE_SELECT) {

    mbmi->tx_size = tx_size_from_tx_mode(mbmi->bsize, tx_mode);

  } else {

    const BLOCK_SIZE bsize = mbmi->bsize;

    const TX_SIZE min_tx_size = depth_to_tx_size(MAX_TX_DEPTH, bsize);

    if (tx_size_wide[min_tx_size] > tx_size_wide[mbmi->tx_size] ||

        tx_size_high[min_tx_size] > tx_size_high[mbmi->tx_size])

      mbmi->tx_size = min_tx_size;

    const TX_SIZE max_tx_size = get_vartx_max_txsize(xd, bsize, 0);

    if (tx_size_wide[max_tx_size] < tx_size_wide[mbmi->tx_size] ||

        tx_size_high[max_tx_size] < tx_size_high[mbmi->tx_size])

      mbmi->tx_size = max_tx_size;

  }

  if (is_inter_block(mbmi)) {

    memset(mbmi->inter_tx_size, mbmi->tx_size, sizeof(mbmi->inter_tx_size));

  }

  const int stride = xd->tx_type_map_stride;

  const int bw = mi_size_wide[mbmi->bsize];

  for (int row = 0; row < mi_size_high[mbmi->bsize]; ++row) {

    memset(xd->tx_type_map + row * stride, DCT_DCT,

           bw * sizeof(xd->tx_type_map[0]));

  }

  av1_zero(txfm_info->blk_skip);

  txfm_info->skip_txfm = 0;

}

// This function will copy the best reference mode information from

// MB_MODE_INFO_EXT_FRAME to MB_MODE_INFO_EXT.

static INLINE void copy_mbmi_ext_frame_to_mbmi_ext(

    MB_MODE_INFO_EXT *mbmi_ext,

    const MB_MODE_INFO_EXT_FRAME *const mbmi_ext_best, uint8_t ref_frame_type) {

  memcpy(mbmi_ext->ref_mv_stack[ref_frame_type], mbmi_ext_best->ref_mv_stack,

         sizeof(mbmi_ext->ref_mv_stack[USABLE_REF_MV_STACK_SIZE]));

  memcpy(mbmi_ext->weight[ref_frame_type], mbmi_ext_best->weight,

         sizeof(mbmi_ext->weight[USABLE_REF_MV_STACK_SIZE]));

  mbmi_ext->mode_context[ref_frame_type] = mbmi_ext_best->mode_context;

  mbmi_ext->ref_mv_count[ref_frame_type] = mbmi_ext_best->ref_mv_count;

  memcpy(mbmi_ext->global_mvs, mbmi_ext_best->global_mvs,

         sizeof(mbmi_ext->global_mvs));

}

void av1_update_state(const AV1_COMP *const cpi, ThreadData *td,

                      const PICK_MODE_CONTEXT *const ctx, int mi_row,

                      int mi_col, BLOCK_SIZE bsize, RUN_TYPE dry_run) {

  int i, x_idx, y;

  const AV1_COMMON *const cm = &cpi->common;

  const CommonModeInfoParams *const mi_params = &cm->mi_params;

  const int num_planes = av1_num_planes(cm);

  MACROBLOCK *const x = &td->mb;

  MACROBLOCKD *const xd = &x->e_mbd;

  struct macroblock_plane *const p = x->plane;

  struct macroblockd_plane *const pd = xd->plane;

  const MB_MODE_INFO *const mi = &ctx->mic;

  MB_MODE_INFO *const mi_addr = xd->mi[0];

  const struct segmentation *const seg = &cm->seg;

  assert(bsize < BLOCK_SIZES_ALL);

  const int bw = mi_size_wide[mi->bsize];

  const int bh = mi_size_high[mi->bsize];

  const int mis = mi_params->mi_stride;

  const int mi_width = mi_size_wide[bsize];

  const int mi_height = mi_size_high[bsize];

  TxfmSearchInfo *txfm_info = &x->txfm_search_info;

  assert(mi->bsize == bsize);

  *mi_addr = *mi;

  copy_mbmi_ext_frame_to_mbmi_ext(&x->mbmi_ext, &ctx->mbmi_ext_best,

                                  av1_ref_frame_type(ctx->mic.ref_frame));

  memcpy(txfm_info->blk_skip, ctx->blk_skip,

         sizeof(txfm_info->blk_skip[0]) * ctx->num_4x4_blk);

  txfm_info->skip_txfm = ctx->rd_stats.skip_txfm;

  xd->tx_type_map = ctx->tx_type_map;

  xd->tx_type_map_stride = mi_size_wide[bsize];

  // If not dry_run, copy the transform type data into the frame level buffer.

  // Encoder will fetch tx types when writing bitstream.

  if (!dry_run) {

    const int grid_idx = get_mi_grid_idx(mi_params, mi_row, mi_col);

    uint8_t *const tx_type_map = mi_params->tx_type_map + grid_idx;

    const int mi_stride = mi_params->mi_stride;

    for (int blk_row = 0; blk_row < bh; ++blk_row) {

      av1_copy_array(tx_type_map + blk_row * mi_stride,

                     xd->tx_type_map + blk_row * xd->tx_type_map_stride, bw);

    }

    xd->tx_type_map = tx_type_map;

    xd->tx_type_map_stride = mi_stride;

  }

  // If segmentation in use

  if (seg->enabled) {

    // For in frame complexity AQ copy the segment id from the segment map.

    if (cpi->oxcf.q_cfg.aq_mode == COMPLEXITY_AQ) {

      const uint8_t *const map =

          seg->update_map ? cpi->enc_seg.map : cm->last_frame_seg_map;

      mi_addr->segment_id =

          map ? get_segment_id(mi_params, map, bsize, mi_row, mi_col) : 0;

      reset_tx_size(x, mi_addr, x->txfm_search_params.tx_mode_search_type);

    }

    // Else for cyclic refresh mode update the segment map, set the segment id

    // and then update the quantizer.

    if (cpi->oxcf.q_cfg.aq_mode == CYCLIC_REFRESH_AQ &&

        !cpi->rc.rtc_external_ratectrl) {

      av1_cyclic_refresh_update_segment(cpi, x, mi_row, mi_col, bsize,

                                        ctx->rd_stats.rate, ctx->rd_stats.dist,

                                        txfm_info->skip_txfm, dry_run);

    }

    if (mi_addr->uv_mode == UV_CFL_PRED && !is_cfl_allowed(xd))

      mi_addr->uv_mode = UV_DC_PRED;

  }

  // Count zero motion vector.

  if (!dry_run && cpi->oxcf.q_cfg.aq_mode == CYCLIC_REFRESH_AQ &&

      !frame_is_intra_only(cm)) {

    const MV mv = mi->mv[0].as_mv;

    if (is_inter_block(mi) && mi->ref_frame[0] == LAST_FRAME &&

        abs(mv.row) < 8 && abs(mv.col) < 8) {

      const int ymis = AOMMIN(cm->mi_params.mi_rows - mi_row, bh);

      // Accumulate low_content_frame.

      for (int mi_y = 0; mi_y < ymis; mi_y += 2) x->cnt_zeromv += bw << 1;

    }

  }

  for (i = 0; i < num_planes; ++i) {

    p[i].coeff = ctx->coeff[i];

    p[i].qcoeff = ctx->qcoeff[i];

    p[i].dqcoeff = ctx->dqcoeff[i];

    p[i].eobs = ctx->eobs[i];

    p[i].txb_entropy_ctx = ctx->txb_entropy_ctx[i];

  }

  for (i = 0; i < 2; ++i) pd[i].color_index_map = ctx->color_index_map[i];

  // Restore the coding context of the MB to that that was in place

  // when the mode was picked for it

  for (y = 0; y < mi_height; y++) {

    for (x_idx = 0; x_idx < mi_width; x_idx++) {

      if ((xd->mb_to_right_edge >> (3 + MI_SIZE_LOG2)) + mi_width > x_idx &&

          (xd->mb_to_bottom_edge >> (3 + MI_SIZE_LOG2)) + mi_height > y) {

        xd->mi[x_idx + y * mis] = mi_addr;

      }

    }

  }

  if (cpi->oxcf.q_cfg.aq_mode)

    av1_init_plane_quantizers(cpi, x, mi_addr->segment_id);

  if (dry_run) return;

#if CONFIG_INTERNAL_STATS

  {

    unsigned int *const mode_chosen_counts =

        (unsigned int *)cpi->mode_chosen_counts;  // Cast const away.

    if (frame_is_intra_only(cm)) {

      static const int kf_mode_index[] = {

        THR_DC /*DC_PRED*/,

        THR_V_PRED /*V_PRED*/,

        THR_H_PRED /*H_PRED*/,

        THR_D45_PRED /*D45_PRED*/,

        THR_D135_PRED /*D135_PRED*/,

        THR_D113_PRED /*D113_PRED*/,

        THR_D157_PRED /*D157_PRED*/,

        THR_D203_PRED /*D203_PRED*/,

        THR_D67_PRED /*D67_PRED*/,

        THR_SMOOTH,   /*SMOOTH_PRED*/

        THR_SMOOTH_V, /*SMOOTH_V_PRED*/

        THR_SMOOTH_H, /*SMOOTH_H_PRED*/

        THR_PAETH /*PAETH_PRED*/,

      };

      ++mode_chosen_counts[kf_mode_index[mi_addr->mode]];

    } else {

      // Note how often each mode chosen as best

      ++mode_chosen_counts[ctx->best_mode_index];

    }

  }

#endif

  if (!frame_is_intra_only(cm)) {

    if (cm->features.interp_filter == SWITCHABLE &&

        mi_addr->motion_mode != WARPED_CAUSAL &&

        !is_nontrans_global_motion(xd, xd->mi[0])) {

      update_filter_type_count(td->counts, xd, mi_addr);

    }

  }

  const int x_mis = AOMMIN(bw, mi_params->mi_cols - mi_col);

  const int y_mis = AOMMIN(bh, mi_params->mi_rows - mi_row);

  if (cm->seq_params->order_hint_info.enable_ref_frame_mvs)

    av1_copy_frame_mvs(cm, mi, mi_row, mi_col, x_mis, y_mis);

}

void av1_update_inter_mode_stats(FRAME_CONTEXT *fc, FRAME_COUNTS *counts,

                                 PREDICTION_MODE mode, int16_t mode_context) {

  (void)counts;

  int16_t mode_ctx = mode_context & NEWMV_CTX_MASK;

  if (mode == NEWMV) {

#if CONFIG_ENTROPY_STATS

    ++counts->newmv_mode[mode_ctx][0];

#endif

    update_cdf(fc->newmv_cdf[mode_ctx], 0, 2);

    return;

  }

#if CONFIG_ENTROPY_STATS

  ++counts->newmv_mode[mode_ctx][1];

#endif

  update_cdf(fc->newmv_cdf[mode_ctx], 1, 2);

  mode_ctx = (mode_context >> GLOBALMV_OFFSET) & GLOBALMV_CTX_MASK;

  if (mode == GLOBALMV) {

#if CONFIG_ENTROPY_STATS

    ++counts->zeromv_mode[mode_ctx][0];

#endif

    update_cdf(fc->zeromv_cdf[mode_ctx], 0, 2);

    return;

  }

#if CONFIG_ENTROPY_STATS

  ++counts->zeromv_mode[mode_ctx][1];

#endif

  update_cdf(fc->zeromv_cdf[mode_ctx], 1, 2);

  mode_ctx = (mode_context >> REFMV_OFFSET) & REFMV_CTX_MASK;

#if CONFIG_ENTROPY_STATS

  ++counts->refmv_mode[mode_ctx][mode != NEARESTMV];

#endif

  update_cdf(fc->refmv_cdf[mode_ctx], mode != NEARESTMV, 2);

}

static void update_palette_cdf(MACROBLOCKD *xd, const MB_MODE_INFO *const mbmi,

                               FRAME_COUNTS *counts) {

  FRAME_CONTEXT *fc = xd->tile_ctx;

  const BLOCK_SIZE bsize = mbmi->bsize;

  const PALETTE_MODE_INFO *const pmi = &mbmi->palette_mode_info;

  const int palette_bsize_ctx = av1_get_palette_bsize_ctx(bsize);

  (void)counts;

  if (mbmi->mode == DC_PRED) {

    const int n = pmi->palette_size[0];

    const int palette_mode_ctx = av1_get_palette_mode_ctx(xd);

#if CONFIG_ENTROPY_STATS

    ++counts->palette_y_mode[palette_bsize_ctx][palette_mode_ctx][n > 0];

#endif

    update_cdf(fc->palette_y_mode_cdf[palette_bsize_ctx][palette_mode_ctx],

               n > 0, 2);

    if (n > 0) {

#if CONFIG_ENTROPY_STATS

      ++counts->palette_y_size[palette_bsize_ctx][n - PALETTE_MIN_SIZE];

#endif

      update_cdf(fc->palette_y_size_cdf[palette_bsize_ctx],

                 n - PALETTE_MIN_SIZE, PALETTE_SIZES);

    }

  }

  if (mbmi->uv_mode == UV_DC_PRED) {

    const int n = pmi->palette_size[1];

    const int palette_uv_mode_ctx = (pmi->palette_size[0] > 0);

#if CONFIG_ENTROPY_STATS

    ++counts->palette_uv_mode[palette_uv_mode_ctx][n > 0];

#endif

    update_cdf(fc->palette_uv_mode_cdf[palette_uv_mode_ctx], n > 0, 2);

    if (n > 0) {

#if CONFIG_ENTROPY_STATS

      ++counts->palette_uv_size[palette_bsize_ctx][n - PALETTE_MIN_SIZE];

#endif

      update_cdf(fc->palette_uv_size_cdf[palette_bsize_ctx],

                 n - PALETTE_MIN_SIZE, PALETTE_SIZES);

    }

  }

}

void av1_sum_intra_stats(const AV1_COMMON *const cm, FRAME_COUNTS *counts,

                         MACROBLOCKD *xd, const MB_MODE_INFO *const mbmi,

                         const MB_MODE_INFO *above_mi,

                         const MB_MODE_INFO *left_mi, const int intraonly) {

  FRAME_CONTEXT *fc = xd->tile_ctx;

  const PREDICTION_MODE y_mode = mbmi->mode;

  (void)counts;

  const BLOCK_SIZE bsize = mbmi->bsize;

  if (intraonly) {

#if CONFIG_ENTROPY_STATS

    const PREDICTION_MODE above = av1_above_block_mode(above_mi);

    const PREDICTION_MODE left = av1_left_block_mode(left_mi);

    const int above_ctx = intra_mode_context[above];

    const int left_ctx = intra_mode_context[left];

    ++counts->kf_y_mode[above_ctx][left_ctx][y_mode];

#endif  // CONFIG_ENTROPY_STATS

    update_cdf(get_y_mode_cdf(fc, above_mi, left_mi), y_mode, INTRA_MODES);

  } else {

#if CONFIG_ENTROPY_STATS

    ++counts->y_mode[size_group_lookup[bsize]][y_mode];

#endif  // CONFIG_ENTROPY_STATS

    update_cdf(fc->y_mode_cdf[size_group_lookup[bsize]], y_mode, INTRA_MODES);

  }

  if (av1_filter_intra_allowed(cm, mbmi)) {

    const int use_filter_intra_mode =

        mbmi->filter_intra_mode_info.use_filter_intra;

#if CONFIG_ENTROPY_STATS

    ++counts->filter_intra[mbmi->bsize][use_filter_intra_mode];

    if (use_filter_intra_mode) {

      ++counts

            ->filter_intra_mode[mbmi->filter_intra_mode_info.filter_intra_mode];

    }

#endif  // CONFIG_ENTROPY_STATS

    update_cdf(fc->filter_intra_cdfs[mbmi->bsize], use_filter_intra_mode, 2);

    if (use_filter_intra_mode) {

      update_cdf(fc->filter_intra_mode_cdf,

                 mbmi->filter_intra_mode_info.filter_intra_mode,

                 FILTER_INTRA_MODES);

    }

  }

  if (av1_is_directional_mode(mbmi->mode) && av1_use_angle_delta(bsize)) {

#if CONFIG_ENTROPY_STATS

    ++counts->angle_delta[mbmi->mode - V_PRED]

                         [mbmi->angle_delta[PLANE_TYPE_Y] + MAX_ANGLE_DELTA];

#endif

    update_cdf(fc->angle_delta_cdf[mbmi->mode - V_PRED],

               mbmi->angle_delta[PLANE_TYPE_Y] + MAX_ANGLE_DELTA,

               2 * MAX_ANGLE_DELTA + 1);

  }

  if (!xd->is_chroma_ref) return;

  const UV_PREDICTION_MODE uv_mode = mbmi->uv_mode;

  const CFL_ALLOWED_TYPE cfl_allowed = is_cfl_allowed(xd);

#if CONFIG_ENTROPY_STATS

  ++counts->uv_mode[cfl_allowed][y_mode][uv_mode];

#endif  // CONFIG_ENTROPY_STATS

  update_cdf(fc->uv_mode_cdf[cfl_allowed][y_mode], uv_mode,

             UV_INTRA_MODES - !cfl_allowed);

  if (uv_mode == UV_CFL_PRED) {

    const int8_t joint_sign = mbmi->cfl_alpha_signs;

    const uint8_t idx = mbmi->cfl_alpha_idx;

#if CONFIG_ENTROPY_STATS

    ++counts->cfl_sign[joint_sign];

#endif

    update_cdf(fc->cfl_sign_cdf, joint_sign, CFL_JOINT_SIGNS);

    if (CFL_SIGN_U(joint_sign) != CFL_SIGN_ZERO) {

      aom_cdf_prob *cdf_u = fc->cfl_alpha_cdf[CFL_CONTEXT_U(joint_sign)];

#if CONFIG_ENTROPY_STATS

      ++counts->cfl_alpha[CFL_CONTEXT_U(joint_sign)][CFL_IDX_U(idx)];

#endif

      update_cdf(cdf_u, CFL_IDX_U(idx), CFL_ALPHABET_SIZE);

    }

    if (CFL_SIGN_V(joint_sign) != CFL_SIGN_ZERO) {

      aom_cdf_prob *cdf_v = fc->cfl_alpha_cdf[CFL_CONTEXT_V(joint_sign)];

#if CONFIG_ENTROPY_STATS

      ++counts->cfl_alpha[CFL_CONTEXT_V(joint_sign)][CFL_IDX_V(idx)];

#endif

      update_cdf(cdf_v, CFL_IDX_V(idx), CFL_ALPHABET_SIZE);

    }

  }

  if (av1_is_directional_mode(get_uv_mode(uv_mode)) &&

      av1_use_angle_delta(bsize)) {

#if CONFIG_ENTROPY_STATS

    ++counts->angle_delta[uv_mode - UV_V_PRED]

                         [mbmi->angle_delta[PLANE_TYPE_UV] + MAX_ANGLE_DELTA];

#endif

    update_cdf(fc->angle_delta_cdf[uv_mode - UV_V_PRED],

               mbmi->angle_delta[PLANE_TYPE_UV] + MAX_ANGLE_DELTA,

               2 * MAX_ANGLE_DELTA + 1);

  }

  if (av1_allow_palette(cm->features.allow_screen_content_tools, bsize)) {

    update_palette_cdf(xd, mbmi, counts);

  }

}

void av1_restore_context(MACROBLOCK *x, const RD_SEARCH_MACROBLOCK_CONTEXT *ctx,

                         int mi_row, int mi_col, BLOCK_SIZE bsize,

                         const int num_planes) {

  MACROBLOCKD *xd = &x->e_mbd;

  int p;

  const int num_4x4_blocks_wide = mi_size_wide[bsize];

  const int num_4x4_blocks_high = mi_size_high[bsize];

  int mi_width = mi_size_wide[bsize];

  int mi_height = mi_size_high[bsize];

  for (p = 0; p < num_planes; p++) {

    int tx_col = mi_col;

    int tx_row = mi_row & MAX_MIB_MASK;

    memcpy(

        xd->above_entropy_context[p] + (tx_col >> xd->plane[p].subsampling_x),

        ctx->a + num_4x4_blocks_wide * p,

        (sizeof(ENTROPY_CONTEXT) * num_4x4_blocks_wide) >>

            xd->plane[p].subsampling_x);

    memcpy(xd->left_entropy_context[p] + (tx_row >> xd->plane[p].subsampling_y),

           ctx->l + num_4x4_blocks_high * p,

           (sizeof(ENTROPY_CONTEXT) * num_4x4_blocks_high) >>

               xd->plane[p].subsampling_y);

  }

  memcpy(xd->above_partition_context + mi_col, ctx->sa,

         sizeof(*xd->above_partition_context) * mi_width);

  memcpy(xd->left_partition_context + (mi_row & MAX_MIB_MASK), ctx->sl,

         sizeof(xd->left_partition_context[0]) * mi_height);

  xd->above_txfm_context = ctx->p_ta;

  xd->left_txfm_context = ctx->p_tl;

  memcpy(xd->above_txfm_context, ctx->ta,

         sizeof(*xd->above_txfm_context) * mi_width);

  memcpy(xd->left_txfm_context, ctx->tl,

         sizeof(*xd->left_txfm_context) * mi_height);

}

void av1_save_context(const MACROBLOCK *x, RD_SEARCH_MACROBLOCK_CONTEXT *ctx,

                      int mi_row, int mi_col, BLOCK_SIZE bsize,

                      const int num_planes) {

  const MACROBLOCKD *xd = &x->e_mbd;

  int p;

  int mi_width = mi_size_wide[bsize];

  int mi_height = mi_size_high[bsize];

  // buffer the above/left context information of the block in search.

  for (p = 0; p < num_planes; ++p) {

    int tx_col = mi_col;

    int tx_row = mi_row & MAX_MIB_MASK;

    memcpy(

        ctx->a + mi_width * p,

        xd->above_entropy_context[p] + (tx_col >> xd->plane[p].subsampling_x),

        (sizeof(ENTROPY_CONTEXT) * mi_width) >> xd->plane[p].subsampling_x);

    memcpy(ctx->l + mi_height * p,

           xd->left_entropy_context[p] + (tx_row >> xd->plane[p].subsampling_y),

           (sizeof(ENTROPY_CONTEXT) * mi_height) >> xd->plane[p].subsampling_y);

  }

  memcpy(ctx->sa, xd->above_partition_context + mi_col,

         sizeof(*xd->above_partition_context) * mi_width);

  memcpy(ctx->sl, xd->left_partition_context + (mi_row & MAX_MIB_MASK),

         sizeof(xd->left_partition_context[0]) * mi_height);

  memcpy(ctx->ta, xd->above_txfm_context,

         sizeof(*xd->above_txfm_context) * mi_width);

  memcpy(ctx->tl, xd->left_txfm_context,

         sizeof(*xd->left_txfm_context) * mi_height);

  ctx->p_ta = xd->above_txfm_context;

  ctx->p_tl = xd->left_txfm_context;

}

static void set_partial_sb_partition(const AV1_COMMON *const cm,

                                     MB_MODE_INFO *mi, int bh_in, int bw_in,

                                     int mi_rows_remaining,

                                     int mi_cols_remaining, BLOCK_SIZE bsize,

                                     MB_MODE_INFO **mib) {

  int bh = bh_in;

  int r, c;

  for (r = 0; r < cm->seq_params->mib_size; r += bh) {

    int bw = bw_in;

    for (c = 0; c < cm->seq_params->mib_size; c += bw) {

      const int grid_index = get_mi_grid_idx(&cm->mi_params, r, c);

      const int mi_index = get_alloc_mi_idx(&cm->mi_params, r, c);

      mib[grid_index] = mi + mi_index;

      mib[grid_index]->bsize = find_partition_size(

          bsize, mi_rows_remaining - r, mi_cols_remaining - c, &bh, &bw);

    }

  }

}

// This function attempts to set all mode info entries in a given superblock

// to the same block partition size.

// However, at the bottom and right borders of the image the requested size

// may not be allowed in which case this code attempts to choose the largest

// allowable partition.

void av1_set_fixed_partitioning(AV1_COMP *cpi, const TileInfo *const tile,

                                MB_MODE_INFO **mib, int mi_row, int mi_col,

                                BLOCK_SIZE bsize) {

  AV1_COMMON *const cm = &cpi->common;

  const CommonModeInfoParams *const mi_params = &cm->mi_params;

  const int mi_rows_remaining = tile->mi_row_end - mi_row;

  const int mi_cols_remaining = tile->mi_col_end - mi_col;

  MB_MODE_INFO *const mi_upper_left =

      mi_params->mi_alloc + get_alloc_mi_idx(mi_params, mi_row, mi_col);

  int bh = mi_size_high[bsize];

  int bw = mi_size_wide[bsize];

  assert(bsize >= mi_params->mi_alloc_bsize &&

         "Attempted to use bsize < mi_params->mi_alloc_bsize");

  assert((mi_rows_remaining > 0) && (mi_cols_remaining > 0));

  // Apply the requested partition size to the SB if it is all "in image"

  if ((mi_cols_remaining >= cm->seq_params->mib_size) &&

      (mi_rows_remaining >= cm->seq_params->mib_size)) {

    for (int block_row = 0; block_row < cm->seq_params->mib_size;

         block_row += bh) {

      for (int block_col = 0; block_col < cm->seq_params->mib_size;

           block_col += bw) {

        const int grid_index = get_mi_grid_idx(mi_params, block_row, block_col);

        const int mi_index = get_alloc_mi_idx(mi_params, block_row, block_col);

        mib[grid_index] = mi_upper_left + mi_index;

        mib[grid_index]->bsize = bsize;

      }

    }

  } else {

    // Else this is a partial SB.

    set_partial_sb_partition(cm, mi_upper_left, bh, bw, mi_rows_remaining,

                             mi_cols_remaining, bsize, mib);

  }

}

int av1_is_leaf_split_partition(AV1_COMMON *cm, int mi_row, int mi_col,

                                BLOCK_SIZE bsize) {

  const int bs = mi_size_wide[bsize];

  const int hbs = bs / 2;

  assert(bsize >= BLOCK_8X8);

  const BLOCK_SIZE subsize = get_partition_subsize(bsize, PARTITION_SPLIT);

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

    int x_idx = (i & 1) * hbs;

    int y_idx = (i >> 1) * hbs;

    if ((mi_row + y_idx >= cm->mi_params.mi_rows) ||

        (mi_col + x_idx >= cm->mi_params.mi_cols))

      return 0;

    if (get_partition(cm, mi_row + y_idx, mi_col + x_idx, subsize) !=

            PARTITION_NONE &&

        subsize != BLOCK_8X8)

      return 0;

  }

  return 1;

}

#if !CONFIG_REALTIME_ONLY

int av1_get_rdmult_delta(AV1_COMP *cpi, BLOCK_SIZE bsize, int mi_row,

                         int mi_col, int orig_rdmult) {

  AV1_COMMON *const cm = &cpi->common;

  const GF_GROUP *const gf_group = &cpi->ppi->gf_group;

  assert(IMPLIES(cpi->ppi->gf_group.size > 0,

                 cpi->gf_frame_index < cpi->ppi->gf_group.size));

  const int tpl_idx = cpi->gf_frame_index;

  TplParams *const tpl_data = &cpi->ppi->tpl_data;

  const uint8_t block_mis_log2 = tpl_data->tpl_stats_block_mis_log2;

  int64_t intra_cost = 0;

  int64_t mc_dep_cost = 0;

  const int mi_wide = mi_size_wide[bsize];

  const int mi_high = mi_size_high[bsize];

  TplDepFrame *tpl_frame = &tpl_data->tpl_frame[tpl_idx];

  TplDepStats *tpl_stats = tpl_frame->tpl_stats_ptr;

  int tpl_stride = tpl_frame->stride;

  if (!av1_tpl_stats_ready(&cpi->ppi->tpl_data, cpi->gf_frame_index)) {

    return orig_rdmult;

  }

  if (!is_frame_tpl_eligible(gf_group, cpi->gf_frame_index)) {

    return orig_rdmult;

  }

  int mi_count = 0;

  const int mi_col_sr =

      coded_to_superres_mi(mi_col, cm->superres_scale_denominator);

  const int mi_col_end_sr =

      coded_to_superres_mi(mi_col + mi_wide, cm->superres_scale_denominator);

  const int mi_cols_sr = av1_pixels_to_mi(cm->superres_upscaled_width);

  const int step = 1 << block_mis_log2;

  const int row_step = step;

  const int col_step_sr =

      coded_to_superres_mi(step, cm->superres_scale_denominator);

  for (int row = mi_row; row < mi_row + mi_high; row += row_step) {

    for (int col = mi_col_sr; col < mi_col_end_sr; col += col_step_sr) {

      if (row >= cm->mi_params.mi_rows || col >= mi_cols_sr) continue;

      TplDepStats *this_stats =

          &tpl_stats[av1_tpl_ptr_pos(row, col, tpl_stride, block_mis_log2)];

      int64_t mc_dep_delta =

          RDCOST(tpl_frame->base_rdmult, this_stats->mc_dep_rate,

                 this_stats->mc_dep_dist);

      intra_cost += this_stats->recrf_dist << RDDIV_BITS;

      mc_dep_cost += (this_stats->recrf_dist << RDDIV_BITS) + mc_dep_delta;

      mi_count++;

    }

  }

  assert(mi_count <= MAX_TPL_BLK_IN_SB * MAX_TPL_BLK_IN_SB);

  double beta = 1.0;

  if (mc_dep_cost > 0 && intra_cost > 0) {

    const double r0 = cpi->rd.r0;

    const double rk = (double)intra_cost / mc_dep_cost;

    beta = (r0 / rk);

  }

  int rdmult = av1_get_adaptive_rdmult(cpi, beta);

  rdmult = AOMMIN(rdmult, orig_rdmult * 3 / 2);

  rdmult = AOMMAX(rdmult, orig_rdmult * 1 / 2);

  rdmult = AOMMAX(1, rdmult);

  return rdmult;

}

// Checks to see if a super block is on a horizontal image edge.

// In most cases this is the "real" edge unless there are formatting

// bars embedded in the stream.

int av1_active_h_edge(const AV1_COMP *cpi, int mi_row, int mi_step) {

  int top_edge = 0;

  int bottom_edge = cpi->common.mi_params.mi_rows;

  int is_active_h_edge = 0;

  // For two pass account for any formatting bars detected.

  if (is_stat_consumption_stage_twopass(cpi)) {

    const AV1_COMMON *const cm = &cpi->common;

    const FIRSTPASS_STATS *const this_frame_stats = read_one_frame_stats(

        &cpi->ppi->twopass, cm->current_frame.display_order_hint);

    if (this_frame_stats == NULL) return AOM_CODEC_ERROR;

    // The inactive region is specified in MBs not mi units.

    // The image edge is in the following MB row.

    top_edge += (int)(this_frame_stats->inactive_zone_rows * 4);

    bottom_edge -= (int)(this_frame_stats->inactive_zone_rows * 4);

    bottom_edge = AOMMAX(top_edge, bottom_edge);

  }

  if (((top_edge >= mi_row) && (top_edge < (mi_row + mi_step))) ||

      ((bottom_edge >= mi_row) && (bottom_edge < (mi_row + mi_step)))) {

    is_active_h_edge = 1;

  }

  return is_active_h_edge;

}

// Checks to see if a super block is on a vertical image edge.

// In most cases this is the "real" edge unless there are formatting

// bars embedded in the stream.

int av1_active_v_edge(const AV1_COMP *cpi, int mi_col, int mi_step) {

  int left_edge = 0;

  int right_edge = cpi->common.mi_params.mi_cols;

  int is_active_v_edge = 0;

  // For two pass account for any formatting bars detected.

  if (is_stat_consumption_stage_twopass(cpi)) {

    const AV1_COMMON *const cm = &cpi->common;

    const FIRSTPASS_STATS *const this_frame_stats = read_one_frame_stats(

        &cpi->ppi->twopass, cm->current_frame.display_order_hint);

    if (this_frame_stats == NULL) return AOM_CODEC_ERROR;

    // The inactive region is specified in MBs not mi units.

    // The image edge is in the following MB row.

    left_edge += (int)(this_frame_stats->inactive_zone_cols * 4);

    right_edge -= (int)(this_frame_stats->inactive_zone_cols * 4);

    right_edge = AOMMAX(left_edge, right_edge);

  }

  if (((left_edge >= mi_col) && (left_edge < (mi_col + mi_step))) ||

      ((right_edge >= mi_col) && (right_edge < (mi_col + mi_step)))) {

    is_active_v_edge = 1;

  }

  return is_active_v_edge;

}

void av1_get_tpl_stats_sb(AV1_COMP *cpi, BLOCK_SIZE bsize, int mi_row,

                          int mi_col, SuperBlockEnc *sb_enc) {

  sb_enc->tpl_data_count = 0;

  if (!cpi->oxcf.algo_cfg.enable_tpl_model) return;

  if (cpi->common.current_frame.frame_type == KEY_FRAME) return;

  const FRAME_UPDATE_TYPE update_type =

      get_frame_update_type(&cpi->ppi->gf_group, cpi->gf_frame_index);

  if (update_type == INTNL_OVERLAY_UPDATE || update_type == OVERLAY_UPDATE)

    return;

  assert(IMPLIES(cpi->ppi->gf_group.size > 0,

                 cpi->gf_frame_index < cpi->ppi->gf_group.size));

  AV1_COMMON *const cm = &cpi->common;

  const int gf_group_index = cpi->gf_frame_index;

  TplParams *const tpl_data = &cpi->ppi->tpl_data;

  if (!av1_tpl_stats_ready(tpl_data, gf_group_index)) return;

  const int mi_wide = mi_size_wide[bsize];

  const int mi_high = mi_size_high[bsize];

  TplDepFrame *tpl_frame = &tpl_data->tpl_frame[gf_group_index];

  TplDepStats *tpl_stats = tpl_frame->tpl_stats_ptr;

  int tpl_stride = tpl_frame->stride;

  int mi_count = 0;

  int count = 0;

  const int mi_col_sr =

      coded_to_superres_mi(mi_col, cm->superres_scale_denominator);

  const int mi_col_end_sr =

      coded_to_superres_mi(mi_col + mi_wide, cm->superres_scale_denominator);

  // mi_cols_sr is mi_cols at superres case.

  const int mi_cols_sr = av1_pixels_to_mi(cm->superres_upscaled_width);

  // TPL store unit size is not the same as the motion estimation unit size.

  // Here always use motion estimation size to avoid getting repetitive inter/

  // intra cost.

  const BLOCK_SIZE tpl_bsize = convert_length_to_bsize(tpl_data->tpl_bsize_1d);

  assert(mi_size_wide[tpl_bsize] == mi_size_high[tpl_bsize]);

  const int row_step = mi_size_high[tpl_bsize];

  const int col_step_sr = coded_to_superres_mi(mi_size_wide[tpl_bsize],

                                               cm->superres_scale_denominator);

  // Stride is only based on SB size, and we fill in values for every 16x16

  // block in a SB.

  sb_enc->tpl_stride = (mi_col_end_sr - mi_col_sr) / col_step_sr;

  for (int row = mi_row; row < mi_row + mi_high; row += row_step) {

    for (int col = mi_col_sr; col < mi_col_end_sr; col += col_step_sr) {

      assert(count < MAX_TPL_BLK_IN_SB * MAX_TPL_BLK_IN_SB);

      // Handle partial SB, so that no invalid values are used later.

      if (row >= cm->mi_params.mi_rows || col >= mi_cols_sr) {

        sb_enc->tpl_inter_cost[count] = INT64_MAX;

        sb_enc->tpl_intra_cost[count] = INT64_MAX;

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

          sb_enc->tpl_mv[count][i].as_int = INVALID_MV;

        }

        count++;

        continue;

      }

      TplDepStats *this_stats = &tpl_stats[av1_tpl_ptr_pos(

          row, col, tpl_stride, tpl_data->tpl_stats_block_mis_log2)];