/*

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

#include "av1/common/reconintra.h"

#include "av1/encoder/block.h"

#include "av1/encoder/hybrid_fwd_txfm.h"

#include "av1/common/idct.h"

#include "av1/encoder/model_rd.h"

#include "av1/encoder/random.h"

#include "av1/encoder/rdopt_utils.h"

#include "av1/encoder/sorting_network.h"

#include "av1/encoder/tx_prune_model_weights.h"

#include "av1/encoder/tx_search.h"

#include "av1/encoder/txb_rdopt.h"

#define PROB_THRESH_OFFSET_TX_TYPE 100

struct rdcost_block_args {

  const AV1_COMP *cpi;

  MACROBLOCK *x;

  ENTROPY_CONTEXT t_above[MAX_MIB_SIZE];

  ENTROPY_CONTEXT t_left[MAX_MIB_SIZE];

  RD_STATS rd_stats;

  int64_t current_rd;

  int64_t best_rd;

  int exit_early;

  int incomplete_exit;

  FAST_TX_SEARCH_MODE ftxs_mode;

  int skip_trellis;

};

typedef struct {

  int64_t rd;

  int txb_entropy_ctx;

  TX_TYPE tx_type;

} TxCandidateInfo;

// origin_threshold * 128 / 100

static const uint32_t skip_pred_threshold[3][BLOCK_SIZES_ALL] = {

  {

      64, 64, 64, 70, 60, 60, 68, 68, 68, 68, 68,

      68, 68, 68, 68, 68, 64, 64, 70, 70, 68, 68,

  },

  {

      88, 88, 88, 86, 87, 87, 68, 68, 68, 68, 68,

      68, 68, 68, 68, 68, 88, 88, 86, 86, 68, 68,

  },

  {

      90, 93, 93, 90, 93, 93, 74, 74, 74, 74, 74,

      74, 74, 74, 74, 74, 90, 90, 90, 90, 74, 74,

  },

};

// lookup table for predict_skip_txfm

// int max_tx_size = max_txsize_rect_lookup[bsize];

// if (tx_size_high[max_tx_size] > 16 || tx_size_wide[max_tx_size] > 16)

//   max_tx_size = AOMMIN(max_txsize_lookup[bsize], TX_16X16);

static const TX_SIZE max_predict_sf_tx_size[BLOCK_SIZES_ALL] = {

  TX_4X4,   TX_4X8,   TX_8X4,   TX_8X8,   TX_8X16,  TX_16X8,

  TX_16X16, TX_16X16, TX_16X16, TX_16X16, TX_16X16, TX_16X16,

  TX_16X16, TX_16X16, TX_16X16, TX_16X16, TX_4X16,  TX_16X4,

  TX_8X8,   TX_8X8,   TX_16X16, TX_16X16,

};

// look-up table for sqrt of number of pixels in a transform block

// rounded up to the nearest integer.

static const int sqrt_tx_pixels_2d[TX_SIZES_ALL] = { 4,  8,  16, 32, 32, 6,  6,

                                                     12, 12, 23, 23, 32, 32, 8,

                                                     8,  16, 16, 23, 23 };

static INLINE uint32_t get_block_residue_hash(MACROBLOCK *x, BLOCK_SIZE bsize) {

  const int rows = block_size_high[bsize];

  const int cols = block_size_wide[bsize];

  const int16_t *diff = x->plane[0].src_diff;

  const uint32_t hash =

      av1_get_crc32c_value(&x->txfm_search_info.mb_rd_record->crc_calculator,

                           (uint8_t *)diff, 2 * rows * cols);

  return (hash << 5) + bsize;

}

static INLINE int32_t find_mb_rd_info(const MB_RD_RECORD *const mb_rd_record,

                                      const int64_t ref_best_rd,

                                      const uint32_t hash) {

  int32_t match_index = -1;

  if (ref_best_rd != INT64_MAX) {

    for (int i = 0; i < mb_rd_record->num; ++i) {

      const int index = (mb_rd_record->index_start + i) % RD_RECORD_BUFFER_LEN;

      // If there is a match in the mb_rd_record, fetch the RD decision and

      // terminate early.

      if (mb_rd_record->mb_rd_info[index].hash_value == hash) {

        match_index = index;

        break;

      }

    }

  }

  return match_index;

}

static AOM_INLINE void fetch_mb_rd_info(int n4,

                                        const MB_RD_INFO *const mb_rd_info,

                                        RD_STATS *const rd_stats,

                                        MACROBLOCK *const x) {

  MACROBLOCKD *const xd = &x->e_mbd;

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

  mbmi->tx_size = mb_rd_info->tx_size;

  memcpy(x->txfm_search_info.blk_skip, mb_rd_info->blk_skip,

         sizeof(mb_rd_info->blk_skip[0]) * n4);

  av1_copy(mbmi->inter_tx_size, mb_rd_info->inter_tx_size);

  av1_copy_array(xd->tx_type_map, mb_rd_info->tx_type_map, n4);

  *rd_stats = mb_rd_info->rd_stats;

}

// Compute the pixel domain distortion from diff on all visible 4x4s in the

// transform block.

static INLINE int64_t pixel_diff_dist(const MACROBLOCK *x, int plane,

                                      int blk_row, int blk_col,

                                      const BLOCK_SIZE plane_bsize,

                                      const BLOCK_SIZE tx_bsize,

                                      unsigned int *block_mse_q8) {

  int visible_rows, visible_cols;

  const MACROBLOCKD *xd = &x->e_mbd;

  get_txb_dimensions(xd, plane, plane_bsize, blk_row, blk_col, tx_bsize, NULL,

                     NULL, &visible_cols, &visible_rows);

  const int diff_stride = block_size_wide[plane_bsize];

  const int16_t *diff = x->plane[plane].src_diff;

  diff += ((blk_row * diff_stride + blk_col) << MI_SIZE_LOG2);

  uint64_t sse =

      aom_sum_squares_2d_i16(diff, diff_stride, visible_cols, visible_rows);

  if (block_mse_q8 != NULL) {

    if (visible_cols > 0 && visible_rows > 0)

      *block_mse_q8 =

          (unsigned int)((256 * sse) / (visible_cols * visible_rows));

    else

      *block_mse_q8 = UINT_MAX;

  }

  return sse;

}

// Computes the residual block's SSE and mean on all visible 4x4s in the

// transform block

static INLINE int64_t pixel_diff_stats(

    MACROBLOCK *x, int plane, int blk_row, int blk_col,

    const BLOCK_SIZE plane_bsize, const BLOCK_SIZE tx_bsize,

    unsigned int *block_mse_q8, int64_t *per_px_mean, uint64_t *block_var) {

  int visible_rows, visible_cols;

  const MACROBLOCKD *xd = &x->e_mbd;

  get_txb_dimensions(xd, plane, plane_bsize, blk_row, blk_col, tx_bsize, NULL,

                     NULL, &visible_cols, &visible_rows);

  const int diff_stride = block_size_wide[plane_bsize];

  const int16_t *diff = x->plane[plane].src_diff;

  diff += ((blk_row * diff_stride + blk_col) << MI_SIZE_LOG2);

  uint64_t sse = 0;

  int sum = 0;

  sse = aom_sum_sse_2d_i16(diff, diff_stride, visible_cols, visible_rows, &sum);

  if (visible_cols > 0 && visible_rows > 0) {

    double norm_factor = 1.0 / (visible_cols * visible_rows);

    int sign_sum = sum > 0 ? 1 : -1;

    // Conversion to transform domain

    *per_px_mean = (int64_t)(norm_factor * abs(sum)) << 7;

    *per_px_mean = sign_sum * (*per_px_mean);

    *block_mse_q8 = (unsigned int)(norm_factor * (256 * sse));

    *block_var = (uint64_t)(sse - (uint64_t)(norm_factor * sum * sum));

  } else {

    *block_mse_q8 = UINT_MAX;

  }

  return sse;

}

// Uses simple features on top of DCT coefficients to quickly predict

// whether optimal RD decision is to skip encoding the residual.

// The sse value is stored in dist.

static int predict_skip_txfm(MACROBLOCK *x, BLOCK_SIZE bsize, int64_t *dist,

                             int reduced_tx_set) {

  const TxfmSearchParams *txfm_params = &x->txfm_search_params;

  const int bw = block_size_wide[bsize];

  const int bh = block_size_high[bsize];

  const MACROBLOCKD *xd = &x->e_mbd;

  const int16_t dc_q = av1_dc_quant_QTX(x->qindex, 0, xd->bd);

  *dist = pixel_diff_dist(x, 0, 0, 0, bsize, bsize, NULL);

  const int64_t mse = *dist / bw / bh;

  // Normalized quantizer takes the transform upscaling factor (8 for tx size

  // smaller than 32) into account.

  const int16_t normalized_dc_q = dc_q >> 3;

  const int64_t mse_thresh = (int64_t)normalized_dc_q * normalized_dc_q / 8;

  // For faster early skip decision, use dist to compare against threshold so

  // that quality risk is less for the skip=1 decision. Otherwise, use mse

  // since the fwd_txfm coeff checks will take care of quality

  // TODO(any): Use dist to return 0 when skip_txfm_level is 1

  int64_t pred_err = (txfm_params->skip_txfm_level >= 2) ? *dist : mse;

  // Predict not to skip when error is larger than threshold.

  if (pred_err > mse_thresh) return 0;

  // Return as skip otherwise for aggressive early skip

  else if (txfm_params->skip_txfm_level >= 2)

    return 1;

  const int max_tx_size = max_predict_sf_tx_size[bsize];

  const int tx_h = tx_size_high[max_tx_size];

  const int tx_w = tx_size_wide[max_tx_size];

  DECLARE_ALIGNED(32, tran_low_t, coefs[32 * 32]);

  TxfmParam param;

  param.tx_type = DCT_DCT;

  param.tx_size = max_tx_size;

  param.bd = xd->bd;

  param.is_hbd = is_cur_buf_hbd(xd);

  param.lossless = 0;

  param.tx_set_type = av1_get_ext_tx_set_type(

      param.tx_size, is_inter_block(xd->mi[0]), reduced_tx_set);

  const int bd_idx = (xd->bd == 8) ? 0 : ((xd->bd == 10) ? 1 : 2);

  const uint32_t max_qcoef_thresh = skip_pred_threshold[bd_idx][bsize];

  const int16_t *src_diff = x->plane[0].src_diff;

  const int n_coeff = tx_w * tx_h;

  const int16_t ac_q = av1_ac_quant_QTX(x->qindex, 0, xd->bd);

  const uint32_t dc_thresh = max_qcoef_thresh * dc_q;

  const uint32_t ac_thresh = max_qcoef_thresh * ac_q;

  for (int row = 0; row < bh; row += tx_h) {

    for (int col = 0; col < bw; col += tx_w) {

      av1_fwd_txfm(src_diff + col, coefs, bw, &param);

      // Operating on TX domain, not pixels; we want the QTX quantizers

      const uint32_t dc_coef = (((uint32_t)abs(coefs[0])) << 7);

      if (dc_coef >= dc_thresh) return 0;

      for (int i = 1; i < n_coeff; ++i) {

        const uint32_t ac_coef = (((uint32_t)abs(coefs[i])) << 7);

        if (ac_coef >= ac_thresh) return 0;

      }

    }

    src_diff += tx_h * bw;

  }

  return 1;

}

// Used to set proper context for early termination with skip = 1.

static AOM_INLINE void set_skip_txfm(MACROBLOCK *x, RD_STATS *rd_stats,

                                     int bsize, int64_t dist) {

  MACROBLOCKD *const xd = &x->e_mbd;

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

  const int n4 = bsize_to_num_blk(bsize);

  const TX_SIZE tx_size = max_txsize_rect_lookup[bsize];

  memset(xd->tx_type_map, DCT_DCT, sizeof(xd->tx_type_map[0]) * n4);

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

  mbmi->tx_size = tx_size;

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

    set_blk_skip(x->txfm_search_info.blk_skip, 0, i, 1);

  rd_stats->skip_txfm = 1;

  if (is_cur_buf_hbd(xd)) dist = ROUND_POWER_OF_TWO(dist, (xd->bd - 8) * 2);

  rd_stats->dist = rd_stats->sse = (dist << 4);

  // Though decision is to make the block as skip based on luma stats,

  // it is possible that block becomes non skip after chroma rd. In addition

  // intermediate non skip costs calculated by caller function will be

  // incorrect, if rate is set as  zero (i.e., if zero_blk_rate is not

  // accounted). Hence intermediate rate is populated to code the luma tx blks

  // as skip, the caller function based on final rd decision (i.e., skip vs

  // non-skip) sets the final rate accordingly. Here the rate populated

  // corresponds to coding all the tx blocks with zero_blk_rate (based on max tx

  // size possible) in the current block. Eg: For 128*128 block, rate would be

  // 4 * zero_blk_rate where zero_blk_rate corresponds to coding of one 64x64 tx

  // block as 'all zeros'

  ENTROPY_CONTEXT ctxa[MAX_MIB_SIZE];

  ENTROPY_CONTEXT ctxl[MAX_MIB_SIZE];

  av1_get_entropy_contexts(bsize, &xd->plane[0], ctxa, ctxl);

  ENTROPY_CONTEXT *ta = ctxa;

  ENTROPY_CONTEXT *tl = ctxl;

  const TX_SIZE txs_ctx = get_txsize_entropy_ctx(tx_size);

  TXB_CTX txb_ctx;

  get_txb_ctx(bsize, tx_size, 0, ta, tl, &txb_ctx);

  const int zero_blk_rate = x->coeff_costs.coeff_costs[txs_ctx][PLANE_TYPE_Y]

                                .txb_skip_cost[txb_ctx.txb_skip_ctx][1];

  rd_stats->rate = zero_blk_rate *

                   (block_size_wide[bsize] >> tx_size_wide_log2[tx_size]) *

                   (block_size_high[bsize] >> tx_size_high_log2[tx_size]);

}

static AOM_INLINE void save_mb_rd_info(int n4, uint32_t hash,

                                       const MACROBLOCK *const x,

                                       const RD_STATS *const rd_stats,

                                       MB_RD_RECORD *mb_rd_record) {

  int index;

  if (mb_rd_record->num < RD_RECORD_BUFFER_LEN) {

    index =

        (mb_rd_record->index_start + mb_rd_record->num) % RD_RECORD_BUFFER_LEN;

    ++mb_rd_record->num;

  } else {

    index = mb_rd_record->index_start;

    mb_rd_record->index_start =

        (mb_rd_record->index_start + 1) % RD_RECORD_BUFFER_LEN;

  }

  MB_RD_INFO *const mb_rd_info = &mb_rd_record->mb_rd_info[index];

  const MACROBLOCKD *const xd = &x->e_mbd;

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

  mb_rd_info->hash_value = hash;

  mb_rd_info->tx_size = mbmi->tx_size;

  memcpy(mb_rd_info->blk_skip, x->txfm_search_info.blk_skip,

         sizeof(mb_rd_info->blk_skip[0]) * n4);

  av1_copy(mb_rd_info->inter_tx_size, mbmi->inter_tx_size);

  av1_copy_array(mb_rd_info->tx_type_map, xd->tx_type_map, n4);

  mb_rd_info->rd_stats = *rd_stats;

}

static int get_search_init_depth(int mi_width, int mi_height, int is_inter,

                                 const SPEED_FEATURES *sf,

                                 int tx_size_search_method) {

  if (tx_size_search_method == USE_LARGESTALL) return MAX_VARTX_DEPTH;

  if (sf->tx_sf.tx_size_search_lgr_block) {

    if (mi_width > mi_size_wide[BLOCK_64X64] ||

        mi_height > mi_size_high[BLOCK_64X64])

      return MAX_VARTX_DEPTH;

  }

  if (is_inter) {

    return (mi_height != mi_width)

               ? sf->tx_sf.inter_tx_size_search_init_depth_rect

               : sf->tx_sf.inter_tx_size_search_init_depth_sqr;

  } else {

    return (mi_height != mi_width)

               ? sf->tx_sf.intra_tx_size_search_init_depth_rect

               : sf->tx_sf.intra_tx_size_search_init_depth_sqr;

  }

}

static AOM_INLINE void select_tx_block(

    const AV1_COMP *cpi, MACROBLOCK *x, int blk_row, int blk_col, int block,

    TX_SIZE tx_size, int depth, BLOCK_SIZE plane_bsize, ENTROPY_CONTEXT *ta,

    ENTROPY_CONTEXT *tl, TXFM_CONTEXT *tx_above, TXFM_CONTEXT *tx_left,

    RD_STATS *rd_stats, int64_t prev_level_rd, int64_t ref_best_rd,

    int *is_cost_valid, FAST_TX_SEARCH_MODE ftxs_mode);

// NOTE: CONFIG_COLLECT_RD_STATS has 3 possible values

// 0: Do not collect any RD stats

// 1: Collect RD stats for transform units

// 2: Collect RD stats for partition units

#if CONFIG_COLLECT_RD_STATS

static AOM_INLINE void get_energy_distribution_fine(

    const AV1_COMP *cpi, BLOCK_SIZE bsize, const uint8_t *src, int src_stride,

    const uint8_t *dst, int dst_stride, int need_4th, double *hordist,

    double *verdist) {

  const int bw = block_size_wide[bsize];

  const int bh = block_size_high[bsize];

  unsigned int esq[16] = { 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0 };

  if (bsize < BLOCK_16X16 || (bsize >= BLOCK_4X16 && bsize <= BLOCK_32X8)) {

    // Special cases: calculate 'esq' values manually, as we don't have 'vf'

    // functions for the 16 (very small) sub-blocks of this block.

    const int w_shift = (bw == 4) ? 0 : (bw == 8) ? 1 : (bw == 16) ? 2 : 3;

    const int h_shift = (bh == 4) ? 0 : (bh == 8) ? 1 : (bh == 16) ? 2 : 3;

    assert(bw <= 32);

    assert(bh <= 32);

    assert(((bw - 1) >> w_shift) + (((bh - 1) >> h_shift) << 2) == 15);

    if (cpi->common.seq_params->use_highbitdepth) {

      const uint16_t *src16 = CONVERT_TO_SHORTPTR(src);

      const uint16_t *dst16 = CONVERT_TO_SHORTPTR(dst);

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

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

          const int index = (j >> w_shift) + ((i >> h_shift) << 2);

          esq[index] +=

              (src16[j + i * src_stride] - dst16[j + i * dst_stride]) *

              (src16[j + i * src_stride] - dst16[j + i * dst_stride]);

        }

    } else {

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

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

          const int index = (j >> w_shift) + ((i >> h_shift) << 2);

          esq[index] += (src[j + i * src_stride] - dst[j + i * dst_stride]) *

                        (src[j + i * src_stride] - dst[j + i * dst_stride]);

        }

    }

  } else {  // Calculate 'esq' values using 'vf' functions on the 16 sub-blocks.

    const int f_index =

        (bsize < BLOCK_SIZES) ? bsize - BLOCK_16X16 : bsize - BLOCK_8X16;

    assert(f_index >= 0 && f_index < BLOCK_SIZES_ALL);

    const BLOCK_SIZE subsize = (BLOCK_SIZE)f_index;

    assert(block_size_wide[bsize] == 4 * block_size_wide[subsize]);

    assert(block_size_high[bsize] == 4 * block_size_high[subsize]);

    cpi->ppi->fn_ptr[subsize].vf(src, src_stride, dst, dst_stride, &esq[0]);

    cpi->ppi->fn_ptr[subsize].vf(src + bw / 4, src_stride, dst + bw / 4,

                                 dst_stride, &esq[1]);

    cpi->ppi->fn_ptr[subsize].vf(src + bw / 2, src_stride, dst + bw / 2,

                                 dst_stride, &esq[2]);

    cpi->ppi->fn_ptr[subsize].vf(src + 3 * bw / 4, src_stride, dst + 3 * bw / 4,

                                 dst_stride, &esq[3]);

    src += bh / 4 * src_stride;

    dst += bh / 4 * dst_stride;

    cpi->ppi->fn_ptr[subsize].vf(src, src_stride, dst, dst_stride, &esq[4]);

    cpi->ppi->fn_ptr[subsize].vf(src + bw / 4, src_stride, dst + bw / 4,

                                 dst_stride, &esq[5]);

    cpi->ppi->fn_ptr[subsize].vf(src + bw / 2, src_stride, dst + bw / 2,

                                 dst_stride, &esq[6]);

    cpi->ppi->fn_ptr[subsize].vf(src + 3 * bw / 4, src_stride, dst + 3 * bw / 4,

                                 dst_stride, &esq[7]);

    src += bh / 4 * src_stride;

    dst += bh / 4 * dst_stride;

    cpi->ppi->fn_ptr[subsize].vf(src, src_stride, dst, dst_stride, &esq[8]);

    cpi->ppi->fn_ptr[subsize].vf(src + bw / 4, src_stride, dst + bw / 4,

                                 dst_stride, &esq[9]);

    cpi->ppi->fn_ptr[subsize].vf(src + bw / 2, src_stride, dst + bw / 2,

                                 dst_stride, &esq[10]);

    cpi->ppi->fn_ptr[subsize].vf(src + 3 * bw / 4, src_stride, dst + 3 * bw / 4,

                                 dst_stride, &esq[11]);

    src += bh / 4 * src_stride;

    dst += bh / 4 * dst_stride;

    cpi->ppi->fn_ptr[subsize].vf(src, src_stride, dst, dst_stride, &esq[12]);

    cpi->ppi->fn_ptr[subsize].vf(src + bw / 4, src_stride, dst + bw / 4,

                                 dst_stride, &esq[13]);

    cpi->ppi->fn_ptr[subsize].vf(src + bw / 2, src_stride, dst + bw / 2,

                                 dst_stride, &esq[14]);

    cpi->ppi->fn_ptr[subsize].vf(src + 3 * bw / 4, src_stride, dst + 3 * bw / 4,

                                 dst_stride, &esq[15]);

  }

  double total = (double)esq[0] + esq[1] + esq[2] + esq[3] + esq[4] + esq[5] +

                 esq[6] + esq[7] + esq[8] + esq[9] + esq[10] + esq[11] +

                 esq[12] + esq[13] + esq[14] + esq[15];

  if (total > 0) {

    const double e_recip = 1.0 / total;

    hordist[0] = ((double)esq[0] + esq[4] + esq[8] + esq[12]) * e_recip;

    hordist[1] = ((double)esq[1] + esq[5] + esq[9] + esq[13]) * e_recip;

    hordist[2] = ((double)esq[2] + esq[6] + esq[10] + esq[14]) * e_recip;

    if (need_4th) {

      hordist[3] = ((double)esq[3] + esq[7] + esq[11] + esq[15]) * e_recip;

    }

    verdist[0] = ((double)esq[0] + esq[1] + esq[2] + esq[3]) * e_recip;

    verdist[1] = ((double)esq[4] + esq[5] + esq[6] + esq[7]) * e_recip;

    verdist[2] = ((double)esq[8] + esq[9] + esq[10] + esq[11]) * e_recip;

    if (need_4th) {

      verdist[3] = ((double)esq[12] + esq[13] + esq[14] + esq[15]) * e_recip;

    }

  } else {

    hordist[0] = verdist[0] = 0.25;

    hordist[1] = verdist[1] = 0.25;

    hordist[2] = verdist[2] = 0.25;

    if (need_4th) {

      hordist[3] = verdist[3] = 0.25;

    }

  }

}

static double get_sse_norm(const int16_t *diff, int stride, int w, int h) {

  double sum = 0.0;

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

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

      const int err = diff[j * stride + i];

      sum += err * err;

    }

  }

  assert(w > 0 && h > 0);

  return sum / (w * h);

}

static double get_sad_norm(const int16_t *diff, int stride, int w, int h) {

  double sum = 0.0;

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

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

      sum += abs(diff[j * stride + i]);

    }

  }

  assert(w > 0 && h > 0);

  return sum / (w * h);

}

static AOM_INLINE void get_2x2_normalized_sses_and_sads(

    const AV1_COMP *const cpi, BLOCK_SIZE tx_bsize, const uint8_t *const src,

    int src_stride, const uint8_t *const dst, int dst_stride,

    const int16_t *const src_diff, int diff_stride, double *const sse_norm_arr,

    double *const sad_norm_arr) {

  const BLOCK_SIZE tx_bsize_half =

      get_partition_subsize(tx_bsize, PARTITION_SPLIT);

  if (tx_bsize_half == BLOCK_INVALID) {  // manually calculate stats

    const int half_width = block_size_wide[tx_bsize] / 2;

    const int half_height = block_size_high[tx_bsize] / 2;

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

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

        const int16_t *const this_src_diff =

            src_diff + row * half_height * diff_stride + col * half_width;

        if (sse_norm_arr) {

          sse_norm_arr[row * 2 + col] =

              get_sse_norm(this_src_diff, diff_stride, half_width, half_height);

        }

        if (sad_norm_arr) {

          sad_norm_arr[row * 2 + col] =

              get_sad_norm(this_src_diff, diff_stride, half_width, half_height);

        }

      }

    }

  } else {  // use function pointers to calculate stats

    const int half_width = block_size_wide[tx_bsize_half];

    const int half_height = block_size_high[tx_bsize_half];

    const int num_samples_half = half_width * half_height;

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

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

        const uint8_t *const this_src =

            src + row * half_height * src_stride + col * half_width;

        const uint8_t *const this_dst =

            dst + row * half_height * dst_stride + col * half_width;

        if (sse_norm_arr) {

          unsigned int this_sse;

          cpi->ppi->fn_ptr[tx_bsize_half].vf(this_src, src_stride, this_dst,

                                             dst_stride, &this_sse);

          sse_norm_arr[row * 2 + col] = (double)this_sse / num_samples_half;

        }

        if (sad_norm_arr) {

          const unsigned int this_sad = cpi->ppi->fn_ptr[tx_bsize_half].sdf(

              this_src, src_stride, this_dst, dst_stride);

          sad_norm_arr[row * 2 + col] = (double)this_sad / num_samples_half;

        }

      }

    }

  }

}

#if CONFIG_COLLECT_RD_STATS == 1

static double get_mean(const int16_t *diff, int stride, int w, int h) {

  double sum = 0.0;

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

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

      sum += diff[j * stride + i];

    }

  }

  assert(w > 0 && h > 0);

  return sum / (w * h);

}

static AOM_INLINE void PrintTransformUnitStats(

    const AV1_COMP *const cpi, MACROBLOCK *x, const RD_STATS *const rd_stats,

    int blk_row, int blk_col, BLOCK_SIZE plane_bsize, TX_SIZE tx_size,

    TX_TYPE tx_type, int64_t rd) {

  if (rd_stats->rate == INT_MAX || rd_stats->dist == INT64_MAX) return;

  // Generate small sample to restrict output size.

  static unsigned int seed = 21743;

  if (lcg_rand16(&seed) % 256 > 0) return;

  const char output_file[] = "tu_stats.txt";

  FILE *fout = fopen(output_file, "a");

  if (!fout) return;

  const BLOCK_SIZE tx_bsize = txsize_to_bsize[tx_size];

  const MACROBLOCKD *const xd = &x->e_mbd;

  const int plane = 0;

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

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

  const int txw = tx_size_wide[tx_size];

  const int txh = tx_size_high[tx_size];

  const int dequant_shift = (is_cur_buf_hbd(xd)) ? xd->bd - 5 : 3;

  const int q_step = p->dequant_QTX[1] >> dequant_shift;

  const int num_samples = txw * txh;

  const double rate_norm = (double)rd_stats->rate / num_samples;

  const double dist_norm = (double)rd_stats->dist / num_samples;

  fprintf(fout, "%g %g", rate_norm, dist_norm);

  const int src_stride = p->src.stride;

  const uint8_t *const src =

      &p->src.buf[(blk_row * src_stride + blk_col) << MI_SIZE_LOG2];

  const int dst_stride = pd->dst.stride;

  const uint8_t *const dst =

      &pd->dst.buf[(blk_row * dst_stride + blk_col) << MI_SIZE_LOG2];

  unsigned int sse;

  cpi->ppi->fn_ptr[tx_bsize].vf(src, src_stride, dst, dst_stride, &sse);

  const double sse_norm = (double)sse / num_samples;

  const unsigned int sad =

      cpi->ppi->fn_ptr[tx_bsize].sdf(src, src_stride, dst, dst_stride);

  const double sad_norm = (double)sad / num_samples;

  fprintf(fout, " %g %g", sse_norm, sad_norm);

  const int diff_stride = block_size_wide[plane_bsize];

  const int16_t *const src_diff =

      &p->src_diff[(blk_row * diff_stride + blk_col) << MI_SIZE_LOG2];

  double sse_norm_arr[4], sad_norm_arr[4];

  get_2x2_normalized_sses_and_sads(cpi, tx_bsize, src, src_stride, dst,

                                   dst_stride, src_diff, diff_stride,

                                   sse_norm_arr, sad_norm_arr);

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

    fprintf(fout, " %g", sse_norm_arr[i]);

  }

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

    fprintf(fout, " %g", sad_norm_arr[i]);

  }

  const TX_TYPE_1D tx_type_1d_row = htx_tab[tx_type];

  const TX_TYPE_1D tx_type_1d_col = vtx_tab[tx_type];

  fprintf(fout, " %d %d %d %d %d", q_step, tx_size_wide[tx_size],

          tx_size_high[tx_size], tx_type_1d_row, tx_type_1d_col);

  int model_rate;

  int64_t model_dist;

  model_rd_sse_fn[MODELRD_CURVFIT](cpi, x, tx_bsize, plane, sse, num_samples,

                                   &model_rate, &model_dist);

  const double model_rate_norm = (double)model_rate / num_samples;

  const double model_dist_norm = (double)model_dist / num_samples;

  fprintf(fout, " %g %g", model_rate_norm, model_dist_norm);

  const double mean = get_mean(src_diff, diff_stride, txw, txh);

  float hor_corr, vert_corr;

  av1_get_horver_correlation_full(src_diff, diff_stride, txw, txh, &hor_corr,

                                  &vert_corr);

  fprintf(fout, " %g %g %g", mean, hor_corr, vert_corr);

  double hdist[4] = { 0 }, vdist[4] = { 0 };

  get_energy_distribution_fine(cpi, tx_bsize, src, src_stride, dst, dst_stride,

                               1, hdist, vdist);

  fprintf(fout, " %g %g %g %g %g %g %g %g", hdist[0], hdist[1], hdist[2],

          hdist[3], vdist[0], vdist[1], vdist[2], vdist[3]);

  fprintf(fout, " %d %" PRId64, x->rdmult, rd);

  fprintf(fout, "\n");

  fclose(fout);

}

#endif  // CONFIG_COLLECT_RD_STATS == 1

#if CONFIG_COLLECT_RD_STATS >= 2

static int64_t get_sse(const AV1_COMP *cpi, const MACROBLOCK *x) {

  const AV1_COMMON *cm = &cpi->common;

  const int num_planes = av1_num_planes(cm);

  const MACROBLOCKD *xd = &x->e_mbd;

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

  int64_t total_sse = 0;

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

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

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

    const BLOCK_SIZE bs =

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

    unsigned int sse;

    if (x->skip_chroma_rd && plane) continue;

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

                            pd->dst.stride, &sse);

    total_sse += sse;

  }

  total_sse <<= 4;

  return total_sse;

}

static int get_est_rate_dist(const TileDataEnc *tile_data, BLOCK_SIZE bsize,

                             int64_t sse, int *est_residue_cost,

                             int64_t *est_dist) {

  const InterModeRdModel *md = &tile_data->inter_mode_rd_models[bsize];

  if (md->ready) {

    if (sse < md->dist_mean) {

      *est_residue_cost = 0;

      *est_dist = sse;

    } else {

      *est_dist = (int64_t)round(md->dist_mean);

      const double est_ld = md->a * sse + md->b;

      // Clamp estimated rate cost by INT_MAX / 2.

      // TODO(angiebird@google.com): find better solution than clamping.

      if (fabs(est_ld) < 1e-2) {

        *est_residue_cost = INT_MAX / 2;

      } else {

        double est_residue_cost_dbl = ((sse - md->dist_mean) / est_ld);

        if (est_residue_cost_dbl < 0) {

          *est_residue_cost = 0;

        } else {

          *est_residue_cost =

              (int)AOMMIN((int64_t)round(est_residue_cost_dbl), INT_MAX / 2);

        }

      }

      if (*est_residue_cost <= 0) {

        *est_residue_cost = 0;

        *est_dist = sse;

      }

    }

    return 1;

  }

  return 0;

}

static double get_highbd_diff_mean(const uint8_t *src8, int src_stride,

                                   const uint8_t *dst8, int dst_stride, int w,

                                   int h) {

  const uint16_t *src = CONVERT_TO_SHORTPTR(src8);

  const uint16_t *dst = CONVERT_TO_SHORTPTR(dst8);

  double sum = 0.0;

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

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

      const int diff = src[j * src_stride + i] - dst[j * dst_stride + i];

      sum += diff;

    }

  }

  assert(w > 0 && h > 0);

  return sum / (w * h);

}

static double get_diff_mean(const uint8_t *src, int src_stride,

                            const uint8_t *dst, int dst_stride, int w, int h) {

  double sum = 0.0;

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

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

      const int diff = src[j * src_stride + i] - dst[j * dst_stride + i];

      sum += diff;

    }

  }

  assert(w > 0 && h > 0);

  return sum / (w * h);

}

static AOM_INLINE void PrintPredictionUnitStats(const AV1_COMP *const cpi,

                                                const TileDataEnc *tile_data,

                                                MACROBLOCK *x,

                                                const RD_STATS *const rd_stats,

                                                BLOCK_SIZE plane_bsize) {

  if (rd_stats->rate == INT_MAX || rd_stats->dist == INT64_MAX) return;

  if (cpi->sf.inter_sf.inter_mode_rd_model_estimation == 1 &&

      (tile_data == NULL ||

       !tile_data->inter_mode_rd_models[plane_bsize].ready))

    return;

  (void)tile_data;

  // Generate small sample to restrict output size.

  static unsigned int seed = 95014;

  if ((lcg_rand16(&seed) % (1 << (14 - num_pels_log2_lookup[plane_bsize]))) !=

      1)

    return;

  const char output_file[] = "pu_stats.txt";

  FILE *fout = fopen(output_file, "a");

  if (!fout) return;

  MACROBLOCKD *const xd = &x->e_mbd;

  const int plane = 0;

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

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

  const int diff_stride = block_size_wide[plane_bsize];

  int bw, bh;

  get_txb_dimensions(xd, plane, plane_bsize, 0, 0, plane_bsize, NULL, NULL, &bw,

                     &bh);

  const int num_samples = bw * bh;

  const int dequant_shift = (is_cur_buf_hbd(xd)) ? xd->bd - 5 : 3;

  const int q_step = p->dequant_QTX[1] >> dequant_shift;

  const int shift = (xd->bd - 8);

  const double rate_norm = (double)rd_stats->rate / num_samples;

  const double dist_norm = (double)rd_stats->dist / num_samples;

  const double rdcost_norm =

      (double)RDCOST(x->rdmult, rd_stats->rate, rd_stats->dist) / num_samples;

  fprintf(fout, "%g %g %g", rate_norm, dist_norm, rdcost_norm);

  const int src_stride = p->src.stride;

  const uint8_t *const src = p->src.buf;

  const int dst_stride = pd->dst.stride;

  const uint8_t *const dst = pd->dst.buf;

  const int16_t *const src_diff = p->src_diff;

  int64_t sse = calculate_sse(xd, p, pd, bw, bh);

  const double sse_norm = (double)sse / num_samples;

  const unsigned int sad =

      cpi->ppi->fn_ptr[plane_bsize].sdf(src, src_stride, dst, dst_stride);

  const double sad_norm =

      (double)sad / (1 << num_pels_log2_lookup[plane_bsize]);

  fprintf(fout, " %g %g", sse_norm, sad_norm);

  double sse_norm_arr[4], sad_norm_arr[4];

  get_2x2_normalized_sses_and_sads(cpi, plane_bsize, src, src_stride, dst,

                                   dst_stride, src_diff, diff_stride,

                                   sse_norm_arr, sad_norm_arr);

  if (shift) {

    for (int k = 0; k < 4; ++k) sse_norm_arr[k] /= (1 << (2 * shift));

    for (int k = 0; k < 4; ++k) sad_norm_arr[k] /= (1 << shift);

  }

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

    fprintf(fout, " %g", sse_norm_arr[i]);

  }

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

    fprintf(fout, " %g", sad_norm_arr[i]);

  }

  fprintf(fout, " %d %d %d %d", q_step, x->rdmult, bw, bh);

  int model_rate;

  int64_t model_dist;

  model_rd_sse_fn[MODELRD_CURVFIT](cpi, x, plane_bsize, plane, sse, num_samples,

                                   &model_rate, &model_dist);

  const double model_rdcost_norm =

      (double)RDCOST(x->rdmult, model_rate, model_dist) / num_samples;

  const double model_rate_norm = (double)model_rate / num_samples;

  const double model_dist_norm = (double)model_dist / num_samples;

  fprintf(fout, " %g %g %g", model_rate_norm, model_dist_norm,

          model_rdcost_norm);

  double mean;

  if (is_cur_buf_hbd(xd)) {

    mean = get_highbd_diff_mean(p->src.buf, p->src.stride, pd->dst.buf,

                                pd->dst.stride, bw, bh);

  } else {

    mean = get_diff_mean(p->src.buf, p->src.stride, pd->dst.buf, pd->dst.stride,

                         bw, bh);

  }

  mean /= (1 << shift);

  float hor_corr, vert_corr;

  av1_get_horver_correlation_full(src_diff, diff_stride, bw, bh, &hor_corr,

                                  &vert_corr);

  fprintf(fout, " %g %g %g", mean, hor_corr, vert_corr);

  double hdist[4] = { 0 }, vdist[4] = { 0 };

  get_energy_distribution_fine(cpi, plane_bsize, src, src_stride, dst,

                               dst_stride, 1, hdist, vdist);

  fprintf(fout, " %g %g %g %g %g %g %g %g", hdist[0], hdist[1], hdist[2],

          hdist[3], vdist[0], vdist[1], vdist[2], vdist[3]);

  if (cpi->sf.inter_sf.inter_mode_rd_model_estimation == 1) {

    assert(tile_data->inter_mode_rd_models[plane_bsize].ready);

    const int64_t overall_sse = get_sse(cpi, x);

    int est_residue_cost = 0;

    int64_t est_dist = 0;

    get_est_rate_dist(tile_data, plane_bsize, overall_sse, &est_residue_cost,

                      &est_dist);

    const double est_residue_cost_norm = (double)est_residue_cost / num_samples;

    const double est_dist_norm = (double)est_dist / num_samples;

    const double est_rdcost_norm =

        (double)RDCOST(x->rdmult, est_residue_cost, est_dist) / num_samples;

    fprintf(fout, " %g %g %g", est_residue_cost_norm, est_dist_norm,

            est_rdcost_norm);

  }

  fprintf(fout, "\n");

  fclose(fout);

}

#endif  // CONFIG_COLLECT_RD_STATS >= 2

#endif  // CONFIG_COLLECT_RD_STATS

static AOM_INLINE void inverse_transform_block_facade(MACROBLOCK *const x,

                                                      int plane, int block,

                                                      int blk_row, int blk_col,

                                                      int eob,

                                                      int reduced_tx_set) {

  if (!eob) return;

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

  MACROBLOCKD *const xd = &x->e_mbd;

  tran_low_t *dqcoeff = p->dqcoeff + BLOCK_OFFSET(block);

  const PLANE_TYPE plane_type = get_plane_type(plane);

  const TX_SIZE tx_size = av1_get_tx_size(plane, xd);

  const TX_TYPE tx_type = av1_get_tx_type(xd, plane_type, blk_row, blk_col,

                                          tx_size, reduced_tx_set);

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

  const int dst_stride = pd->dst.stride;

  uint8_t *dst = &pd->dst.buf[(blk_row * dst_stride + blk_col) << MI_SIZE_LOG2];

  av1_inverse_transform_block(xd, dqcoeff, plane, tx_type, tx_size, dst,

                              dst_stride, eob, reduced_tx_set);

}

static INLINE void recon_intra(const AV1_COMP *cpi, MACROBLOCK *x, int plane,

                               int block, int blk_row, int blk_col,

                               BLOCK_SIZE plane_bsize, TX_SIZE tx_size,

                               const TXB_CTX *const txb_ctx, int skip_trellis,

                               TX_TYPE best_tx_type, int do_quant,

                               int *rate_cost, uint16_t best_eob) {

  const AV1_COMMON *cm = &cpi->common;

  MACROBLOCKD *xd = &x->e_mbd;

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

  const int is_inter = is_inter_block(mbmi);

  if (!is_inter && best_eob &&

      (blk_row + tx_size_high_unit[tx_size] < mi_size_high[plane_bsize] ||

       blk_col + tx_size_wide_unit[tx_size] < mi_size_wide[plane_bsize])) {

    // if the quantized coefficients are stored in the dqcoeff buffer, we don't

    // need to do transform and quantization again.

    if (do_quant) {

      TxfmParam txfm_param_intra;

      QUANT_PARAM quant_param_intra;

      av1_setup_xform(cm, x, tx_size, best_tx_type, &txfm_param_intra);

      av1_setup_quant(tx_size, !skip_trellis,

                      skip_trellis

                          ? (USE_B_QUANT_NO_TRELLIS ? AV1_XFORM_QUANT_B

                                                    : AV1_XFORM_QUANT_FP)

                          : AV1_XFORM_QUANT_FP,

                      cpi->oxcf.q_cfg.quant_b_adapt, &quant_param_intra);

      av1_setup_qmatrix(&cm->quant_params, xd, plane, tx_size, best_tx_type,

                        &quant_param_intra);

      av1_xform_quant(x, plane, block, blk_row, blk_col, plane_bsize,

                      &txfm_param_intra, &quant_param_intra);

      if (quant_param_intra.use_optimize_b) {

        av1_optimize_b(cpi, x, plane, block, tx_size, best_tx_type, txb_ctx,

                       rate_cost);

      }

    }

    inverse_transform_block_facade(x, plane, block, blk_row, blk_col,

                                   x->plane[plane].eobs[block],

                                   cm->features.reduced_tx_set_used);

    // This may happen because of hash collision. The eob stored in the hash

    // table is non-zero, but the real eob is zero. We need to make sure tx_type

    // is DCT_DCT in this case.

    if (plane == 0 && x->plane[plane].eobs[block] == 0 &&

        best_tx_type != DCT_DCT) {

      update_txk_array(xd, blk_row, blk_col, tx_size, DCT_DCT);

    }

  }

}

static unsigned pixel_dist_visible_only(

    const AV1_COMP *const cpi, const MACROBLOCK *x, const uint8_t *src,

    const int src_stride, const uint8_t *dst, const int dst_stride,

    const BLOCK_SIZE tx_bsize, int txb_rows, int txb_cols, int visible_rows,

    int visible_cols) {

  unsigned sse;

  if (txb_rows == visible_rows && txb_cols == visible_cols) {

    cpi->ppi->fn_ptr[tx_bsize].vf(src, src_stride, dst, dst_stride, &sse);

    return sse;

  }

#if CONFIG_AV1_HIGHBITDEPTH

  const MACROBLOCKD *xd = &x->e_mbd;

  if (is_cur_buf_hbd(xd)) {

    uint64_t sse64 = aom_highbd_sse_odd_size(src, src_stride, dst, dst_stride,

                                             visible_cols, visible_rows);

    return (unsigned int)ROUND_POWER_OF_TWO(sse64, (xd->bd - 8) * 2);

  }

#else

  (void)x;

#endif

  sse = aom_sse_odd_size(src, src_stride, dst, dst_stride, visible_cols,

                         visible_rows);

  return sse;

}

// Compute the pixel domain distortion from src and dst on all visible 4x4s in

// the

// transform block.

static unsigned pixel_dist(const AV1_COMP *const cpi, const MACROBLOCK *x,

                           int plane, const uint8_t *src, const int src_stride,

                           const uint8_t *dst, const int dst_stride,

                           int blk_row, int blk_col,

                           const BLOCK_SIZE plane_bsize,

                           const BLOCK_SIZE tx_bsize) {

  int txb_rows, txb_cols, visible_rows, visible_cols;

  const MACROBLOCKD *xd = &x->e_mbd;

  get_txb_dimensions(xd, plane, plane_bsize, blk_row, blk_col, tx_bsize,

                     &txb_cols, &txb_rows, &visible_cols, &visible_rows);

  assert(visible_rows > 0);

  assert(visible_cols > 0);

  unsigned sse = pixel_dist_visible_only(cpi, x, src, src_stride, dst,

                                         dst_stride, tx_bsize, txb_rows,

                                         txb_cols, visible_rows, visible_cols);

  return sse;

}

static INLINE int64_t dist_block_px_domain(const AV1_COMP *cpi, MACROBLOCK *x,

                                           int plane, BLOCK_SIZE plane_bsize,

                                           int block, int blk_row, int blk_col,

                                           TX_SIZE tx_size) {

  MACROBLOCKD *const xd = &x->e_mbd;

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

  const uint16_t eob = p->eobs[block];

  const BLOCK_SIZE tx_bsize = txsize_to_bsize[tx_size];

  const int bsw = block_size_wide[tx_bsize];

  const int bsh = block_size_high[tx_bsize];

  const int src_stride = x->plane[plane].src.stride;

  const int dst_stride = xd->plane[plane].dst.stride;

  // Scale the transform block index to pixel unit.

  const int src_idx = (blk_row * src_stride + blk_col) << MI_SIZE_LOG2;

  const int dst_idx = (blk_row * dst_stride + blk_col) << MI_SIZE_LOG2;

  const uint8_t *src = &x->plane[plane].src.buf[src_idx];

  const uint8_t *dst = &xd->plane[plane].dst.buf[dst_idx];

  const tran_low_t *dqcoeff = p->dqcoeff + BLOCK_OFFSET(block);

  assert(cpi != NULL);

  assert(tx_size_wide_log2[0] == tx_size_high_log2[0]);

  uint8_t *recon;

  DECLARE_ALIGNED(16, uint16_t, recon16[MAX_TX_SQUARE]);

#if CONFIG_AV1_HIGHBITDEPTH

  if (is_cur_buf_hbd(xd)) {

    recon = CONVERT_TO_BYTEPTR(recon16);

    aom_highbd_convolve_copy(CONVERT_TO_SHORTPTR(dst), dst_stride,

                             CONVERT_TO_SHORTPTR(recon), MAX_TX_SIZE, bsw, bsh);

  } else {

    recon = (uint8_t *)recon16;

    aom_convolve_copy(dst, dst_stride, recon, MAX_TX_SIZE, bsw, bsh);

  }

#else

  recon = (uint8_t *)recon16;

  aom_convolve_copy(dst, dst_stride, recon, MAX_TX_SIZE, bsw, bsh);

#endif

  const PLANE_TYPE plane_type = get_plane_type(plane);

  TX_TYPE tx_type = av1_get_tx_type(xd, plane_type, blk_row, blk_col, tx_size,

                                    cpi->common.features.reduced_tx_set_used);

  av1_inverse_transform_block(xd, dqcoeff, plane, tx_type, tx_size, recon,

                              MAX_TX_SIZE, eob,

                              cpi->common.features.reduced_tx_set_used);

  return 16 * pixel_dist(cpi, x, plane, src, src_stride, recon, MAX_TX_SIZE,

                         blk_row, blk_col, plane_bsize, tx_bsize);

}

// pruning thresholds for prune_txk_type and prune_txk_type_separ

static const int prune_factors[5] = { 200, 200, 120, 80, 40 };  // scale 1000

static const int mul_factors[5] = { 80, 80, 70, 50, 30 };       // scale 100

// R-D costs are sorted in ascending order.

static INLINE void sort_rd(int64_t rds[], int txk[], int len) {

  int i, j, k;

  for (i = 1; i <= len - 1; ++i) {

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

      if (rds[j] > rds[i]) {

        int64_t temprd;

        int tempi;

        temprd = rds[i];

        tempi = txk[i];

        for (k = i; k > j; k--) {

          rds[k] = rds[k - 1];

          txk[k] = txk[k - 1];

        }

        rds[j] = temprd;

        txk[j] = tempi;

        break;