/*

 *  Copyright (c) 2023 The WebM project authors. All Rights Reserved.

 *

 *  Use of this source code is governed by a BSD-style license

 *  that can be found in the LICENSE file in the root of the source

 *  tree. An additional intellectual property rights grant can be found

 *  in the file PATENTS.  All contributing project authors may

 *  be found in the AUTHORS file in the root of the source tree.

 */

#include <math.h>

#include "./vpx_dsp_rtcd.h"

#if CONFIG_NON_GREEDY_MV

#include "vp9/common/vp9_mvref_common.h"

#endif

#include "vp9/common/vp9_reconinter.h"

#include "vp9/common/vp9_reconintra.h"

#include "vp9/common/vp9_scan.h"

#include "vp9/encoder/vp9_encoder.h"

#include "vp9/encoder/vp9_ext_ratectrl.h"

#include "vp9/encoder/vp9_firstpass.h"

#include "vp9/encoder/vp9_ratectrl.h"

#include "vp9/encoder/vp9_tpl_model.h"

#include "vpx/internal/vpx_codec_internal.h"

#include "vpx/vpx_codec.h"

#include "vpx/vpx_ext_ratectrl.h"

static int init_gop_frames_rc(VP9_COMP *cpi, GF_PICTURE *gf_picture,

                              const GF_GROUP *gf_group, int *tpl_group_frames) {

  VP9_COMMON *cm = &cpi->common;

  int frame_idx = 0;

  int i;

  int extend_frame_count = 0;

  int pframe_qindex = cpi->tpl_stats[2].base_qindex;

  int frame_gop_offset = 0;

  int added_overlay = 0;

  RefCntBuffer *frame_bufs = cm->buffer_pool->frame_bufs;

  int8_t recon_frame_index[REFS_PER_FRAME + MAX_ARF_LAYERS];

  memset(recon_frame_index, -1, sizeof(recon_frame_index));

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

    if (frame_bufs[i].ref_count == 0) {

      alloc_frame_mvs(cm, i);

      if (vpx_realloc_frame_buffer(&frame_bufs[i].buf, cm->width, cm->height,

                                   cm->subsampling_x, cm->subsampling_y,

#if CONFIG_VP9_HIGHBITDEPTH

                                   cm->use_highbitdepth,

#endif

                                   VP9_ENC_BORDER_IN_PIXELS, cm->byte_alignment,

                                   NULL, NULL, NULL))

        vpx_internal_error(&cm->error, VPX_CODEC_MEM_ERROR,

                           "Failed to allocate frame buffer");

      recon_frame_index[frame_idx] = i;

      ++frame_idx;

      if (frame_idx >= REFS_PER_FRAME + cpi->oxcf.enable_auto_arf) break;

    }

  }

  for (i = 0; i < REFS_PER_FRAME + 1; ++i) {

    assert(recon_frame_index[i] >= 0);

    cpi->tpl_recon_frames[i] = &frame_bufs[recon_frame_index[i]].buf;

  }

  *tpl_group_frames = 0;

  int ref_table[3];

  if (gf_group->index == 1 && gf_group->update_type[1] == ARF_UPDATE) {

    if (gf_group->update_type[0] == KF_UPDATE) {

      // This is the only frame in ref buffer. We need it to be on

      // gf_picture[0].

      for (i = 0; i < 3; ++i) ref_table[i] = -REFS_PER_FRAME;

      gf_picture[0].frame =

          &cm->buffer_pool->frame_bufs[gf_group->update_ref_idx[0]].buf;

      ref_table[gf_group->update_ref_idx[0]] = 0;

      for (i = 0; i < 3; ++i) gf_picture[0].ref_frame[i] = -REFS_PER_FRAME;

      gf_picture[0].update_type = gf_group->update_type[0];

    } else {

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

        if (cm->ref_frame_map[i] != -1) {

          gf_picture[-i].frame =

              &cm->buffer_pool->frame_bufs[cm->ref_frame_map[i]].buf;

          ref_table[i] = -i;

        } else {

          ref_table[i] = -REFS_PER_FRAME;

        }

      }

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

        gf_picture[0].ref_frame[i] = ref_table[i];

      }

    }

    ++*tpl_group_frames;

    // Initialize base layer ARF frame

    gf_picture[1].frame = cpi->Source;

    for (i = 0; i < 3; ++i) gf_picture[1].ref_frame[i] = ref_table[i];

    gf_picture[1].update_type = gf_group->update_type[1];

    ref_table[gf_group->update_ref_idx[1]] = 1;

    ++*tpl_group_frames;

  } else {

    assert(gf_group->index == 0);

    if (gf_group->update_type[0] == KF_UPDATE) {

      // This is the only frame in ref buffer. We need it to be on

      // gf_picture[0].

      gf_picture[0].frame = cpi->Source;

      for (i = 0; i < 3; ++i) gf_picture[0].ref_frame[i] = -REFS_PER_FRAME;

      gf_picture[0].update_type = gf_group->update_type[0];

      for (i = 0; i < 3; ++i) ref_table[i] = -REFS_PER_FRAME;

      ref_table[gf_group->update_ref_idx[0]] = 0;

    } else {

      // Initialize ref table

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

        if (cm->ref_frame_map[i] != -1) {

          gf_picture[-i].frame =

              &cm->buffer_pool->frame_bufs[cm->ref_frame_map[i]].buf;

          ref_table[i] = -i;

        } else {

          ref_table[i] = -REFS_PER_FRAME;

        }

      }

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

        gf_picture[0].ref_frame[i] = ref_table[i];

      }

      gf_picture[0].update_type = gf_group->update_type[0];

      if (gf_group->update_type[0] != OVERLAY_UPDATE &&

          gf_group->update_ref_idx[0] != -1) {

        ref_table[gf_group->update_ref_idx[0]] = 0;

      }

    }

    ++*tpl_group_frames;

  }

  int has_arf =

      gf_group->gf_group_size > 1 && gf_group->update_type[1] == ARF_UPDATE &&

      gf_group->update_type[gf_group->gf_group_size] == OVERLAY_UPDATE;

  // Initialize P frames

  for (frame_idx = *tpl_group_frames; frame_idx < MAX_ARF_GOP_SIZE;

       ++frame_idx) {

    if (frame_idx >= gf_group->gf_group_size && !has_arf) break;

    struct lookahead_entry *buf;

    frame_gop_offset = gf_group->frame_gop_index[frame_idx];

    buf = vp9_lookahead_peek(cpi->lookahead, frame_gop_offset - 1);

    if (buf == NULL) break;

    gf_picture[frame_idx].frame = &buf->img;

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

      gf_picture[frame_idx].ref_frame[i] = ref_table[i];

    }

    if (gf_group->update_type[frame_idx] != OVERLAY_UPDATE &&

        gf_group->update_ref_idx[frame_idx] != -1) {

      ref_table[gf_group->update_ref_idx[frame_idx]] = frame_idx;

    }

    gf_picture[frame_idx].update_type = gf_group->update_type[frame_idx];

    ++*tpl_group_frames;

    // The length of group of pictures is baseline_gf_interval, plus the

    // beginning golden frame from last GOP, plus the last overlay frame in

    // the same GOP.

    if (frame_idx == gf_group->gf_group_size) {

      added_overlay = 1;

      ++frame_idx;

      ++frame_gop_offset;

      break;

    }

    if (frame_idx == gf_group->gf_group_size - 1 &&

        gf_group->update_type[gf_group->gf_group_size] != OVERLAY_UPDATE) {

      ++frame_idx;

      ++frame_gop_offset;

      break;

    }

  }

  int lst_index = frame_idx - 1;

  // Extend two frames outside the current gf group.

  for (; has_arf && frame_idx < MAX_LAG_BUFFERS && extend_frame_count < 2;

       ++frame_idx) {

    struct lookahead_entry *buf =

        vp9_lookahead_peek(cpi->lookahead, frame_gop_offset - 1);

    if (buf == NULL) break;

    cpi->tpl_stats[frame_idx].base_qindex = pframe_qindex;

    gf_picture[frame_idx].frame = &buf->img;

    gf_picture[frame_idx].ref_frame[0] = gf_picture[lst_index].ref_frame[0];

    gf_picture[frame_idx].ref_frame[1] = gf_picture[lst_index].ref_frame[1];

    gf_picture[frame_idx].ref_frame[2] = gf_picture[lst_index].ref_frame[2];

    if (gf_picture[frame_idx].ref_frame[0] >

            gf_picture[frame_idx].ref_frame[1] &&

        gf_picture[frame_idx].ref_frame[0] >

            gf_picture[frame_idx].ref_frame[2]) {

      gf_picture[frame_idx].ref_frame[0] = lst_index;

    } else if (gf_picture[frame_idx].ref_frame[1] >

                   gf_picture[frame_idx].ref_frame[0] &&

               gf_picture[frame_idx].ref_frame[1] >

                   gf_picture[frame_idx].ref_frame[2]) {

      gf_picture[frame_idx].ref_frame[1] = lst_index;

    } else {

      gf_picture[frame_idx].ref_frame[2] = lst_index;

    }

    gf_picture[frame_idx].update_type = LF_UPDATE;

    lst_index = frame_idx;

    ++*tpl_group_frames;

    ++extend_frame_count;

    ++frame_gop_offset;

  }

  return extend_frame_count + added_overlay;

}

static int init_gop_frames(VP9_COMP *cpi, GF_PICTURE *gf_picture,

                           const GF_GROUP *gf_group, int *tpl_group_frames) {

  if (cpi->ext_ratectrl.ready &&

      (cpi->ext_ratectrl.funcs.rc_type & VPX_RC_GOP) != 0) {

    return init_gop_frames_rc(cpi, gf_picture, gf_group, tpl_group_frames);

  }

  VP9_COMMON *cm = &cpi->common;

  int frame_idx = 0;

  int i;

  int gld_index = -1;

  int alt_index = -2;

  int lst_index = -1;

  int arf_index_stack[MAX_ARF_LAYERS];

  int arf_stack_size = 0;

  int extend_frame_count = 0;

  int pframe_qindex = cpi->tpl_stats[2].base_qindex;

  int frame_gop_offset = 0;

  RefCntBuffer *frame_bufs = cm->buffer_pool->frame_bufs;

  int8_t recon_frame_index[REFS_PER_FRAME + MAX_ARF_LAYERS];

  memset(recon_frame_index, -1, sizeof(recon_frame_index));

  stack_init(arf_index_stack, MAX_ARF_LAYERS);

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

    if (frame_bufs[i].ref_count == 0) {

      alloc_frame_mvs(cm, i);

      if (vpx_realloc_frame_buffer(&frame_bufs[i].buf, cm->width, cm->height,

                                   cm->subsampling_x, cm->subsampling_y,

#if CONFIG_VP9_HIGHBITDEPTH

                                   cm->use_highbitdepth,

#endif

                                   VP9_ENC_BORDER_IN_PIXELS, cm->byte_alignment,

                                   NULL, NULL, NULL))

        vpx_internal_error(&cm->error, VPX_CODEC_MEM_ERROR,

                           "Failed to allocate frame buffer");

      recon_frame_index[frame_idx] = i;

      ++frame_idx;

      if (frame_idx >= REFS_PER_FRAME + cpi->oxcf.enable_auto_arf) break;

    }

  }

  for (i = 0; i < REFS_PER_FRAME + 1; ++i) {

    assert(recon_frame_index[i] >= 0);

    cpi->tpl_recon_frames[i] = &frame_bufs[recon_frame_index[i]].buf;

  }

  *tpl_group_frames = 0;

  // Initialize Golden reference frame.

  gf_picture[0].frame = get_ref_frame_buffer(cpi, GOLDEN_FRAME);

  for (i = 0; i < 3; ++i) gf_picture[0].ref_frame[i] = -REFS_PER_FRAME;

  gf_picture[0].update_type = gf_group->update_type[0];

  gld_index = 0;

  ++*tpl_group_frames;

  gf_picture[-1].frame = get_ref_frame_buffer(cpi, LAST_FRAME);

  gf_picture[-2].frame = get_ref_frame_buffer(cpi, ALTREF_FRAME);

  // Initialize base layer ARF frame

  gf_picture[1].frame = cpi->Source;

  gf_picture[1].ref_frame[0] = gld_index;

  gf_picture[1].ref_frame[1] = lst_index;

  gf_picture[1].ref_frame[2] = alt_index;

  gf_picture[1].update_type = gf_group->update_type[1];

  alt_index = 1;

  ++*tpl_group_frames;

  // Initialize P frames

  for (frame_idx = 2; frame_idx < MAX_ARF_GOP_SIZE; ++frame_idx) {

    struct lookahead_entry *buf;

    frame_gop_offset = gf_group->frame_gop_index[frame_idx];

    buf = vp9_lookahead_peek(cpi->lookahead, frame_gop_offset - 1);

    if (buf == NULL) break;

    gf_picture[frame_idx].frame = &buf->img;

    gf_picture[frame_idx].ref_frame[0] = gld_index;

    gf_picture[frame_idx].ref_frame[1] = lst_index;

    gf_picture[frame_idx].ref_frame[2] = alt_index;

    gf_picture[frame_idx].update_type = gf_group->update_type[frame_idx];

    switch (gf_group->update_type[frame_idx]) {

      case ARF_UPDATE:

        stack_push(arf_index_stack, alt_index, arf_stack_size);

        ++arf_stack_size;

        alt_index = frame_idx;

        break;

      case LF_UPDATE: lst_index = frame_idx; break;

      case OVERLAY_UPDATE:

        gld_index = frame_idx;

        alt_index = stack_pop(arf_index_stack, arf_stack_size);

        --arf_stack_size;

        break;

      case USE_BUF_FRAME:

        lst_index = alt_index;

        alt_index = stack_pop(arf_index_stack, arf_stack_size);

        --arf_stack_size;

        break;

      default: break;

    }

    ++*tpl_group_frames;

    // The length of group of pictures is baseline_gf_interval, plus the

    // beginning golden frame from last GOP, plus the last overlay frame in

    // the same GOP.

    if (frame_idx == gf_group->gf_group_size) break;

  }

  alt_index = -1;

  ++frame_idx;

  ++frame_gop_offset;

  // Extend two frames outside the current gf group.

  for (; frame_idx < MAX_LAG_BUFFERS && extend_frame_count < 2; ++frame_idx) {

    struct lookahead_entry *buf =

        vp9_lookahead_peek(cpi->lookahead, frame_gop_offset - 1);

    if (buf == NULL) break;

    cpi->tpl_stats[frame_idx].base_qindex = pframe_qindex;

    gf_picture[frame_idx].frame = &buf->img;

    gf_picture[frame_idx].ref_frame[0] = gld_index;

    gf_picture[frame_idx].ref_frame[1] = lst_index;

    gf_picture[frame_idx].ref_frame[2] = alt_index;

    gf_picture[frame_idx].update_type = LF_UPDATE;

    lst_index = frame_idx;

    ++*tpl_group_frames;

    ++extend_frame_count;

    ++frame_gop_offset;

  }

  return extend_frame_count;

}

static void init_tpl_stats(VP9_COMP *cpi) {

  int frame_idx;

  for (frame_idx = 0; frame_idx < MAX_ARF_GOP_SIZE; ++frame_idx) {

    TplDepFrame *tpl_frame = &cpi->tpl_stats[frame_idx];

    memset(tpl_frame->tpl_stats_ptr, 0,

           tpl_frame->height * tpl_frame->width *

               sizeof(*tpl_frame->tpl_stats_ptr));

    tpl_frame->is_valid = 0;

  }

}

static void free_tpl_frame_stats_list(VpxTplGopStats *tpl_gop_stats) {

  int frame_idx;

  for (frame_idx = 0; frame_idx < tpl_gop_stats->size; ++frame_idx) {

    vpx_free(tpl_gop_stats->frame_stats_list[frame_idx].block_stats_list);

  }

  vpx_free(tpl_gop_stats->frame_stats_list);

}

static void init_tpl_stats_before_propagation(

    struct vpx_internal_error_info *error_info, VpxTplGopStats *tpl_gop_stats,

    TplDepFrame *tpl_stats, int tpl_gop_frames, int frame_width,

    int frame_height) {

  int frame_idx;

  free_tpl_frame_stats_list(tpl_gop_stats);

  CHECK_MEM_ERROR(

      error_info, tpl_gop_stats->frame_stats_list,

      vpx_calloc(tpl_gop_frames, sizeof(*tpl_gop_stats->frame_stats_list)));

  tpl_gop_stats->size = tpl_gop_frames;

  for (frame_idx = 0; frame_idx < tpl_gop_frames; ++frame_idx) {

    const int mi_rows = tpl_stats[frame_idx].mi_rows;

    const int mi_cols = tpl_stats[frame_idx].mi_cols;

    CHECK_MEM_ERROR(

        error_info, tpl_gop_stats->frame_stats_list[frame_idx].block_stats_list,

        vpx_calloc(

            mi_rows * mi_cols,

            sizeof(

                *tpl_gop_stats->frame_stats_list[frame_idx].block_stats_list)));

    tpl_gop_stats->frame_stats_list[frame_idx].num_blocks = mi_rows * mi_cols;

    tpl_gop_stats->frame_stats_list[frame_idx].frame_width = frame_width;

    tpl_gop_stats->frame_stats_list[frame_idx].frame_height = frame_height;

  }

}

#if CONFIG_NON_GREEDY_MV

static uint32_t full_pixel_motion_search(VP9_COMP *cpi, ThreadData *td,

                                         MotionField *motion_field,

                                         int frame_idx, uint8_t *cur_frame_buf,

                                         uint8_t *ref_frame_buf, int stride,

                                         BLOCK_SIZE bsize, int mi_row,

                                         int mi_col, MV *mv) {

  MACROBLOCK *const x = &td->mb;

  MACROBLOCKD *const xd = &x->e_mbd;

  MV_SPEED_FEATURES *const mv_sf = &cpi->sf.mv;

  int step_param;

  uint32_t bestsme = UINT_MAX;

  const MvLimits tmp_mv_limits = x->mv_limits;

  // lambda is used to adjust the importance of motion vector consistency.

  // TODO(angiebird): Figure out lambda's proper value.

  const int lambda = cpi->tpl_stats[frame_idx].lambda;

  int_mv nb_full_mvs[NB_MVS_NUM];

  int nb_full_mv_num;

  MV best_ref_mv1 = { 0, 0 };

  MV best_ref_mv1_full; /* full-pixel value of best_ref_mv1 */

  best_ref_mv1_full.col = best_ref_mv1.col >> 3;

  best_ref_mv1_full.row = best_ref_mv1.row >> 3;

  // Setup frame pointers

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

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

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

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

  step_param = mv_sf->reduce_first_step_size;

  step_param = VPXMIN(step_param, MAX_MVSEARCH_STEPS - 2);

  vp9_set_mv_search_range(&x->mv_limits, &best_ref_mv1);

  nb_full_mv_num =

      vp9_prepare_nb_full_mvs(motion_field, mi_row, mi_col, nb_full_mvs);

  vp9_full_pixel_diamond_new(cpi, x, bsize, &best_ref_mv1_full, step_param,

                             lambda, 1, nb_full_mvs, nb_full_mv_num, mv);

  /* restore UMV window */

  x->mv_limits = tmp_mv_limits;

  return bestsme;

}

static uint32_t sub_pixel_motion_search(VP9_COMP *cpi, ThreadData *td,

                                        uint8_t *cur_frame_buf,

                                        uint8_t *ref_frame_buf, int stride,

                                        BLOCK_SIZE bsize, MV *mv) {

  MACROBLOCK *const x = &td->mb;

  MACROBLOCKD *const xd = &x->e_mbd;

  MV_SPEED_FEATURES *const mv_sf = &cpi->sf.mv;

  uint32_t bestsme = UINT_MAX;

  uint32_t distortion;

  uint32_t sse;

  int cost_list[5];

  MV best_ref_mv1 = { 0, 0 };

  // Setup frame pointers

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

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

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

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

  // TODO(yunqing): may use higher tap interp filter than 2 taps.

  // Ignore mv costing by sending NULL pointer instead of cost array

  bestsme = cpi->find_fractional_mv_step(

      x, mv, &best_ref_mv1, cpi->common.allow_high_precision_mv, x->errorperbit,

      &cpi->fn_ptr[bsize], 0, mv_sf->subpel_search_level,

      cond_cost_list(cpi, cost_list), NULL, NULL, &distortion, &sse, NULL, 0, 0,

      USE_2_TAPS);

  return bestsme;

}

#else  // CONFIG_NON_GREEDY_MV

static uint32_t motion_compensated_prediction(VP9_COMP *cpi, ThreadData *td,

                                              uint8_t *cur_frame_buf,

                                              uint8_t *ref_frame_buf,

                                              int stride, BLOCK_SIZE bsize,

                                              MV *mv) {

  MACROBLOCK *const x = &td->mb;

  MACROBLOCKD *const xd = &x->e_mbd;

  MV_SPEED_FEATURES *const mv_sf = &cpi->sf.mv;

  const SEARCH_METHODS search_method = NSTEP;

  int step_param;

  int sadpb = x->sadperbit16;

  uint32_t bestsme = UINT_MAX;

  uint32_t distortion;

  uint32_t sse;

  int cost_list[5];

  const MvLimits tmp_mv_limits = x->mv_limits;

  MV best_ref_mv1 = { 0, 0 };

  MV best_ref_mv1_full; /* full-pixel value of best_ref_mv1 */

  best_ref_mv1_full.col = best_ref_mv1.col >> 3;

  best_ref_mv1_full.row = best_ref_mv1.row >> 3;

  // Setup frame pointers

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

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

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

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

  step_param = mv_sf->reduce_first_step_size;

  step_param = VPXMIN(step_param, MAX_MVSEARCH_STEPS - 2);

  vp9_set_mv_search_range(&x->mv_limits, &best_ref_mv1);

  vp9_full_pixel_search(cpi, x, bsize, &best_ref_mv1_full, step_param,

                        search_method, sadpb, cond_cost_list(cpi, cost_list),

                        &best_ref_mv1, mv, 0, 0);

  /* restore UMV window */

  x->mv_limits = tmp_mv_limits;

  // TODO(yunqing): may use higher tap interp filter than 2 taps.

  // Ignore mv costing by sending NULL pointer instead of cost array

  bestsme = cpi->find_fractional_mv_step(

      x, mv, &best_ref_mv1, cpi->common.allow_high_precision_mv, x->errorperbit,

      &cpi->fn_ptr[bsize], 0, mv_sf->subpel_search_level,

      cond_cost_list(cpi, cost_list), NULL, NULL, &distortion, &sse, NULL, 0, 0,

      USE_2_TAPS);

  return bestsme;

}

#endif

static int get_overlap_area(int grid_pos_row, int grid_pos_col, int ref_pos_row,

                            int ref_pos_col, int block, BLOCK_SIZE bsize) {

  int width = 0, height = 0;

  int bw = 4 << b_width_log2_lookup[bsize];

  int bh = 4 << b_height_log2_lookup[bsize];

  switch (block) {

    case 0:

      width = grid_pos_col + bw - ref_pos_col;

      height = grid_pos_row + bh - ref_pos_row;

      break;

    case 1:

      width = ref_pos_col + bw - grid_pos_col;

      height = grid_pos_row + bh - ref_pos_row;

      break;

    case 2:

      width = grid_pos_col + bw - ref_pos_col;

      height = ref_pos_row + bh - grid_pos_row;

      break;

    case 3:

      width = ref_pos_col + bw - grid_pos_col;

      height = ref_pos_row + bh - grid_pos_row;

      break;

    default: assert(0);

  }

  return width * height;

}

static int round_floor(int ref_pos, int bsize_pix) {

  int round;

  if (ref_pos < 0)

    round = -(1 + (-ref_pos - 1) / bsize_pix);

  else

    round = ref_pos / bsize_pix;

  return round;

}

static void tpl_model_store(TplDepStats *tpl_stats, int mi_row, int mi_col,

                            BLOCK_SIZE bsize, int stride) {

  const int mi_height = num_8x8_blocks_high_lookup[bsize];

  const int mi_width = num_8x8_blocks_wide_lookup[bsize];

  const TplDepStats *src_stats = &tpl_stats[mi_row * stride + mi_col];

  int idx, idy;

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

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

      TplDepStats *tpl_ptr = &tpl_stats[(mi_row + idy) * stride + mi_col + idx];

      const int64_t mc_flow = tpl_ptr->mc_flow;

      const int64_t mc_ref_cost = tpl_ptr->mc_ref_cost;

      *tpl_ptr = *src_stats;

      tpl_ptr->mc_flow = mc_flow;

      tpl_ptr->mc_ref_cost = mc_ref_cost;

      tpl_ptr->mc_dep_cost = tpl_ptr->intra_cost + tpl_ptr->mc_flow;

    }

  }

}

static void tpl_store_before_propagation(VpxTplBlockStats *tpl_block_stats,

                                         TplDepStats *tpl_stats, int mi_row,

                                         int mi_col, BLOCK_SIZE bsize,

                                         int src_stride, int64_t recon_error,

                                         int64_t rate_cost, int ref_frame_idx,

                                         int mi_rows, int mi_cols) {

  const int mi_height = num_8x8_blocks_high_lookup[bsize];

  const int mi_width = num_8x8_blocks_wide_lookup[bsize];

  const TplDepStats *src_stats = &tpl_stats[mi_row * src_stride + mi_col];

  int idx, idy;

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

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

      if (mi_row + idy >= mi_rows || mi_col + idx >= mi_cols) continue;

      VpxTplBlockStats *tpl_block_stats_ptr =

          &tpl_block_stats[(mi_row + idy) * mi_cols + mi_col + idx];

      tpl_block_stats_ptr->row = mi_row * 8 + idy * 8;

      tpl_block_stats_ptr->col = mi_col * 8 + idx * 8;

      tpl_block_stats_ptr->inter_cost = src_stats->inter_cost;

      tpl_block_stats_ptr->intra_cost = src_stats->intra_cost;

      // inter/intra_cost here is calculated with SATD which should be close

      // enough to be used as inter/intra_pred_error

      tpl_block_stats_ptr->inter_pred_err = src_stats->inter_cost;

      tpl_block_stats_ptr->intra_pred_err = src_stats->intra_cost;

      tpl_block_stats_ptr->srcrf_dist = recon_error << TPL_DEP_COST_SCALE_LOG2;

      tpl_block_stats_ptr->srcrf_rate = rate_cost << TPL_DEP_COST_SCALE_LOG2;

      tpl_block_stats_ptr->mv_r = src_stats->mv.as_mv.row;

      tpl_block_stats_ptr->mv_c = src_stats->mv.as_mv.col;

      tpl_block_stats_ptr->ref_frame_index = ref_frame_idx;

    }

  }

}

static void tpl_model_update_b(TplDepFrame *tpl_frame, TplDepStats *tpl_stats,

                               int mi_row, int mi_col, const BLOCK_SIZE bsize) {

  if (tpl_stats->ref_frame_index < 0) return;

  TplDepFrame *ref_tpl_frame = &tpl_frame[tpl_stats->ref_frame_index];

  TplDepStats *ref_stats = ref_tpl_frame->tpl_stats_ptr;

  MV mv = tpl_stats->mv.as_mv;

  int mv_row = mv.row >> 3;

  int mv_col = mv.col >> 3;

  int ref_pos_row = mi_row * MI_SIZE + mv_row;

  int ref_pos_col = mi_col * MI_SIZE + mv_col;

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

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

  const int mi_height = num_8x8_blocks_high_lookup[bsize];

  const int mi_width = num_8x8_blocks_wide_lookup[bsize];

  const int pix_num = bw * bh;

  // top-left on grid block location in pixel

  int grid_pos_row_base = round_floor(ref_pos_row, bh) * bh;

  int grid_pos_col_base = round_floor(ref_pos_col, bw) * bw;

  int block;

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

    int grid_pos_row = grid_pos_row_base + bh * (block >> 1);

    int grid_pos_col = grid_pos_col_base + bw * (block & 0x01);

    if (grid_pos_row >= 0 && grid_pos_row < ref_tpl_frame->mi_rows * MI_SIZE &&

        grid_pos_col >= 0 && grid_pos_col < ref_tpl_frame->mi_cols * MI_SIZE) {

      int overlap_area = get_overlap_area(

          grid_pos_row, grid_pos_col, ref_pos_row, ref_pos_col, block, bsize);

      int ref_mi_row = round_floor(grid_pos_row, bh) * mi_height;

      int ref_mi_col = round_floor(grid_pos_col, bw) * mi_width;

      int64_t mc_flow = tpl_stats->mc_dep_cost -

                        (tpl_stats->mc_dep_cost * tpl_stats->inter_cost) /

                            tpl_stats->intra_cost;

      int idx, idy;

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

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

          TplDepStats *des_stats =

              &ref_stats[(ref_mi_row + idy) * ref_tpl_frame->stride +

                         (ref_mi_col + idx)];

          des_stats->mc_flow += (mc_flow * overlap_area) / pix_num;

          des_stats->mc_ref_cost +=

              ((tpl_stats->intra_cost - tpl_stats->inter_cost) * overlap_area) /

              pix_num;

          assert(overlap_area >= 0);

        }

      }

    }

  }

}

static void tpl_model_update(TplDepFrame *tpl_frame, TplDepStats *tpl_stats,

                             int mi_row, int mi_col, const BLOCK_SIZE bsize) {

  int idx, idy;

  const int mi_height = num_8x8_blocks_high_lookup[bsize];

  const int mi_width = num_8x8_blocks_wide_lookup[bsize];

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

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

      TplDepStats *tpl_ptr =

          &tpl_stats[(mi_row + idy) * tpl_frame->stride + (mi_col + idx)];

      tpl_model_update_b(tpl_frame, tpl_ptr, mi_row + idy, mi_col + idx,

                         BLOCK_8X8);

    }

  }

}

static void get_quantize_error(MACROBLOCK *x, int plane, tran_low_t *coeff,

                               tran_low_t *qcoeff, tran_low_t *dqcoeff,

                               TX_SIZE tx_size, int64_t *recon_error,

                               int64_t *sse, uint16_t *eob) {

  MACROBLOCKD *const xd = &x->e_mbd;

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

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

  const ScanOrder *const scan_order = &vp9_default_scan_orders[tx_size];

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

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

  // skip block condition should be handled before this is called.

  assert(!x->skip_block);

#if CONFIG_VP9_HIGHBITDEPTH

  if (xd->cur_buf->flags & YV12_FLAG_HIGHBITDEPTH) {

    vp9_highbd_quantize_fp_32x32(coeff, pix_num, p, qcoeff, dqcoeff,

                                 pd->dequant, eob, scan_order);

  } else {

    vp9_quantize_fp_32x32(coeff, pix_num, p, qcoeff, dqcoeff, pd->dequant, eob,

                          scan_order);

  }

#else

  vp9_quantize_fp_32x32(coeff, pix_num, p, qcoeff, dqcoeff, pd->dequant, eob,

                        scan_order);

#endif  // CONFIG_VP9_HIGHBITDEPTH

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

  *recon_error = VPXMAX(*recon_error, 1);

  *sse = (*sse) >> shift;

  *sse = VPXMAX(*sse, 1);

}

#if CONFIG_VP9_HIGHBITDEPTH

void vp9_highbd_wht_fwd_txfm(int16_t *src_diff, int bw, tran_low_t *coeff,

                             TX_SIZE tx_size) {

  // TODO(sdeng): Implement SIMD based high bit-depth Hadamard transforms.

  switch (tx_size) {

    case TX_8X8: vpx_highbd_hadamard_8x8(src_diff, bw, coeff); break;

    case TX_16X16: vpx_highbd_hadamard_16x16(src_diff, bw, coeff); break;

    case TX_32X32: vpx_highbd_hadamard_32x32(src_diff, bw, coeff); break;

    default: assert(0);

  }

}

#endif  // CONFIG_VP9_HIGHBITDEPTH

void vp9_wht_fwd_txfm(int16_t *src_diff, int bw, tran_low_t *coeff,

                      TX_SIZE tx_size) {

  switch (tx_size) {

    case TX_8X8: vpx_hadamard_8x8(src_diff, bw, coeff); break;

    case TX_16X16: vpx_hadamard_16x16(src_diff, bw, coeff); break;

    case TX_32X32: vpx_hadamard_32x32(src_diff, bw, coeff); break;

    default: assert(0);

  }

}

static void set_mv_limits(const VP9_COMMON *cm, MACROBLOCK *x, int mi_row,

                          int mi_col) {

  x->mv_limits.row_min = -((mi_row * MI_SIZE) + (17 - 2 * VP9_INTERP_EXTEND));

  x->mv_limits.row_max =

      (cm->mi_rows - 1 - mi_row) * MI_SIZE + (17 - 2 * VP9_INTERP_EXTEND);

  x->mv_limits.col_min = -((mi_col * MI_SIZE) + (17 - 2 * VP9_INTERP_EXTEND));

  x->mv_limits.col_max =

      ((cm->mi_cols - 1 - mi_col) * MI_SIZE) + (17 - 2 * VP9_INTERP_EXTEND);

}

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

  const ScanOrder *const scan_order = &vp9_scan_orders[tx_size][DCT_DCT];

  int rate_cost = 1;

  int idx;

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

  for (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 << VP9_PROB_COST_SHIFT);

}

static void mode_estimation(VP9_COMP *cpi, MACROBLOCK *x, MACROBLOCKD *xd,

                            struct scale_factors *sf, GF_PICTURE *gf_picture,

                            int frame_idx, TplDepFrame *tpl_frame,

                            int16_t *src_diff, tran_low_t *coeff,

                            tran_low_t *qcoeff, tran_low_t *dqcoeff, int mi_row,

                            int mi_col, BLOCK_SIZE bsize, TX_SIZE tx_size,

                            YV12_BUFFER_CONFIG *ref_frame[], uint8_t *predictor,

                            int64_t *recon_error, int64_t *rate_cost,

                            int64_t *sse, int *ref_frame_idx) {

  VP9_COMMON *cm = &cpi->common;

  ThreadData *td = &cpi->td;

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

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

  const int pix_num = bw * bh;

  int best_rf_idx = -1;

  int_mv best_mv;

  int64_t best_inter_cost = INT64_MAX;

  int64_t inter_cost;

  int rf_idx;

  const InterpKernel *const kernel = vp9_filter_kernels[EIGHTTAP];

  int64_t best_intra_cost = INT64_MAX;

  int64_t intra_cost;

  PREDICTION_MODE mode;

  int mb_y_offset = mi_row * MI_SIZE * xd->cur_buf->y_stride + mi_col * MI_SIZE;

  MODE_INFO mi_above, mi_left;

  const int mi_height = num_8x8_blocks_high_lookup[bsize];

  const int mi_width = num_8x8_blocks_wide_lookup[bsize];

  TplDepStats *tpl_stats =

      &tpl_frame->tpl_stats_ptr[mi_row * tpl_frame->stride + mi_col];

  xd->mb_to_top_edge = -((mi_row * MI_SIZE) * 8);

  xd->mb_to_bottom_edge = ((cm->mi_rows - 1 - mi_row) * MI_SIZE) * 8;

  xd->mb_to_left_edge = -((mi_col * MI_SIZE) * 8);

  xd->mb_to_right_edge = ((cm->mi_cols - 1 - mi_col) * MI_SIZE) * 8;

  xd->above_mi = (mi_row > 0) ? &mi_above : NULL;

  xd->left_mi = (mi_col > 0) ? &mi_left : NULL;

  // Intra prediction search

  for (mode = DC_PRED; mode <= TM_PRED; ++mode) {

    uint8_t *src, *dst;

    int src_stride, dst_stride;

    src = xd->cur_buf->y_buffer + mb_y_offset;

    src_stride = xd->cur_buf->y_stride;

    dst = &predictor[0];

    dst_stride = bw;

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

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

    vp9_predict_intra_block(xd, b_width_log2_lookup[bsize], tx_size, mode, src,

                            src_stride, dst, dst_stride, 0, 0, 0);

#if CONFIG_VP9_HIGHBITDEPTH

    if (xd->cur_buf->flags & YV12_FLAG_HIGHBITDEPTH) {

      vpx_highbd_subtract_block(bh, bw, src_diff, bw, src, src_stride, dst,

                                dst_stride, xd->bd);

      vp9_highbd_wht_fwd_txfm(src_diff, bw, coeff, tx_size);

      intra_cost = vpx_highbd_satd(coeff, pix_num);

    } else {

      vpx_subtract_block(bh, bw, src_diff, bw, src, src_stride, dst,

                         dst_stride);

      vp9_wht_fwd_txfm(src_diff, bw, coeff, tx_size);

      intra_cost = vpx_satd(coeff, pix_num);

    }

#else

    vpx_subtract_block(bh, bw, src_diff, bw, src, src_stride, dst, dst_stride);

    vp9_wht_fwd_txfm(src_diff, bw, coeff, tx_size);

    intra_cost = vpx_satd(coeff, pix_num);

#endif  // CONFIG_VP9_HIGHBITDEPTH

    if (intra_cost < best_intra_cost) best_intra_cost = intra_cost;

  }

  // Motion compensated prediction

  best_mv.as_int = 0;

  set_mv_limits(cm, x, mi_row, mi_col);

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

    int_mv mv;

#if CONFIG_NON_GREEDY_MV

    MotionField *motion_field;

#endif

    if (ref_frame[rf_idx] == NULL) continue;

#if CONFIG_NON_GREEDY_MV

    (void)td;

    motion_field = vp9_motion_field_info_get_motion_field(

        &cpi->motion_field_info, frame_idx, rf_idx, bsize);

    mv = vp9_motion_field_mi_get_mv(motion_field, mi_row, mi_col);

#else

    motion_compensated_prediction(cpi, td, xd->cur_buf->y_buffer + mb_y_offset,

                                  ref_frame[rf_idx]->y_buffer + mb_y_offset,

                                  xd->cur_buf->y_stride, bsize, &mv.as_mv);