/*

 * 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 <assert.h>

#include <float.h>

#include <stdint.h>

#include "config/aom_config.h"

#if CONFIG_THREE_PASS

#include "av1/encoder/thirdpass.h"

#endif

#include "config/aom_dsp_rtcd.h"

#include "config/aom_scale_rtcd.h"

#include "aom/aom_codec.h"

#include "aom_util/aom_pthread.h"

#include "av1/common/av1_common_int.h"

#include "av1/common/enums.h"

#include "av1/common/idct.h"

#include "av1/common/reconintra.h"

#include "av1/encoder/encoder.h"

#include "av1/encoder/ethread.h"

#include "av1/encoder/encodeframe_utils.h"

#include "av1/encoder/encode_strategy.h"

#include "av1/encoder/hybrid_fwd_txfm.h"

#include "av1/encoder/motion_search_facade.h"

#include "av1/encoder/rd.h"

#include "av1/encoder/rdopt.h"

#include "av1/encoder/reconinter_enc.h"

#include "av1/encoder/tpl_model.h"

static inline double exp_bounded(double v) {

  // When v > 700 or <-700, the exp function will be close to overflow

  // For details, see the "Notes" in the following link.

  // https://en.cppreference.com/w/c/numeric/math/exp

  if (v > 700) {

    return DBL_MAX;

  } else if (v < -700) {

    return 0;

  }

  return exp(v);

}

void av1_init_tpl_txfm_stats(TplTxfmStats *tpl_txfm_stats) {

  tpl_txfm_stats->ready = 0;

  tpl_txfm_stats->coeff_num = 256;

  tpl_txfm_stats->txfm_block_count = 0;

  memset(tpl_txfm_stats->abs_coeff_sum, 0,

         sizeof(tpl_txfm_stats->abs_coeff_sum[0]) * tpl_txfm_stats->coeff_num);

  memset(tpl_txfm_stats->abs_coeff_mean, 0,

         sizeof(tpl_txfm_stats->abs_coeff_mean[0]) * tpl_txfm_stats->coeff_num);

}

#if CONFIG_BITRATE_ACCURACY

void av1_accumulate_tpl_txfm_stats(const TplTxfmStats *sub_stats,

                                   TplTxfmStats *accumulated_stats) {

  accumulated_stats->txfm_block_count += sub_stats->txfm_block_count;

  for (int i = 0; i < accumulated_stats->coeff_num; ++i) {

    accumulated_stats->abs_coeff_sum[i] += sub_stats->abs_coeff_sum[i];

  }

}

void av1_record_tpl_txfm_block(TplTxfmStats *tpl_txfm_stats,

                               const tran_low_t *coeff) {

  // For transform larger than 16x16, the scale of coeff need to be adjusted.

  // It's not LOSSLESS_Q_STEP.

  assert(tpl_txfm_stats->coeff_num <= 256);

  for (int i = 0; i < tpl_txfm_stats->coeff_num; ++i) {

    tpl_txfm_stats->abs_coeff_sum[i] += abs(coeff[i]) / (double)LOSSLESS_Q_STEP;

  }

  ++tpl_txfm_stats->txfm_block_count;

}

void av1_tpl_txfm_stats_update_abs_coeff_mean(TplTxfmStats *txfm_stats) {

  if (txfm_stats->txfm_block_count > 0) {

    for (int j = 0; j < txfm_stats->coeff_num; j++) {

      txfm_stats->abs_coeff_mean[j] =

          txfm_stats->abs_coeff_sum[j] / txfm_stats->txfm_block_count;

    }

    txfm_stats->ready = 1;

  } else {

    txfm_stats->ready = 0;

  }

}

static inline void av1_tpl_store_txfm_stats(TplParams *tpl_data,

                                            const TplTxfmStats *tpl_txfm_stats,

                                            const int frame_index) {

  tpl_data->txfm_stats_list[frame_index] = *tpl_txfm_stats;

}

#endif  // CONFIG_BITRATE_ACCURACY

static inline void get_quantize_error(const MACROBLOCK *x, int plane,

                                      const tran_low_t *coeff,

                                      tran_low_t *qcoeff, tran_low_t *dqcoeff,

                                      TX_SIZE tx_size, uint16_t *eob,

                                      int64_t *recon_error, int64_t *sse) {

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

  const MACROBLOCKD *xd = &x->e_mbd;

  const SCAN_ORDER *const scan_order = &av1_scan_orders[tx_size][DCT_DCT];

  int pix_num = 1 << num_pels_log2_lookup[txsize_to_bsize[tx_size]];

  const int shift = tx_size == TX_32X32 ? 0 : 2;

  QUANT_PARAM quant_param;

  av1_setup_quant(tx_size, 0, AV1_XFORM_QUANT_FP, 0, &quant_param);

#if CONFIG_AV1_HIGHBITDEPTH

  if (is_cur_buf_hbd(xd)) {

    av1_highbd_quantize_fp_facade(coeff, pix_num, p, qcoeff, dqcoeff, eob,

                                  scan_order, &quant_param);

    *recon_error =

        av1_highbd_block_error(coeff, dqcoeff, pix_num, sse, xd->bd) >> shift;

  } else {

    av1_quantize_fp_facade(coeff, pix_num, p, qcoeff, dqcoeff, eob, scan_order,

                           &quant_param);

    *recon_error = av1_block_error(coeff, dqcoeff, pix_num, sse) >> shift;

  }

#else

  (void)xd;

  av1_quantize_fp_facade(coeff, pix_num, p, qcoeff, dqcoeff, eob, scan_order,

                         &quant_param);

  *recon_error = av1_block_error(coeff, dqcoeff, pix_num, sse) >> shift;

#endif  // CONFIG_AV1_HIGHBITDEPTH

  *recon_error = AOMMAX(*recon_error, 1);

  *sse = (*sse) >> shift;

  *sse = AOMMAX(*sse, 1);

}

static inline void set_tpl_stats_block_size(uint8_t *block_mis_log2,

                                            uint8_t *tpl_bsize_1d) {

  // tpl stats bsize: 2 means 16x16

  *block_mis_log2 = 2;

  // Block size used in tpl motion estimation

  *tpl_bsize_1d = 16;

  // MIN_TPL_BSIZE_1D = 16;

  assert(*tpl_bsize_1d >= 16);

}

void av1_setup_tpl_buffers(AV1_PRIMARY *const ppi,

                           CommonModeInfoParams *const mi_params, int width,

                           int height, int byte_alignment, int lag_in_frames) {

  SequenceHeader *const seq_params = &ppi->seq_params;

  TplParams *const tpl_data = &ppi->tpl_data;

  set_tpl_stats_block_size(&tpl_data->tpl_stats_block_mis_log2,

                           &tpl_data->tpl_bsize_1d);

  const uint8_t block_mis_log2 = tpl_data->tpl_stats_block_mis_log2;

  tpl_data->border_in_pixels =

      ALIGN_POWER_OF_TWO(tpl_data->tpl_bsize_1d + 2 * AOM_INTERP_EXTEND, 5);

  const int alloc_y_plane_only =

      ppi->cpi->sf.tpl_sf.use_y_only_rate_distortion ? 1 : 0;

  for (int frame = 0; frame < MAX_LENGTH_TPL_FRAME_STATS; ++frame) {

    const int mi_cols =

        ALIGN_POWER_OF_TWO(mi_params->mi_cols, MAX_MIB_SIZE_LOG2);

    const int mi_rows =

        ALIGN_POWER_OF_TWO(mi_params->mi_rows, MAX_MIB_SIZE_LOG2);

    TplDepFrame *tpl_frame = &tpl_data->tpl_stats_buffer[frame];

    tpl_frame->is_valid = 0;

    tpl_frame->width = mi_cols >> block_mis_log2;

    tpl_frame->height = mi_rows >> block_mis_log2;

    tpl_frame->stride = tpl_data->tpl_stats_buffer[frame].width;

    tpl_frame->mi_rows = mi_params->mi_rows;

    tpl_frame->mi_cols = mi_params->mi_cols;

  }

  tpl_data->tpl_frame = &tpl_data->tpl_stats_buffer[REF_FRAMES + 1];

  // If lag_in_frames <= 1, TPL module is not invoked. Hence dynamic memory

  // allocations are avoided for buffers in tpl_data.

  if (lag_in_frames <= 1) return;

  AOM_CHECK_MEM_ERROR(&ppi->error, tpl_data->txfm_stats_list,

                      aom_calloc(MAX_LENGTH_TPL_FRAME_STATS,

                                 sizeof(*tpl_data->txfm_stats_list)));

  for (int frame = 0; frame < lag_in_frames; ++frame) {

    AOM_CHECK_MEM_ERROR(

        &ppi->error, tpl_data->tpl_stats_pool[frame],

        aom_calloc(tpl_data->tpl_stats_buffer[frame].width *

                       tpl_data->tpl_stats_buffer[frame].height,

                   sizeof(*tpl_data->tpl_stats_buffer[frame].tpl_stats_ptr)));

    if (aom_alloc_frame_buffer(

            &tpl_data->tpl_rec_pool[frame], width, height,

            seq_params->subsampling_x, seq_params->subsampling_y,

            seq_params->use_highbitdepth, tpl_data->border_in_pixels,

            byte_alignment, false, alloc_y_plane_only))

      aom_internal_error(&ppi->error, AOM_CODEC_MEM_ERROR,

                         "Failed to allocate frame buffer");

  }

  if (aom_alloc_frame_buffer(

          &tpl_data->prev_gop_arf_src, width, height, seq_params->subsampling_x,

          seq_params->subsampling_y, seq_params->use_highbitdepth,

          tpl_data->border_in_pixels, byte_alignment, false,

          alloc_y_plane_only))

    aom_internal_error(&ppi->error, AOM_CODEC_MEM_ERROR,

                       "Failed to allocate prev gop arf buffer");

  tpl_data->prev_gop_arf_disp_order = -1;

}

static inline int32_t tpl_get_satd_cost(BitDepthInfo bd_info, int16_t *src_diff,

                                        int diff_stride, const uint8_t *src,

                                        int src_stride, const uint8_t *dst,

                                        int dst_stride, tran_low_t *coeff,

                                        int bw, int bh, TX_SIZE tx_size) {

  const int pix_num = bw * bh;

  av1_subtract_block(bd_info, bh, bw, src_diff, diff_stride, src, src_stride,

                     dst, dst_stride);

  av1_quick_txfm(/*use_hadamard=*/0, tx_size, bd_info, src_diff, bw, coeff);

  return aom_satd(coeff, pix_num);

}

static int rate_estimator(const tran_low_t *qcoeff, int eob, TX_SIZE tx_size) {

  const SCAN_ORDER *const scan_order = &av1_scan_orders[tx_size][DCT_DCT];

  assert((1 << num_pels_log2_lookup[txsize_to_bsize[tx_size]]) >= eob);

  int rate_cost = 1;

  for (int idx = 0; idx < eob; ++idx) {

    unsigned int abs_level = abs(qcoeff[scan_order->scan[idx]]);

    rate_cost += get_msb(abs_level + 1) + 1 + (abs_level > 0);

  }

  return (rate_cost << AV1_PROB_COST_SHIFT);

}

static inline void txfm_quant_rdcost(

    const MACROBLOCK *x, int16_t *src_diff, int diff_stride, uint8_t *src,

    int src_stride, uint8_t *dst, int dst_stride, tran_low_t *coeff,

    tran_low_t *qcoeff, tran_low_t *dqcoeff, int bw, int bh, TX_SIZE tx_size,

    int do_recon, int *rate_cost, int64_t *recon_error, int64_t *sse) {

  const MACROBLOCKD *xd = &x->e_mbd;

  const BitDepthInfo bd_info = get_bit_depth_info(xd);

  uint16_t eob;

  av1_subtract_block(bd_info, bh, bw, src_diff, diff_stride, src, src_stride,

                     dst, dst_stride);

  av1_quick_txfm(/*use_hadamard=*/0, tx_size, bd_info, src_diff, bw, coeff);

  get_quantize_error(x, 0, coeff, qcoeff, dqcoeff, tx_size, &eob, recon_error,

                     sse);

  *rate_cost = rate_estimator(qcoeff, eob, tx_size);

  if (do_recon)

    av1_inverse_transform_block(xd, dqcoeff, 0, DCT_DCT, tx_size, dst,

                                dst_stride, eob, 0);

}

static uint32_t motion_estimation(AV1_COMP *cpi, MACROBLOCK *x,

                                  uint8_t *cur_frame_buf,

                                  uint8_t *ref_frame_buf, int stride,

                                  int ref_stride, int width, int ref_width,

                                  BLOCK_SIZE bsize, MV center_mv,

                                  int_mv *best_mv) {

  AV1_COMMON *cm = &cpi->common;

  MACROBLOCKD *const xd = &x->e_mbd;

  TPL_SPEED_FEATURES *tpl_sf = &cpi->sf.tpl_sf;

  int step_param;

  uint32_t bestsme = UINT_MAX;

  FULLPEL_MV_STATS best_mv_stats;

  int distortion;

  uint32_t sse;

  int cost_list[5];

  FULLPEL_MV start_mv = get_fullmv_from_mv(&center_mv);

  // Setup frame pointers

  x->plane[0].src.buf = cur_frame_buf;

  x->plane[0].src.stride = stride;

  x->plane[0].src.width = width;

  xd->plane[0].pre[0].buf = ref_frame_buf;

  xd->plane[0].pre[0].stride = ref_stride;

  xd->plane[0].pre[0].width = ref_width;

  step_param = tpl_sf->reduce_first_step_size;

  step_param = AOMMIN(step_param, MAX_MVSEARCH_STEPS - 2);

  const search_site_config *search_site_cfg =

      cpi->mv_search_params.search_site_cfg[SS_CFG_SRC];

  if (search_site_cfg->stride != ref_stride)

    search_site_cfg = cpi->mv_search_params.search_site_cfg[SS_CFG_LOOKAHEAD];

  assert(search_site_cfg->stride == ref_stride);

  FULLPEL_MOTION_SEARCH_PARAMS full_ms_params;

  av1_make_default_fullpel_ms_params(&full_ms_params, cpi, x, bsize, &center_mv,

                                     start_mv, search_site_cfg,

                                     tpl_sf->search_method,

                                     /*fine_search_interval=*/0);

  bestsme = av1_full_pixel_search(start_mv, &full_ms_params, step_param,

                                  cond_cost_list(cpi, cost_list),

                                  &best_mv->as_fullmv, &best_mv_stats, NULL);

  // When sub-pel motion search is skipped, populate sub-pel precision MV and

  // return.

  if (tpl_sf->subpel_force_stop == FULL_PEL) {

    best_mv->as_mv = get_mv_from_fullmv(&best_mv->as_fullmv);

    return bestsme;

  }

  SUBPEL_MOTION_SEARCH_PARAMS ms_params;

  av1_make_default_subpel_ms_params(&ms_params, cpi, x, bsize, &center_mv,

                                    cost_list);

  ms_params.forced_stop = tpl_sf->subpel_force_stop;

  ms_params.var_params.subpel_search_type = USE_2_TAPS;

  ms_params.mv_cost_params.mv_cost_type = MV_COST_NONE;

  best_mv_stats.err_cost = 0;

  MV subpel_start_mv = get_mv_from_fullmv(&best_mv->as_fullmv);

  assert(av1_is_subpelmv_in_range(&ms_params.mv_limits, subpel_start_mv));

  bestsme = cpi->mv_search_params.find_fractional_mv_step(

      xd, cm, &ms_params, subpel_start_mv, &best_mv_stats, &best_mv->as_mv,

      &distortion, &sse, NULL);

  return bestsme;

}

typedef struct {

  int_mv mv;

  int sad;

} center_mv_t;

static int compare_sad(const void *a, const void *b) {

  const int diff = ((center_mv_t *)a)->sad - ((center_mv_t *)b)->sad;

  if (diff < 0)

    return -1;

  else if (diff > 0)

    return 1;

  return 0;

}

static int is_alike_mv(int_mv candidate_mv, center_mv_t *center_mvs,

                       int center_mvs_count, int skip_alike_starting_mv) {

  // MV difference threshold is in 1/8 precision.

  const int mv_diff_thr[3] = { 1, (8 << 3), (16 << 3) };

  int thr = mv_diff_thr[skip_alike_starting_mv];

  int i;

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

    if (abs(center_mvs[i].mv.as_mv.col - candidate_mv.as_mv.col) < thr &&

        abs(center_mvs[i].mv.as_mv.row - candidate_mv.as_mv.row) < thr)

      return 1;

  }

  return 0;

}

static void get_rate_distortion(

    int *rate_cost, int64_t *recon_error, int64_t *pred_error,

    int16_t *src_diff, tran_low_t *coeff, tran_low_t *qcoeff,

    tran_low_t *dqcoeff, AV1_COMMON *cm, MACROBLOCK *x,

    const YV12_BUFFER_CONFIG *ref_frame_ptr[2], uint8_t *rec_buffer_pool[3],

    const int rec_stride_pool[3], TX_SIZE tx_size, PREDICTION_MODE best_mode,

    int mi_row, int mi_col, int use_y_only_rate_distortion, int do_recon,

    TplTxfmStats *tpl_txfm_stats) {

  const SequenceHeader *seq_params = cm->seq_params;

  *rate_cost = 0;

  *recon_error = 1;

  *pred_error = 1;

  (void)tpl_txfm_stats;

  MACROBLOCKD *xd = &x->e_mbd;

  int is_compound = (best_mode == NEW_NEWMV);

  int num_planes = use_y_only_rate_distortion ? 1 : MAX_MB_PLANE;

  uint8_t *src_buffer_pool[MAX_MB_PLANE] = {

    xd->cur_buf->y_buffer,

    xd->cur_buf->u_buffer,

    xd->cur_buf->v_buffer,

  };

  const int src_stride_pool[MAX_MB_PLANE] = {

    xd->cur_buf->y_stride,

    xd->cur_buf->uv_stride,

    xd->cur_buf->uv_stride,

  };

  const int_interpfilters kernel =

      av1_broadcast_interp_filter(EIGHTTAP_REGULAR);

  for (int plane = 0; plane < num_planes; ++plane) {

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

    BLOCK_SIZE bsize_plane =

        av1_ss_size_lookup[txsize_to_bsize[tx_size]][pd->subsampling_x]

                          [pd->subsampling_y];

    int dst_buffer_stride = rec_stride_pool[plane];

    int dst_mb_offset =

        ((mi_row * MI_SIZE * dst_buffer_stride) >> pd->subsampling_y) +

        ((mi_col * MI_SIZE) >> pd->subsampling_x);

    uint8_t *dst_buffer = rec_buffer_pool[plane] + dst_mb_offset;

    for (int ref = 0; ref < 1 + is_compound; ++ref) {

      if (!is_inter_mode(best_mode)) {

        av1_predict_intra_block(

            xd, seq_params->sb_size, seq_params->enable_intra_edge_filter,

            block_size_wide[bsize_plane], block_size_high[bsize_plane],

            max_txsize_rect_lookup[bsize_plane], best_mode, 0, 0,

            FILTER_INTRA_MODES, dst_buffer, dst_buffer_stride, dst_buffer,

            dst_buffer_stride, 0, 0, plane);

      } else {

        int_mv best_mv = xd->mi[0]->mv[ref];

        uint8_t *ref_buffer_pool[MAX_MB_PLANE] = {

          ref_frame_ptr[ref]->y_buffer,

          ref_frame_ptr[ref]->u_buffer,

          ref_frame_ptr[ref]->v_buffer,

        };

        InterPredParams inter_pred_params;

        struct buf_2d ref_buf = {

          NULL, ref_buffer_pool[plane],

          plane ? ref_frame_ptr[ref]->uv_width : ref_frame_ptr[ref]->y_width,

          plane ? ref_frame_ptr[ref]->uv_height : ref_frame_ptr[ref]->y_height,

          plane ? ref_frame_ptr[ref]->uv_stride : ref_frame_ptr[ref]->y_stride

        };

        av1_init_inter_params(&inter_pred_params, block_size_wide[bsize_plane],

                              block_size_high[bsize_plane],

                              (mi_row * MI_SIZE) >> pd->subsampling_y,

                              (mi_col * MI_SIZE) >> pd->subsampling_x,

                              pd->subsampling_x, pd->subsampling_y, xd->bd,

                              is_cur_buf_hbd(xd), 0,

                              xd->block_ref_scale_factors[0], &ref_buf, kernel);

        if (is_compound) av1_init_comp_mode(&inter_pred_params);

        inter_pred_params.conv_params = get_conv_params_no_round(

            ref, plane, xd->tmp_conv_dst, MAX_SB_SIZE, is_compound, xd->bd);

        av1_enc_build_one_inter_predictor(dst_buffer, dst_buffer_stride,

                                          &best_mv.as_mv, &inter_pred_params);

      }

    }

    int src_stride = src_stride_pool[plane];

    int src_mb_offset = ((mi_row * MI_SIZE * src_stride) >> pd->subsampling_y) +

                        ((mi_col * MI_SIZE) >> pd->subsampling_x);

    int this_rate = 1;

    int64_t this_recon_error = 1;

    int64_t sse;

    txfm_quant_rdcost(

        x, src_diff, block_size_wide[bsize_plane],

        src_buffer_pool[plane] + src_mb_offset, src_stride, dst_buffer,

        dst_buffer_stride, coeff, qcoeff, dqcoeff, block_size_wide[bsize_plane],

        block_size_high[bsize_plane], max_txsize_rect_lookup[bsize_plane],

        do_recon, &this_rate, &this_recon_error, &sse);

#if CONFIG_BITRATE_ACCURACY

    if (plane == 0 && tpl_txfm_stats) {

      // We only collect Y plane's transform coefficient

      av1_record_tpl_txfm_block(tpl_txfm_stats, coeff);

    }

#endif  // CONFIG_BITRATE_ACCURACY

    *recon_error += this_recon_error;

    *pred_error += sse;

    *rate_cost += this_rate;

  }

}

static inline int32_t get_inter_cost(const AV1_COMP *cpi, MACROBLOCKD *xd,

                                     const uint8_t *src_mb_buffer,

                                     int src_stride,

                                     TplBuffers *tpl_tmp_buffers,

                                     BLOCK_SIZE bsize, TX_SIZE tx_size,

                                     int mi_row, int mi_col, int rf_idx,

                                     MV *rfidx_mv, int use_pred_sad) {

  const BitDepthInfo bd_info = get_bit_depth_info(xd);

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

  const YV12_BUFFER_CONFIG *const ref_frame_ptr =

      tpl_data->src_ref_frame[rf_idx];

  int16_t *src_diff = tpl_tmp_buffers->src_diff;

  tran_low_t *coeff = tpl_tmp_buffers->coeff;

  const int bw = 4 << mi_size_wide_log2[bsize];

  const int bh = 4 << mi_size_high_log2[bsize];

  int32_t inter_cost;

  if (cpi->sf.tpl_sf.subpel_force_stop != FULL_PEL) {

    const int_interpfilters kernel =

        av1_broadcast_interp_filter(EIGHTTAP_REGULAR);

    uint8_t *predictor8 = tpl_tmp_buffers->predictor8;

    uint8_t *predictor =

        is_cur_buf_hbd(xd) ? CONVERT_TO_BYTEPTR(predictor8) : predictor8;

    struct buf_2d ref_buf = { NULL, ref_frame_ptr->y_buffer,

                              ref_frame_ptr->y_width, ref_frame_ptr->y_height,

                              ref_frame_ptr->y_stride };

    InterPredParams inter_pred_params;

    av1_init_inter_params(&inter_pred_params, bw, bh, mi_row * MI_SIZE,

                          mi_col * MI_SIZE, 0, 0, xd->bd, is_cur_buf_hbd(xd), 0,

                          &tpl_data->sf, &ref_buf, kernel);

    inter_pred_params.conv_params = get_conv_params(0, 0, xd->bd);

    av1_enc_build_one_inter_predictor(predictor, bw, rfidx_mv,

                                      &inter_pred_params);

    if (use_pred_sad) {

      inter_cost = (int)cpi->ppi->fn_ptr[bsize].sdf(src_mb_buffer, src_stride,

                                                    predictor, bw);

    } else {

      inter_cost =

          tpl_get_satd_cost(bd_info, src_diff, bw, src_mb_buffer, src_stride,

                            predictor, bw, coeff, bw, bh, tx_size);

    }

  } else {

    int ref_mb_offset =

        mi_row * MI_SIZE * ref_frame_ptr->y_stride + mi_col * MI_SIZE;

    uint8_t *ref_mb = ref_frame_ptr->y_buffer + ref_mb_offset;

    int ref_stride = ref_frame_ptr->y_stride;

    const FULLPEL_MV fullmv = get_fullmv_from_mv(rfidx_mv);

    // Since sub-pel motion search is not performed, use the prediction pixels

    // directly from the reference block ref_mb

    if (use_pred_sad) {

      inter_cost = (int)cpi->ppi->fn_ptr[bsize].sdf(

          src_mb_buffer, src_stride,

          &ref_mb[fullmv.row * ref_stride + fullmv.col], ref_stride);

    } else {

      inter_cost =

          tpl_get_satd_cost(bd_info, src_diff, bw, src_mb_buffer, src_stride,

                            &ref_mb[fullmv.row * ref_stride + fullmv.col],

                            ref_stride, coeff, bw, bh, tx_size);

    }

  }

  return inter_cost;

}

static inline void mode_estimation(AV1_COMP *cpi, TplTxfmStats *tpl_txfm_stats,

                                   TplBuffers *tpl_tmp_buffers, MACROBLOCK *x,

                                   int mi_row, int mi_col, BLOCK_SIZE bsize,

                                   TX_SIZE tx_size, TplDepStats *tpl_stats) {

  AV1_COMMON *cm = &cpi->common;

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

  TPL_SPEED_FEATURES *tpl_sf = &cpi->sf.tpl_sf;

  (void)gf_group;

  MACROBLOCKD *xd = &x->e_mbd;

  const BitDepthInfo bd_info = get_bit_depth_info(xd);

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

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

  const uint8_t block_mis_log2 = tpl_data->tpl_stats_block_mis_log2;

  const int bw = 4 << mi_size_wide_log2[bsize];

  const int bh = 4 << mi_size_high_log2[bsize];

  int32_t best_intra_cost = INT32_MAX;

  int32_t intra_cost;

  PREDICTION_MODE best_mode = DC_PRED;

  const int mb_y_offset =

      mi_row * MI_SIZE * xd->cur_buf->y_stride + mi_col * MI_SIZE;

  uint8_t *src_mb_buffer = xd->cur_buf->y_buffer + mb_y_offset;

  const int src_stride = xd->cur_buf->y_stride;

  const int src_width = xd->cur_buf->y_width;

  int dst_mb_offset =

      mi_row * MI_SIZE * tpl_frame->rec_picture->y_stride + mi_col * MI_SIZE;

  uint8_t *dst_buffer = tpl_frame->rec_picture->y_buffer + dst_mb_offset;

  int dst_buffer_stride = tpl_frame->rec_picture->y_stride;

  int use_y_only_rate_distortion = tpl_sf->use_y_only_rate_distortion;

  uint8_t *rec_buffer_pool[3] = {

    tpl_frame->rec_picture->y_buffer,

    tpl_frame->rec_picture->u_buffer,

    tpl_frame->rec_picture->v_buffer,

  };

  const int rec_stride_pool[3] = {

    tpl_frame->rec_picture->y_stride,

    tpl_frame->rec_picture->uv_stride,

    tpl_frame->rec_picture->uv_stride,

  };

  for (int plane = 1; plane < MAX_MB_PLANE; ++plane) {

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

    pd->subsampling_x = xd->cur_buf->subsampling_x;

    pd->subsampling_y = xd->cur_buf->subsampling_y;

  }

  uint8_t *predictor8 = tpl_tmp_buffers->predictor8;

  int16_t *src_diff = tpl_tmp_buffers->src_diff;

  tran_low_t *coeff = tpl_tmp_buffers->coeff;

  tran_low_t *qcoeff = tpl_tmp_buffers->qcoeff;

  tran_low_t *dqcoeff = tpl_tmp_buffers->dqcoeff;

  uint8_t *predictor =

      is_cur_buf_hbd(xd) ? CONVERT_TO_BYTEPTR(predictor8) : predictor8;

  int64_t recon_error = 1;

  int64_t pred_error = 1;

  memset(tpl_stats, 0, sizeof(*tpl_stats));

  tpl_stats->ref_frame_index[0] = -1;

  tpl_stats->ref_frame_index[1] = -1;

  const int mi_width = mi_size_wide[bsize];

  const int mi_height = mi_size_high[bsize];

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

                        mi_row, mi_col);

  set_mi_row_col(xd, &xd->tile, mi_row, mi_height, mi_col, mi_width,

                 cm->mi_params.mi_rows, cm->mi_params.mi_cols);

  set_plane_n4(xd, mi_size_wide[bsize], mi_size_high[bsize],

               av1_num_planes(cm));

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

  xd->mi[0]->motion_mode = SIMPLE_TRANSLATION;

  // Intra prediction search

  xd->mi[0]->ref_frame[0] = INTRA_FRAME;

  // Pre-load the bottom left line.

  if (xd->left_available &&

      mi_row + tx_size_high_unit[tx_size] < xd->tile.mi_row_end) {

    if (is_cur_buf_hbd(xd)) {

      uint16_t *dst = CONVERT_TO_SHORTPTR(dst_buffer);

      for (int i = 0; i < bw; ++i)

        dst[(bw + i) * dst_buffer_stride - 1] =

            dst[(bw - 1) * dst_buffer_stride - 1];

    } else {

      for (int i = 0; i < bw; ++i)

        dst_buffer[(bw + i) * dst_buffer_stride - 1] =

            dst_buffer[(bw - 1) * dst_buffer_stride - 1];

    }

  }

  // if cpi->sf.tpl_sf.prune_intra_modes is on, then search only DC_PRED,

  // H_PRED, and V_PRED

  const PREDICTION_MODE last_intra_mode =

      tpl_sf->prune_intra_modes ? D45_PRED : INTRA_MODE_END;

  const SequenceHeader *seq_params = cm->seq_params;

  for (PREDICTION_MODE mode = INTRA_MODE_START; mode < last_intra_mode;

       ++mode) {

    av1_predict_intra_block(xd, seq_params->sb_size,

                            seq_params->enable_intra_edge_filter,

                            block_size_wide[bsize], block_size_high[bsize],

                            tx_size, mode, 0, 0, FILTER_INTRA_MODES, dst_buffer,

                            dst_buffer_stride, predictor, bw, 0, 0, 0);

    if (tpl_frame->use_pred_sad) {

      intra_cost = (int32_t)cpi->ppi->fn_ptr[bsize].sdf(

          src_mb_buffer, src_stride, predictor, bw);

    } else {

      intra_cost =

          tpl_get_satd_cost(bd_info, src_diff, bw, src_mb_buffer, src_stride,

                            predictor, bw, coeff, bw, bh, tx_size);

    }

    if (intra_cost < best_intra_cost) {

      best_intra_cost = intra_cost;

      best_mode = mode;

    }

  }

  // Calculate SATD of the best intra mode if SAD was used for mode decision

  // as best_intra_cost is used in ML model to skip intra mode evaluation.

  if (tpl_frame->use_pred_sad) {

    av1_predict_intra_block(

        xd, seq_params->sb_size, seq_params->enable_intra_edge_filter,

        block_size_wide[bsize], block_size_high[bsize], tx_size, best_mode, 0,

        0, FILTER_INTRA_MODES, dst_buffer, dst_buffer_stride, predictor, bw, 0,

        0, 0);

    best_intra_cost =

        tpl_get_satd_cost(bd_info, src_diff, bw, src_mb_buffer, src_stride,

                          predictor, bw, coeff, bw, bh, tx_size);

  }

  int rate_cost = 1;

  if (cpi->use_ducky_encode) {

    get_rate_distortion(&rate_cost, &recon_error, &pred_error, src_diff, coeff,

                        qcoeff, dqcoeff, cm, x, NULL, rec_buffer_pool,

                        rec_stride_pool, tx_size, best_mode, mi_row, mi_col,

                        use_y_only_rate_distortion, 1 /*do_recon*/, NULL);

    tpl_stats->intra_dist = recon_error << TPL_DEP_COST_SCALE_LOG2;

    tpl_stats->intra_sse = pred_error << TPL_DEP_COST_SCALE_LOG2;

    tpl_stats->intra_rate = rate_cost;

  }

#if CONFIG_THREE_PASS

  const int frame_offset = tpl_data->frame_idx - cpi->gf_frame_index;

  if (cpi->third_pass_ctx &&

      frame_offset < cpi->third_pass_ctx->frame_info_count &&

      tpl_data->frame_idx < gf_group->size) {

    double ratio_h, ratio_w;

    av1_get_third_pass_ratio(cpi->third_pass_ctx, frame_offset, cm->height,

                             cm->width, &ratio_h, &ratio_w);

    THIRD_PASS_MI_INFO *this_mi = av1_get_third_pass_mi(

        cpi->third_pass_ctx, frame_offset, mi_row, mi_col, ratio_h, ratio_w);

    PREDICTION_MODE third_pass_mode = this_mi->pred_mode;

    if (third_pass_mode >= last_intra_mode &&

        third_pass_mode < INTRA_MODE_END) {

      av1_predict_intra_block(

          xd, seq_params->sb_size, seq_params->enable_intra_edge_filter,

          block_size_wide[bsize], block_size_high[bsize], tx_size,

          third_pass_mode, 0, 0, FILTER_INTRA_MODES, dst_buffer,

          dst_buffer_stride, predictor, bw, 0, 0, 0);

      intra_cost =

          tpl_get_satd_cost(bd_info, src_diff, bw, src_mb_buffer, src_stride,

                            predictor, bw, coeff, bw, bh, tx_size);

      if (intra_cost < best_intra_cost) {

        best_intra_cost = intra_cost;

        best_mode = third_pass_mode;

      }

    }

  }

#endif  // CONFIG_THREE_PASS

  // Motion compensated prediction

  xd->mi[0]->ref_frame[0] = INTRA_FRAME;

  xd->mi[0]->ref_frame[1] = NONE_FRAME;

  xd->mi[0]->compound_idx = 1;

  int best_rf_idx = -1;

  int_mv best_mv[2];

  int32_t inter_cost;

  int32_t best_inter_cost = INT32_MAX;

  int rf_idx;

  int_mv single_mv[INTER_REFS_PER_FRAME];

  best_mv[0].as_int = INVALID_MV;

  best_mv[1].as_int = INVALID_MV;

  for (rf_idx = 0; rf_idx < INTER_REFS_PER_FRAME; ++rf_idx) {

    single_mv[rf_idx].as_int = INVALID_MV;

    if (tpl_data->ref_frame[rf_idx] == NULL ||

        tpl_data->src_ref_frame[rf_idx] == NULL) {

      tpl_stats->mv[rf_idx].as_int = INVALID_MV;

      continue;

    }

    const YV12_BUFFER_CONFIG *ref_frame_ptr = tpl_data->src_ref_frame[rf_idx];

    const int ref_mb_offset =

        mi_row * MI_SIZE * ref_frame_ptr->y_stride + mi_col * MI_SIZE;

    uint8_t *ref_mb = ref_frame_ptr->y_buffer + ref_mb_offset;

    const int ref_stride = ref_frame_ptr->y_stride;

    const int ref_width = ref_frame_ptr->y_width;

    int_mv best_rfidx_mv = { 0 };

    uint32_t bestsme = UINT32_MAX;

    center_mv_t center_mvs[4] = { { { 0 }, INT_MAX },

                                  { { 0 }, INT_MAX },

                                  { { 0 }, INT_MAX },

                                  { { 0 }, INT_MAX } };

    int refmv_count = 1;

    int idx;

    if (xd->up_available) {

      TplDepStats *ref_tpl_stats = &tpl_frame->tpl_stats_ptr[av1_tpl_ptr_pos(

          mi_row - mi_height, mi_col, tpl_frame->stride, block_mis_log2)];

      if (!is_alike_mv(ref_tpl_stats->mv[rf_idx], center_mvs, refmv_count,

                       tpl_sf->skip_alike_starting_mv)) {

        center_mvs[refmv_count].mv.as_int = ref_tpl_stats->mv[rf_idx].as_int;

        ++refmv_count;

      }

    }

    if (xd->left_available) {

      TplDepStats *ref_tpl_stats = &tpl_frame->tpl_stats_ptr[av1_tpl_ptr_pos(

          mi_row, mi_col - mi_width, tpl_frame->stride, block_mis_log2)];

      if (!is_alike_mv(ref_tpl_stats->mv[rf_idx], center_mvs, refmv_count,

                       tpl_sf->skip_alike_starting_mv)) {

        center_mvs[refmv_count].mv.as_int = ref_tpl_stats->mv[rf_idx].as_int;

        ++refmv_count;

      }

    }

    if (xd->up_available && mi_col + mi_width < xd->tile.mi_col_end) {

      TplDepStats *ref_tpl_stats = &tpl_frame->tpl_stats_ptr[av1_tpl_ptr_pos(

          mi_row - mi_height, mi_col + mi_width, tpl_frame->stride,

          block_mis_log2)];

      if (!is_alike_mv(ref_tpl_stats->mv[rf_idx], center_mvs, refmv_count,

                       tpl_sf->skip_alike_starting_mv)) {

        center_mvs[refmv_count].mv.as_int = ref_tpl_stats->mv[rf_idx].as_int;

        ++refmv_count;

      }

    }

#if CONFIG_THREE_PASS

    if (cpi->third_pass_ctx &&

        frame_offset < cpi->third_pass_ctx->frame_info_count &&

        tpl_data->frame_idx < gf_group->size) {

      double ratio_h, ratio_w;

      av1_get_third_pass_ratio(cpi->third_pass_ctx, frame_offset, cm->height,

                               cm->width, &ratio_h, &ratio_w);

      THIRD_PASS_MI_INFO *this_mi = av1_get_third_pass_mi(

          cpi->third_pass_ctx, frame_offset, mi_row, mi_col, ratio_h, ratio_w);

      int_mv tp_mv = av1_get_third_pass_adjusted_mv(this_mi, ratio_h, ratio_w,

                                                    rf_idx + LAST_FRAME);

      if (tp_mv.as_int != INVALID_MV &&

          !is_alike_mv(tp_mv, center_mvs + 1, refmv_count - 1,

                       tpl_sf->skip_alike_starting_mv)) {

        center_mvs[0].mv = tp_mv;

      }

    }

#endif  // CONFIG_THREE_PASS

    // Prune starting mvs

    if (tpl_sf->prune_starting_mv && refmv_count > 1) {

      // Get each center mv's sad.

      for (idx = 0; idx < refmv_count; ++idx) {

        FULLPEL_MV mv = get_fullmv_from_mv(&center_mvs[idx].mv.as_mv);

        clamp_fullmv(&mv, &x->mv_limits);

        center_mvs[idx].sad = (int)cpi->ppi->fn_ptr[bsize].sdf(

            src_mb_buffer, src_stride, &ref_mb[mv.row * ref_stride + mv.col],

            ref_stride);

      }

      // Rank center_mv using sad.

      qsort(center_mvs, refmv_count, sizeof(center_mvs[0]), compare_sad);

      refmv_count = AOMMIN(4 - tpl_sf->prune_starting_mv, refmv_count);

      // Further reduce number of refmv based on sad difference.

      if (refmv_count > 1) {

        int last_sad = center_mvs[refmv_count - 1].sad;

        int second_to_last_sad = center_mvs[refmv_count - 2].sad;

        if ((last_sad - second_to_last_sad) * 5 > second_to_last_sad)

          refmv_count--;

      }

    }

    for (idx = 0; idx < refmv_count; ++idx) {

      int_mv this_mv;

      uint32_t thissme = motion_estimation(

          cpi, x, src_mb_buffer, ref_mb, src_stride, ref_stride, src_width,

          ref_width, bsize, center_mvs[idx].mv.as_mv, &this_mv);

      if (thissme < bestsme) {

        bestsme = thissme;

        best_rfidx_mv = this_mv;

      }

    }

    tpl_stats->mv[rf_idx].as_int = best_rfidx_mv.as_int;

    single_mv[rf_idx] = best_rfidx_mv;

    inter_cost = get_inter_cost(

        cpi, xd, src_mb_buffer, src_stride, tpl_tmp_buffers, bsize, tx_size,

        mi_row, mi_col, rf_idx, &best_rfidx_mv.as_mv, tpl_frame->use_pred_sad);

    // Store inter cost for each ref frame. This is used to prune inter modes.

    tpl_stats->pred_error[rf_idx] = AOMMAX(1, inter_cost);

    if (inter_cost < best_inter_cost) {

      best_rf_idx = rf_idx;

      best_inter_cost = inter_cost;

      best_mv[0].as_int = best_rfidx_mv.as_int;

    }

  }

  // Calculate SATD of the best inter mode if SAD was used for mode decision

  // as best_inter_cost is used in ML model to skip intra mode evaluation.

  if (best_inter_cost < INT32_MAX && tpl_frame->use_pred_sad) {

    assert(best_rf_idx != -1);

    best_inter_cost = get_inter_cost(

        cpi, xd, src_mb_buffer, src_stride, tpl_tmp_buffers, bsize, tx_size,

        mi_row, mi_col, best_rf_idx, &best_mv[0].as_mv, 0 /* use_pred_sad */);

  }

  if (best_rf_idx != -1 && best_inter_cost < best_intra_cost) {

    best_mode = NEWMV;

    xd->mi[0]->ref_frame[0] = best_rf_idx + LAST_FRAME;

    xd->mi[0]->mv[0].as_int = best_mv[0].as_int;

  }

  // Start compound predition search.

  int comp_ref_frames[3][2] = {

    { 0, 4 },

    { 0, 6 },

    { 3, 6 },

  };

  int start_rf = 0;

  int end_rf = 3;

  if (!tpl_sf->allow_compound_pred) end_rf = 0;

#if CONFIG_THREE_PASS

  if (cpi->third_pass_ctx &&

      frame_offset < cpi->third_pass_ctx->frame_info_count &&

      tpl_data->frame_idx < gf_group->size) {

    double ratio_h, ratio_w;

    av1_get_third_pass_ratio(cpi->third_pass_ctx, frame_offset, cm->height,

                             cm->width, &ratio_h, &ratio_w);

    THIRD_PASS_MI_INFO *this_mi = av1_get_third_pass_mi(

        cpi->third_pass_ctx, frame_offset, mi_row, mi_col, ratio_h, ratio_w);

    if (this_mi->ref_frame[0] >= LAST_FRAME &&

        this_mi->ref_frame[1] >= LAST_FRAME) {

      int found = 0;

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

        if (comp_ref_frames[i][0] + LAST_FRAME == this_mi->ref_frame[0] &&

            comp_ref_frames[i][1] + LAST_FRAME == this_mi->ref_frame[1]) {

          found = 1;

          break;

        }

      }

      if (!found || !tpl_sf->allow_compound_pred) {

        comp_ref_frames[2][0] = this_mi->ref_frame[0] - LAST_FRAME;