/*

 * 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/av1_common_int.h"

#include "av1/common/blockd.h"

#include "av1/common/enums.h"

#include "av1/common/reconintra.h"

#include "av1/encoder/aq_complexity.h"

#include "av1/encoder/aq_variance.h"

#include "av1/encoder/context_tree.h"

#include "av1/encoder/encoder.h"

#include "av1/encoder/encodeframe.h"

#include "av1/encoder/encodeframe_utils.h"

#include "av1/encoder/encodemv.h"

#include "av1/encoder/motion_search_facade.h"

#include "av1/encoder/partition_search.h"

#include "av1/encoder/partition_strategy.h"

#include "av1/encoder/reconinter_enc.h"

#include "av1/encoder/tokenize.h"

#include "av1/encoder/var_based_part.h"

#include "av1/encoder/av1_ml_partition_models.h"

#if CONFIG_TUNE_VMAF

#include "av1/encoder/tune_vmaf.h"

#endif

#define COLLECT_MOTION_SEARCH_FEATURE_SB 0

#define ML_PARTITION_WHOLE_TREE_DECISION 0

void av1_reset_part_sf(PARTITION_SPEED_FEATURES *part_sf) {

  part_sf->partition_search_type = SEARCH_PARTITION;

  part_sf->less_rectangular_check_level = 0;

  part_sf->use_square_partition_only_threshold = BLOCK_128X128;

  part_sf->auto_max_partition_based_on_simple_motion = NOT_IN_USE;

  part_sf->default_max_partition_size = BLOCK_LARGEST;

  part_sf->default_min_partition_size = BLOCK_4X4;

  part_sf->adjust_var_based_rd_partitioning = 0;

  part_sf->max_intra_bsize = BLOCK_LARGEST;

  // This setting only takes effect when partition_search_type is set

  // to FIXED_PARTITION.

  part_sf->fixed_partition_size = BLOCK_16X16;

  // Recode loop tolerance %.

  part_sf->partition_search_breakout_dist_thr = 0;

  part_sf->partition_search_breakout_rate_thr = 0;

  part_sf->prune_ext_partition_types_search_level = 0;

  part_sf->prune_part4_search = 0;

  part_sf->ml_prune_partition = 0;

  part_sf->ml_early_term_after_part_split_level = 0;

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

    part_sf->ml_partition_search_breakout_thresh[i] =

        -1;  // -1 means not enabled.

  }

  part_sf->simple_motion_search_prune_agg = SIMPLE_AGG_LVL0;

  part_sf->simple_motion_search_split = 0;

  part_sf->simple_motion_search_prune_rect = 0;

  part_sf->simple_motion_search_early_term_none = 0;

  part_sf->simple_motion_search_reduce_search_steps = 0;

  part_sf->intra_cnn_based_part_prune_level = 0;

  part_sf->ext_partition_eval_thresh = BLOCK_8X8;

  part_sf->rect_partition_eval_thresh = BLOCK_128X128;

  part_sf->prune_ext_part_using_split_info = 0;

  part_sf->prune_rectangular_split_based_on_qidx = 0;

  part_sf->early_term_after_none_split = 0;

  part_sf->ml_predict_breakout_level = 0;

  part_sf->prune_sub_8x8_partition_level = 0;

  part_sf->simple_motion_search_rect_split = 0;

  part_sf->reuse_prev_rd_results_for_part_ab = 0;

  part_sf->reuse_best_prediction_for_part_ab = 0;

  part_sf->use_best_rd_for_pruning = 0;

  part_sf->skip_non_sq_part_based_on_none = 0;

}

// Reset speed features that works for the baseline encoding, but

// blocks the external partition search.

void av1_reset_sf_for_ext_part(AV1_COMP *const cpi) {

  cpi->sf.inter_sf.prune_ref_frame_for_rect_partitions = 0;

}

#if !CONFIG_REALTIME_ONLY

// If input |features| is NULL, write tpl stats to file for each super block.

// Otherwise, store tpl stats to |features|.

// The tpl stats is computed in the unit of tpl_bsize_1d (16x16).

// When writing to text file:

// The first row contains super block position, super block size,

// tpl unit length, number of units in the super block.

// The second row contains the intra prediction cost for each unit.

// The third row contains the inter prediction cost for each unit.

// The forth row contains the motion compensated dependency cost for each unit.

static void collect_tpl_stats_sb(const AV1_COMP *const cpi,

                                 const BLOCK_SIZE bsize, const int mi_row,

                                 const int mi_col,

                                 aom_partition_features_t *features) {

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

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

  if (gf_group->update_type[cpi->gf_frame_index] == INTNL_OVERLAY_UPDATE ||

      gf_group->update_type[cpi->gf_frame_index] == OVERLAY_UPDATE) {

    return;

  }

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

  TplDepFrame *tpl_frame = &tpl_data->tpl_frame[cpi->gf_frame_index];

  TplDepStats *tpl_stats = tpl_frame->tpl_stats_ptr;

  // If tpl stats is not established, early return

  if (!tpl_data->ready || gf_group->max_layer_depth_allowed == 0) {

    if (features != NULL) features->sb_features.tpl_features.available = 0;

    return;

  }

  const int tpl_stride = tpl_frame->stride;

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

  const int mi_width =

      AOMMIN(mi_size_wide[bsize], cm->mi_params.mi_cols - mi_col);

  const int mi_height =

      AOMMIN(mi_size_high[bsize], cm->mi_params.mi_rows - mi_row);

  const int col_steps = (mi_width / step) + ((mi_width % step) > 0);

  const int row_steps = (mi_height / step) + ((mi_height % step) > 0);

  const int num_blocks = col_steps * row_steps;

  if (features == NULL) {

    char filename[256];

    snprintf(filename, sizeof(filename), "%s/tpl_feature_sb%d",

             cpi->oxcf.partition_info_path, cpi->sb_counter);

    FILE *pfile = fopen(filename, "w");

    fprintf(pfile, "%d,%d,%d,%d,%d\n", mi_row, mi_col, bsize,

            tpl_data->tpl_bsize_1d, num_blocks);

    int count = 0;

    for (int row = 0; row < mi_height; row += step) {

      for (int col = 0; col < mi_width; col += step) {

        TplDepStats *this_stats =

            &tpl_stats[av1_tpl_ptr_pos(mi_row + row, mi_col + col, tpl_stride,

                                       tpl_data->tpl_stats_block_mis_log2)];

        fprintf(pfile, "%.0f", (double)this_stats->intra_cost);

        if (count < num_blocks - 1) fprintf(pfile, ",");

        ++count;

      }

    }

    fprintf(pfile, "\n");

    count = 0;

    for (int row = 0; row < mi_height; row += step) {

      for (int col = 0; col < mi_width; col += step) {

        TplDepStats *this_stats =

            &tpl_stats[av1_tpl_ptr_pos(mi_row + row, mi_col + col, tpl_stride,

                                       tpl_data->tpl_stats_block_mis_log2)];

        fprintf(pfile, "%.0f", (double)this_stats->inter_cost);

        if (count < num_blocks - 1) fprintf(pfile, ",");

        ++count;

      }

    }

    fprintf(pfile, "\n");

    count = 0;

    for (int row = 0; row < mi_height; row += step) {

      for (int col = 0; col < mi_width; col += step) {

        TplDepStats *this_stats =

            &tpl_stats[av1_tpl_ptr_pos(mi_row + row, mi_col + col, tpl_stride,

                                       tpl_data->tpl_stats_block_mis_log2)];

        const int64_t mc_dep_delta =

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

                   this_stats->mc_dep_dist);

        fprintf(pfile, "%.0f", (double)mc_dep_delta);

        if (count < num_blocks - 1) fprintf(pfile, ",");

        ++count;

      }

    }

    fclose(pfile);

  } else {

    features->sb_features.tpl_features.available = 1;

    features->sb_features.tpl_features.tpl_unit_length = tpl_data->tpl_bsize_1d;

    features->sb_features.tpl_features.num_units = num_blocks;

    int count = 0;

    for (int row = 0; row < mi_height; row += step) {

      for (int col = 0; col < mi_width; col += step) {

        TplDepStats *this_stats =

            &tpl_stats[av1_tpl_ptr_pos(mi_row + row, mi_col + col, tpl_stride,

                                       tpl_data->tpl_stats_block_mis_log2)];

        const int64_t mc_dep_delta =

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

                   this_stats->mc_dep_dist);

        features->sb_features.tpl_features.intra_cost[count] =

            this_stats->intra_cost;

        features->sb_features.tpl_features.inter_cost[count] =

            this_stats->inter_cost;

        features->sb_features.tpl_features.mc_dep_cost[count] = mc_dep_delta;

        ++count;

      }

    }

  }

}

#endif  // !CONFIG_REALTIME_ONLY

static void update_txfm_count(MACROBLOCK *x, MACROBLOCKD *xd,

                              FRAME_COUNTS *counts, TX_SIZE tx_size, int depth,

                              int blk_row, int blk_col,

                              uint8_t allow_update_cdf) {

  MB_MODE_INFO *mbmi = xd->mi[0];

  const BLOCK_SIZE bsize = mbmi->bsize;

  const int max_blocks_high = max_block_high(xd, bsize, 0);

  const int max_blocks_wide = max_block_wide(xd, bsize, 0);

  int ctx = txfm_partition_context(xd->above_txfm_context + blk_col,

                                   xd->left_txfm_context + blk_row, mbmi->bsize,

                                   tx_size);

  const int txb_size_index = av1_get_txb_size_index(bsize, blk_row, blk_col);

  const TX_SIZE plane_tx_size = mbmi->inter_tx_size[txb_size_index];

  if (blk_row >= max_blocks_high || blk_col >= max_blocks_wide) return;

  assert(tx_size > TX_4X4);

  if (depth == MAX_VARTX_DEPTH) {

    // Don't add to counts in this case

    mbmi->tx_size = tx_size;

    txfm_partition_update(xd->above_txfm_context + blk_col,

                          xd->left_txfm_context + blk_row, tx_size, tx_size);

    return;

  }

  if (tx_size == plane_tx_size) {

#if CONFIG_ENTROPY_STATS

    ++counts->txfm_partition[ctx][0];

#endif

    if (allow_update_cdf)

      update_cdf(xd->tile_ctx->txfm_partition_cdf[ctx], 0, 2);

    mbmi->tx_size = tx_size;

    txfm_partition_update(xd->above_txfm_context + blk_col,

                          xd->left_txfm_context + blk_row, tx_size, tx_size);

  } else {

    const TX_SIZE sub_txs = sub_tx_size_map[tx_size];

    const int bsw = tx_size_wide_unit[sub_txs];

    const int bsh = tx_size_high_unit[sub_txs];

#if CONFIG_ENTROPY_STATS

    ++counts->txfm_partition[ctx][1];

#endif

    if (allow_update_cdf)

      update_cdf(xd->tile_ctx->txfm_partition_cdf[ctx], 1, 2);

    ++x->txfm_search_info.txb_split_count;

    if (sub_txs == TX_4X4) {

      mbmi->inter_tx_size[txb_size_index] = TX_4X4;

      mbmi->tx_size = TX_4X4;

      txfm_partition_update(xd->above_txfm_context + blk_col,

                            xd->left_txfm_context + blk_row, TX_4X4, tx_size);

      return;

    }

    for (int row = 0; row < tx_size_high_unit[tx_size]; row += bsh) {

      for (int col = 0; col < tx_size_wide_unit[tx_size]; col += bsw) {

        int offsetr = row;

        int offsetc = col;

        update_txfm_count(x, xd, counts, sub_txs, depth + 1, blk_row + offsetr,

                          blk_col + offsetc, allow_update_cdf);

      }

    }

  }

}

static void tx_partition_count_update(const AV1_COMMON *const cm, MACROBLOCK *x,

                                      BLOCK_SIZE plane_bsize,

                                      FRAME_COUNTS *td_counts,

                                      uint8_t allow_update_cdf) {

  MACROBLOCKD *xd = &x->e_mbd;

  const int mi_width = mi_size_wide[plane_bsize];

  const int mi_height = mi_size_high[plane_bsize];

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

  const int bh = tx_size_high_unit[max_tx_size];

  const int bw = tx_size_wide_unit[max_tx_size];

  xd->above_txfm_context =

      cm->above_contexts.txfm[xd->tile.tile_row] + xd->mi_col;

  xd->left_txfm_context =

      xd->left_txfm_context_buffer + (xd->mi_row & MAX_MIB_MASK);

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

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

      update_txfm_count(x, xd, td_counts, max_tx_size, 0, idy, idx,

                        allow_update_cdf);

    }

  }

}

static void set_txfm_context(MACROBLOCKD *xd, TX_SIZE tx_size, int blk_row,

                             int blk_col) {

  MB_MODE_INFO *mbmi = xd->mi[0];

  const BLOCK_SIZE bsize = mbmi->bsize;

  const int max_blocks_high = max_block_high(xd, bsize, 0);

  const int max_blocks_wide = max_block_wide(xd, bsize, 0);

  const int txb_size_index = av1_get_txb_size_index(bsize, blk_row, blk_col);

  const TX_SIZE plane_tx_size = mbmi->inter_tx_size[txb_size_index];

  if (blk_row >= max_blocks_high || blk_col >= max_blocks_wide) return;

  if (tx_size == plane_tx_size) {

    mbmi->tx_size = tx_size;

    txfm_partition_update(xd->above_txfm_context + blk_col,

                          xd->left_txfm_context + blk_row, tx_size, tx_size);

  } else {

    if (tx_size == TX_8X8) {

      mbmi->inter_tx_size[txb_size_index] = TX_4X4;

      mbmi->tx_size = TX_4X4;

      txfm_partition_update(xd->above_txfm_context + blk_col,

                            xd->left_txfm_context + blk_row, TX_4X4, tx_size);

      return;

    }

    const TX_SIZE sub_txs = sub_tx_size_map[tx_size];

    const int bsw = tx_size_wide_unit[sub_txs];

    const int bsh = tx_size_high_unit[sub_txs];

    const int row_end =

        AOMMIN(tx_size_high_unit[tx_size], max_blocks_high - blk_row);

    const int col_end =

        AOMMIN(tx_size_wide_unit[tx_size], max_blocks_wide - blk_col);

    for (int row = 0; row < row_end; row += bsh) {

      const int offsetr = blk_row + row;

      for (int col = 0; col < col_end; col += bsw) {

        const int offsetc = blk_col + col;

        set_txfm_context(xd, sub_txs, offsetr, offsetc);

      }

    }

  }

}

static void tx_partition_set_contexts(const AV1_COMMON *const cm,

                                      MACROBLOCKD *xd, BLOCK_SIZE plane_bsize) {

  const int mi_width = mi_size_wide[plane_bsize];

  const int mi_height = mi_size_high[plane_bsize];

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

  const int bh = tx_size_high_unit[max_tx_size];

  const int bw = tx_size_wide_unit[max_tx_size];

  xd->above_txfm_context =

      cm->above_contexts.txfm[xd->tile.tile_row] + xd->mi_col;

  xd->left_txfm_context =

      xd->left_txfm_context_buffer + (xd->mi_row & MAX_MIB_MASK);

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

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

      set_txfm_context(xd, max_tx_size, idy, idx);

    }

  }

}

static void update_zeromv_cnt(const AV1_COMP *const cpi,

                              const MB_MODE_INFO *const mi, int mi_row,

                              int mi_col, BLOCK_SIZE bsize) {

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

      mi->segment_id > CR_SEGMENT_ID_BOOST2) {

    return;

  }

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

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

  const int bw = mi_size_wide[bsize] >> 1;

  const int bh = mi_size_high[bsize] >> 1;

  const int xmis = AOMMIN((cm->mi_params.mi_cols - mi_col) >> 1, bw);

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

  const int block_index =

      (mi_row >> 1) * (cm->mi_params.mi_cols >> 1) + (mi_col >> 1);

  for (int y = 0; y < ymis; y++) {

    for (int x = 0; x < xmis; x++) {

      // consec_zero_mv is in the scale of 8x8 blocks

      const int map_offset = block_index + y * (cm->mi_params.mi_cols >> 1) + x;

      if (abs(mv.row) < 10 && abs(mv.col) < 10) {

        if (cpi->consec_zero_mv[map_offset] < 255)

          cpi->consec_zero_mv[map_offset]++;

      } else {

        cpi->consec_zero_mv[map_offset] = 0;

      }

    }

  }

}

static void encode_superblock(const AV1_COMP *const cpi, TileDataEnc *tile_data,

                              ThreadData *td, TokenExtra **t, RUN_TYPE dry_run,

                              BLOCK_SIZE bsize, int *rate) {

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

  const int num_planes = av1_num_planes(cm);

  MACROBLOCK *const x = &td->mb;

  MACROBLOCKD *const xd = &x->e_mbd;

  MB_MODE_INFO **mi_4x4 = xd->mi;

  MB_MODE_INFO *mbmi = mi_4x4[0];

  const int seg_skip =

      segfeature_active(&cm->seg, mbmi->segment_id, SEG_LVL_SKIP);

  const int mis = cm->mi_params.mi_stride;

  const int mi_width = mi_size_wide[bsize];

  const int mi_height = mi_size_high[bsize];

  const int is_inter = is_inter_block(mbmi);

  // Initialize tx_mode and tx_size_search_method

  TxfmSearchParams *txfm_params = &x->txfm_search_params;

  set_tx_size_search_method(

      cm, &cpi->winner_mode_params, txfm_params,

      cpi->sf.winner_mode_sf.enable_winner_mode_for_tx_size_srch, 1);

  const int mi_row = xd->mi_row;

  const int mi_col = xd->mi_col;

  if (!is_inter) {

    xd->cfl.store_y = store_cfl_required(cm, xd);

    mbmi->skip_txfm = 1;

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

      av1_encode_intra_block_plane(cpi, x, bsize, plane, dry_run,

                                   cpi->optimize_seg_arr[mbmi->segment_id]);

    }

    // If there is at least one lossless segment, force the skip for intra

    // block to be 0, in order to avoid the segment_id to be changed by in

    // write_segment_id().

    if (!cpi->common.seg.segid_preskip && cpi->common.seg.update_map &&

        cpi->enc_seg.has_lossless_segment)

      mbmi->skip_txfm = 0;

    xd->cfl.store_y = 0;

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

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

        if (mbmi->palette_mode_info.palette_size[plane] > 0) {

          if (!dry_run) {

            av1_tokenize_color_map(x, plane, t, bsize, mbmi->tx_size,

                                   PALETTE_MAP, tile_data->allow_update_cdf,

                                   td->counts);

          } else if (dry_run == DRY_RUN_COSTCOEFFS) {

            *rate +=

                av1_cost_color_map(x, plane, bsize, mbmi->tx_size, PALETTE_MAP);

          }

        }

      }

    }

    av1_update_intra_mb_txb_context(cpi, td, dry_run, bsize,

                                    tile_data->allow_update_cdf);

  } else {

    int ref;

    const int is_compound = has_second_ref(mbmi);

    set_ref_ptrs(cm, xd, mbmi->ref_frame[0], mbmi->ref_frame[1]);

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

      const YV12_BUFFER_CONFIG *cfg =

          get_ref_frame_yv12_buf(cm, mbmi->ref_frame[ref]);

      assert(IMPLIES(!is_intrabc_block(mbmi), cfg));

      av1_setup_pre_planes(xd, ref, cfg, mi_row, mi_col,

                           xd->block_ref_scale_factors[ref], num_planes);

    }

    // Predicted sample of inter mode (for Luma plane) cannot be reused if

    // nonrd_check_partition_merge_mode or nonrd_check_partition_split speed

    // feature is enabled, Since in such cases the buffer may not contain the

    // predicted sample of best mode.

    const int start_plane =

        (cpi->sf.rt_sf.reuse_inter_pred_nonrd &&

         (!cpi->sf.rt_sf.nonrd_check_partition_merge_mode) &&

         (!cpi->sf.rt_sf.nonrd_check_partition_split) &&

         cm->seq_params->bit_depth == AOM_BITS_8)

            ? 1

            : 0;

    av1_enc_build_inter_predictor(cm, xd, mi_row, mi_col, NULL, bsize,

                                  start_plane, av1_num_planes(cm) - 1);

    if (mbmi->motion_mode == OBMC_CAUSAL) {

      assert(cpi->oxcf.motion_mode_cfg.enable_obmc);

      av1_build_obmc_inter_predictors_sb(cm, xd);

    }

#if CONFIG_MISMATCH_DEBUG

    if (dry_run == OUTPUT_ENABLED) {

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

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

        int pixel_c, pixel_r;

        mi_to_pixel_loc(&pixel_c, &pixel_r, mi_col, mi_row, 0, 0,

                        pd->subsampling_x, pd->subsampling_y);

        if (!is_chroma_reference(mi_row, mi_col, bsize, pd->subsampling_x,

                                 pd->subsampling_y))

          continue;

        mismatch_record_block_pre(pd->dst.buf, pd->dst.stride,

                                  cm->current_frame.order_hint, plane, pixel_c,

                                  pixel_r, pd->width, pd->height,

                                  xd->cur_buf->flags & YV12_FLAG_HIGHBITDEPTH);

      }

    }

#else

    (void)num_planes;

#endif

    av1_encode_sb(cpi, x, bsize, dry_run);

    av1_tokenize_sb_vartx(cpi, td, dry_run, bsize, rate,

                          tile_data->allow_update_cdf);

  }

  if (!dry_run) {

    if (av1_allow_intrabc(cm) && is_intrabc_block(mbmi)) td->intrabc_used = 1;

    if (txfm_params->tx_mode_search_type == TX_MODE_SELECT &&

        !xd->lossless[mbmi->segment_id] && mbmi->bsize > BLOCK_4X4 &&

        !(is_inter && (mbmi->skip_txfm || seg_skip))) {

      if (is_inter) {

        tx_partition_count_update(cm, x, bsize, td->counts,

                                  tile_data->allow_update_cdf);

      } else {

        if (mbmi->tx_size != max_txsize_rect_lookup[bsize])

          ++x->txfm_search_info.txb_split_count;

        if (block_signals_txsize(bsize)) {

          const int tx_size_ctx = get_tx_size_context(xd);

          const int32_t tx_size_cat = bsize_to_tx_size_cat(bsize);

          const int depth = tx_size_to_depth(mbmi->tx_size, bsize);

          const int max_depths = bsize_to_max_depth(bsize);

          if (tile_data->allow_update_cdf)

            update_cdf(xd->tile_ctx->tx_size_cdf[tx_size_cat][tx_size_ctx],

                       depth, max_depths + 1);

#if CONFIG_ENTROPY_STATS

          ++td->counts->intra_tx_size[tx_size_cat][tx_size_ctx][depth];

#endif

        }

      }

      assert(IMPLIES(is_rect_tx(mbmi->tx_size), is_rect_tx_allowed(xd, mbmi)));

    } else {

      int i, j;

      TX_SIZE intra_tx_size;

      // The new intra coding scheme requires no change of transform size

      if (is_inter) {

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

          intra_tx_size = TX_4X4;

        } else {

          intra_tx_size =

              tx_size_from_tx_mode(bsize, txfm_params->tx_mode_search_type);

        }

      } else {

        intra_tx_size = mbmi->tx_size;

      }

      for (j = 0; j < mi_height; j++)

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

          if (mi_col + i < cm->mi_params.mi_cols &&

              mi_row + j < cm->mi_params.mi_rows)

            mi_4x4[mis * j + i]->tx_size = intra_tx_size;

      if (intra_tx_size != max_txsize_rect_lookup[bsize])

        ++x->txfm_search_info.txb_split_count;

    }

  }

  if (txfm_params->tx_mode_search_type == TX_MODE_SELECT &&

      block_signals_txsize(mbmi->bsize) && is_inter &&

      !(mbmi->skip_txfm || seg_skip) && !xd->lossless[mbmi->segment_id]) {

    if (dry_run) tx_partition_set_contexts(cm, xd, bsize);

  } else {

    TX_SIZE tx_size = mbmi->tx_size;

    // The new intra coding scheme requires no change of transform size

    if (is_inter) {

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

        tx_size = TX_4X4;

      } else {

        tx_size = tx_size_from_tx_mode(bsize, txfm_params->tx_mode_search_type);

      }

    } else {

      tx_size = (bsize > BLOCK_4X4) ? tx_size : TX_4X4;

    }

    mbmi->tx_size = tx_size;

    set_txfm_ctxs(tx_size, xd->width, xd->height,

                  (mbmi->skip_txfm || seg_skip) && is_inter_block(mbmi), xd);

  }

  if (is_inter_block(mbmi) && !xd->is_chroma_ref && is_cfl_allowed(xd)) {

    cfl_store_block(xd, mbmi->bsize, mbmi->tx_size);

  }

  if (!dry_run) {

    if (cpi->oxcf.pass == AOM_RC_ONE_PASS && cpi->svc.temporal_layer_id == 0 &&

        cpi->sf.rt_sf.use_temporal_noise_estimate &&

        (!cpi->ppi->use_svc ||

         (cpi->ppi->use_svc &&

          !cpi->svc.layer_context[cpi->svc.temporal_layer_id].is_key_frame &&

          cpi->svc.spatial_layer_id == cpi->svc.number_spatial_layers - 1)))

      update_zeromv_cnt(cpi, mbmi, mi_row, mi_col, bsize);

  }

}

static void setup_block_rdmult(const AV1_COMP *const cpi, MACROBLOCK *const x,

                               int mi_row, int mi_col, BLOCK_SIZE bsize,

                               AQ_MODE aq_mode, MB_MODE_INFO *mbmi) {

  x->rdmult = cpi->rd.RDMULT;

  if (aq_mode != NO_AQ) {

    assert(mbmi != NULL);

    if (aq_mode == VARIANCE_AQ) {

      if (cpi->vaq_refresh) {

        const int energy = bsize <= BLOCK_16X16

                               ? x->mb_energy

                               : av1_log_block_var(cpi, x, bsize);

        mbmi->segment_id = energy;

      }

      x->rdmult = set_segment_rdmult(cpi, x, mbmi->segment_id);

    } else if (aq_mode == COMPLEXITY_AQ) {

      x->rdmult = set_segment_rdmult(cpi, x, mbmi->segment_id);

    } else if (aq_mode == CYCLIC_REFRESH_AQ) {

      // If segment is boosted, use rdmult for that segment.

      if (cyclic_refresh_segment_id_boosted(mbmi->segment_id))

        x->rdmult = av1_cyclic_refresh_get_rdmult(cpi->cyclic_refresh);

    }

  }

#if !CONFIG_REALTIME_ONLY

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

  if (cm->delta_q_info.delta_q_present_flag &&

      !cpi->sf.rt_sf.use_nonrd_pick_mode) {

    x->rdmult = av1_get_cb_rdmult(cpi, x, bsize, mi_row, mi_col);

  }

#endif  // !CONFIG_REALTIME_ONLY

  if (cpi->oxcf.tune_cfg.tuning == AOM_TUNE_SSIM) {

    av1_set_ssim_rdmult(cpi, &x->errorperbit, bsize, mi_row, mi_col,

                        &x->rdmult);

  }

#if CONFIG_TUNE_VMAF

  if (cpi->oxcf.tune_cfg.tuning == AOM_TUNE_VMAF_WITHOUT_PREPROCESSING ||

      cpi->oxcf.tune_cfg.tuning == AOM_TUNE_VMAF_MAX_GAIN ||

      cpi->oxcf.tune_cfg.tuning == AOM_TUNE_VMAF_NEG_MAX_GAIN) {

    av1_set_vmaf_rdmult(cpi, x, bsize, mi_row, mi_col, &x->rdmult);

  }

#endif

#if CONFIG_TUNE_BUTTERAUGLI

  if (cpi->oxcf.tune_cfg.tuning == AOM_TUNE_BUTTERAUGLI) {

    av1_set_butteraugli_rdmult(cpi, x, bsize, mi_row, mi_col, &x->rdmult);

  }

#endif

}

void av1_set_offsets_without_segment_id(const AV1_COMP *const cpi,

                                        const TileInfo *const tile,

                                        MACROBLOCK *const x, int mi_row,

                                        int mi_col, BLOCK_SIZE bsize) {

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

  const int num_planes = av1_num_planes(cm);

  MACROBLOCKD *const xd = &x->e_mbd;

  assert(bsize < BLOCK_SIZES_ALL);

  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_entropy_context(xd, mi_row, mi_col, num_planes);

  xd->above_txfm_context = cm->above_contexts.txfm[tile->tile_row] + mi_col;

  xd->left_txfm_context =

      xd->left_txfm_context_buffer + (mi_row & MAX_MIB_MASK);

  // Set up destination pointers.

  av1_setup_dst_planes(xd->plane, bsize, &cm->cur_frame->buf, mi_row, mi_col, 0,

                       num_planes);

  // Set up limit values for MV components.

  // Mv beyond the range do not produce new/different prediction block.

  av1_set_mv_limits(&cm->mi_params, &x->mv_limits, mi_row, mi_col, mi_height,

                    mi_width, cpi->oxcf.border_in_pixels);

  set_plane_n4(xd, mi_width, mi_height, num_planes);

  // Set up distance of MB to edge of frame in 1/8th pel units.

  assert(!(mi_col & (mi_width - 1)) && !(mi_row & (mi_height - 1)));

  set_mi_row_col(xd, tile, mi_row, mi_height, mi_col, mi_width,

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

  // Set up source buffers.

  av1_setup_src_planes(x, cpi->source, mi_row, mi_col, num_planes, bsize);

  // required by av1_append_sub8x8_mvs_for_idx() and av1_find_best_ref_mvs()

  xd->tile = *tile;

}

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

                     MACROBLOCK *const x, int mi_row, int mi_col,

                     BLOCK_SIZE bsize) {

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

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

  MACROBLOCKD *const xd = &x->e_mbd;

  MB_MODE_INFO *mbmi;

  av1_set_offsets_without_segment_id(cpi, tile, x, mi_row, mi_col, bsize);

  // Setup segment ID.

  mbmi = xd->mi[0];

  mbmi->segment_id = 0;

  if (seg->enabled) {

    if (seg->enabled && !cpi->vaq_refresh) {

      const uint8_t *const map =

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

      mbmi->segment_id =

          map ? get_segment_id(&cm->mi_params, map, bsize, mi_row, mi_col) : 0;

    }

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

  }

}

/*!\brief Hybrid intra mode search.

 *

 * \ingroup intra_mode_search

 * \callgraph

 * \callergraph

 * This is top level function for mode search for intra frames in non-RD

 * optimized case. Depending on speed feature and block size it calls

 * either non-RD or RD optimized intra mode search.

 *

 * \param[in]    cpi            Top-level encoder structure

 * \param[in]    x              Pointer to structure holding all the data for

                                the current macroblock

 * \param[in]    rd_cost        Struct to keep track of the RD information

 * \param[in]    bsize          Current block size

 * \param[in]    ctx            Structure to hold snapshot of coding context

                                during the mode picking process

 *

 * \return Nothing is returned. Instead, the MB_MODE_INFO struct inside x

 * is modified to store information about the best mode computed

 * in this function. The rd_cost struct is also updated with the RD stats

 * corresponding to the best mode found.

 */

static AOM_INLINE void hybrid_intra_mode_search(AV1_COMP *cpi,

                                                MACROBLOCK *const x,

                                                RD_STATS *rd_cost,

                                                BLOCK_SIZE bsize,

                                                PICK_MODE_CONTEXT *ctx) {

  if (cpi->sf.rt_sf.hybrid_intra_pickmode && bsize < BLOCK_16X16)

    av1_rd_pick_intra_mode_sb(cpi, x, rd_cost, bsize, ctx, INT64_MAX);

  else

    av1_nonrd_pick_intra_mode(cpi, x, rd_cost, bsize, ctx);

}

// For real time/allintra row-mt enabled multi-threaded encoding with cost

// update frequency set to COST_UPD_TILE/COST_UPD_OFF, tile ctxt is not updated

// at superblock level. Thus, it is not required for the encoding of top-right

// superblock be complete for updating tile ctxt. However, when encoding a block

// whose right edge is also the superblock edge, intra and inter mode evaluation

// (ref mv list population) require the encoding of the top-right superblock to

// be complete. So, here, we delay the waiting of threads until the need for the

// data from the top-right superblock region.

static AOM_INLINE void wait_for_top_right_sb(

    AV1EncRowMultiThreadInfo *enc_row_mt, AV1EncRowMultiThreadSync *row_mt_sync,

    TileInfo *tile_info, BLOCK_SIZE sb_size, int sb_mi_size_log2,

    BLOCK_SIZE bsize, int mi_row, int mi_col) {

  const int sb_size_in_mi = mi_size_wide[sb_size];

  const int bw_in_mi = mi_size_wide[bsize];

  const int blk_row_in_sb = mi_row & (sb_size_in_mi - 1);

  const int blk_col_in_sb = mi_col & (sb_size_in_mi - 1);

  const int top_right_block_in_sb =

      (blk_row_in_sb == 0) && (blk_col_in_sb + bw_in_mi >= sb_size_in_mi);

  // Don't wait if the block is the not the top-right block in the superblock.

  if (!top_right_block_in_sb) return;

  // Wait for the top-right superblock to finish encoding.

  const int sb_row_in_tile =

      (mi_row - tile_info->mi_row_start) >> sb_mi_size_log2;

  const int sb_col_in_tile =

      (mi_col - tile_info->mi_col_start) >> sb_mi_size_log2;

  (*(enc_row_mt->sync_read_ptr))(row_mt_sync, sb_row_in_tile, sb_col_in_tile);

}

/*!\brief Interface for AV1 mode search for an individual coding block

 *

 * \ingroup partition_search

 * \callgraph

 * \callergraph

 * Searches prediction modes, transform, and coefficient coding modes for an

 * individual coding block. This function is the top-level interface that

 * directs the encoder to the proper mode search function, among these

 * implemented for inter/intra + rd/non-rd + non-skip segment/skip segment.

 *

 * \param[in]    cpi            Top-level encoder structure

 * \param[in]    tile_data      Pointer to struct holding adaptive

 *                              data/contexts/models for the tile during

 *                              encoding

 * \param[in]    x              Pointer to structure holding all the data for

 *                              the current macroblock

 * \param[in]    mi_row         Row coordinate of the block in a step size of

 *                              MI_SIZE

 * \param[in]    mi_col         Column coordinate of the block in a step size of

 *                              MI_SIZE

 * \param[in]    rd_cost        Pointer to structure holding rate and distortion

 *                              stats for the current block

 * \param[in]    partition      Partition mode of the parent block

 * \param[in]    bsize          Current block size

 * \param[in]    ctx            Pointer to structure holding coding contexts and

 *                              chosen modes for the current block

 * \param[in]    best_rd        Upper bound of rd cost of a valid partition

 *

 * \return Nothing is returned. Instead, the chosen modes and contexts necessary

 * for reconstruction are stored in ctx, the rate-distortion stats are stored in

 * rd_cost. If no valid mode leading to rd_cost <= best_rd, the status will be

 * signalled by an INT64_MAX rd_cost->rdcost.

 */

static void pick_sb_modes(AV1_COMP *const cpi, TileDataEnc *tile_data,

                          MACROBLOCK *const x, int mi_row, int mi_col,

                          RD_STATS *rd_cost, PARTITION_TYPE partition,

                          BLOCK_SIZE bsize, PICK_MODE_CONTEXT *ctx,

                          RD_STATS best_rd) {

  if (cpi->sf.part_sf.use_best_rd_for_pruning && best_rd.rdcost < 0) {

    ctx->rd_stats.rdcost = INT64_MAX;

    ctx->rd_stats.skip_txfm = 0;

    av1_invalid_rd_stats(rd_cost);

    return;

  }

  av1_set_offsets(cpi, &tile_data->tile_info, x, mi_row, mi_col, bsize);

  if (cpi->sf.part_sf.reuse_prev_rd_results_for_part_ab &&

      ctx->rd_mode_is_ready) {

    assert(ctx->mic.bsize == bsize);

    assert(ctx->mic.partition == partition);

    rd_cost->rate = ctx->rd_stats.rate;

    rd_cost->dist = ctx->rd_stats.dist;

    rd_cost->rdcost = ctx->rd_stats.rdcost;

    return;

  }

  AV1_COMMON *const cm = &cpi->common;

  const int num_planes = av1_num_planes(cm);

  MACROBLOCKD *const xd = &x->e_mbd;

  MB_MODE_INFO *mbmi;

  struct macroblock_plane *const p = x->plane;

  struct macroblockd_plane *const pd = xd->plane;

  const AQ_MODE aq_mode = cpi->oxcf.q_cfg.aq_mode;

  TxfmSearchInfo *txfm_info = &x->txfm_search_info;

  int i;

  // This is only needed for real time/allintra row-mt enabled multi-threaded

  // encoding with cost update frequency set to COST_UPD_TILE/COST_UPD_OFF.

  wait_for_top_right_sb(&cpi->mt_info.enc_row_mt, &tile_data->row_mt_sync,

                        &tile_data->tile_info, cm->seq_params->sb_size,

                        cm->seq_params->mib_size_log2, bsize, mi_row, mi_col);

#if CONFIG_COLLECT_COMPONENT_TIMING

  start_timing(cpi, rd_pick_sb_modes_time);

#endif

  mbmi = xd->mi[0];

  mbmi->bsize = bsize;

  mbmi->partition = partition;

#if CONFIG_RD_DEBUG

  mbmi->mi_row = mi_row;

  mbmi->mi_col = mi_col;

#endif

  // Sets up the tx_type_map buffer in MACROBLOCKD.

  xd->tx_type_map = txfm_info->tx_type_map_;

  xd->tx_type_map_stride = mi_size_wide[bsize];

  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];

  ctx->skippable = 0;

  // Set to zero to make sure we do not use the previous encoded frame stats

  mbmi->skip_txfm = 0;

  // Reset skip mode flag.

  mbmi->skip_mode = 0;

  if (is_cur_buf_hbd(xd)) {

    x->source_variance = av1_high_get_sby_perpixel_variance(

        cpi, &x->plane[0].src, bsize, xd->bd);

  } else {

    x->source_variance =

        av1_get_sby_perpixel_variance(cpi, &x->plane[0].src, bsize);

  }

  // Initialize default mode evaluation params

  set_mode_eval_params(cpi, x, DEFAULT_EVAL);

  // Save rdmult before it might be changed, so it can be restored later.

  const int orig_rdmult = x->rdmult;

  setup_block_rdmult(cpi, x, mi_row, mi_col, bsize, aq_mode, mbmi);

  // Set error per bit for current rdmult

  av1_set_error_per_bit(&x->errorperbit, x->rdmult);

  av1_rd_cost_update(x->rdmult, &best_rd);

  // If set best_rd.rdcost to INT64_MAX, the encoder will not use any previous

  // rdcost information for the following mode search.

  // Disabling the feature could get some coding gain, with encoder slowdown.

  if (!cpi->sf.part_sf.use_best_rd_for_pruning) {

    av1_invalid_rd_stats(&best_rd);

  }

  // Find best coding mode & reconstruct the MB so it is available

  // as a predictor for MBs that follow in the SB

  if (frame_is_intra_only(cm)) {

#if CONFIG_COLLECT_COMPONENT_TIMING

    start_timing(cpi, av1_rd_pick_intra_mode_sb_time);

#endif

    av1_rd_pick_intra_mode_sb(cpi, x, rd_cost, bsize, ctx, best_rd.rdcost);

#if CONFIG_COLLECT_COMPONENT_TIMING

    end_timing(cpi, av1_rd_pick_intra_mode_sb_time);

#endif

  } else {

#if CONFIG_COLLECT_COMPONENT_TIMING

    start_timing(cpi, av1_rd_pick_inter_mode_sb_time);

#endif

    if (segfeature_active(&cm->seg, mbmi->segment_id, SEG_LVL_SKIP)) {

      av1_rd_pick_inter_mode_sb_seg_skip(cpi, tile_data, x, mi_row, mi_col,

                                         rd_cost, bsize, ctx, best_rd.rdcost);

    } else {

      av1_rd_pick_inter_mode(cpi, tile_data, x, rd_cost, bsize, ctx,

                             best_rd.rdcost);

    }

#if CONFIG_COLLECT_COMPONENT_TIMING

    end_timing(cpi, av1_rd_pick_inter_mode_sb_time);

#endif

  }

  // Examine the resulting rate and for AQ mode 2 make a segment choice.

  if (rd_cost->rate != INT_MAX && aq_mode == COMPLEXITY_AQ &&

      bsize >= BLOCK_16X16) {

    av1_caq_select_segment(cpi, x, bsize, mi_row, mi_col, rd_cost->rate);

  }

  x->rdmult = orig_rdmult;

  // TODO(jingning) The rate-distortion optimization flow needs to be

  // refactored to provide proper exit/return handle.

  if (rd_cost->rate == INT_MAX) rd_cost->rdcost = INT64_MAX;

  ctx->rd_stats.rate = rd_cost->rate;

  ctx->rd_stats.dist = rd_cost->dist;

  ctx->rd_stats.rdcost = rd_cost->rdcost;

#if CONFIG_COLLECT_COMPONENT_TIMING

  end_timing(cpi, rd_pick_sb_modes_time);

#endif

}

static void update_stats(const AV1_COMMON *const cm, ThreadData *td) {

  MACROBLOCK *x = &td->mb;

  MACROBLOCKD *const xd = &x->e_mbd;

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

  const MB_MODE_INFO_EXT *const mbmi_ext = &x->mbmi_ext;

  const CurrentFrame *const current_frame = &cm->current_frame;

  const BLOCK_SIZE bsize = mbmi->bsize;

  FRAME_CONTEXT *fc = xd->tile_ctx;

  const int seg_ref_active =

      segfeature_active(&cm->seg, mbmi->segment_id, SEG_LVL_REF_FRAME);