/*

 * Copyright (c) 2021, Alliance for Open Media. All rights reserved

 *

 * This source code is subject to the terms of the BSD 3-Clause Clear License

 * and the Alliance for Open Media Patent License 1.0. If the BSD 3-Clause Clear

 * License was not distributed with this source code in the LICENSE file, you

 * can obtain it at aomedia.org/license/software-license/bsd-3-c-c/.  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

 * aomedia.org/license/patent-license/.

 */

#include "av2/decoder/decodetxb.h"

#include "avm_ports/mem.h"

#include "av2/common/hr_coding.h"

#include "av2/common/idct.h"

#include "av2/common/pred_common.h"

#include "av2/common/reconintra.h"

#include "av2/common/scan.h"

#include "av2/common/txb_common.h"

#include "av2/decoder/decodemv.h"

#if CONFIG_PARAKIT_COLLECT_DATA

#include "av2/common/cost.h"

#endif

#if CONFIG_PARAKIT_COLLECT_DATA

static int get_q_ctx(int q) {

  if (q <= 90) return 0;

  if (q <= 140) return 1;

  if (q <= 190) return 2;

  return 3;

}

#endif

/*!\brief Read and decode from bitstream an integer value following Exp-Golomb

 * coding with order k

 *

 * \ingroup coefficient_coding

 *

 * This function reads and decodes from the bitstream r an integer value which

 * was coded in Exp-Golomb with order k

 *

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

 *                         current macroblockd

 * \param[in]    r         Pointer to the bitstream reader

 * \param[in]    k         Order of Exp-Golomb coding

 *

 */

static int read_exp_golomb(MACROBLOCKD *xd, avm_reader *r, int k) {

  int length = avm_read_unary(r, 21, ACCT_INFO("hr"));

  if (length > 20) {

    avm_internal_error(xd->error_info, AVM_CODEC_CORRUPT_FRAME,

                       "Invalid length in read_exp_golomb");

  }

  length += k;

  int x = 1 << length;

  x += avm_read_literal(r, length, ACCT_INFO("hr"));

  return x - (1 << k);

}

/*!\brief Read and decode from the bitstream a Truncated-Rice coded integer

 * value.

 *

 * \ingroup coefficient_coding

 *

 * This function reads and decodes from the bitstream a Truncated-Rice coded

 * integer value, with Rice parameter m. It first reads and decodes the unary

 * prefix q from the input bitstream r. Given a shorter prefix (q < cmax), it

 * decodes the remaining offset of the integer using Golomb-Rice coding;

 * otherwise (i.e., when q == cmax) it shifts to use Exp-Golomb to decode the

 * remainder value.

 *

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

 *                         current macroblockd

 * \param[in]    r         Pointer to the reader of the bitstream

 * \param[in]    m         Parameter of the Rice distribution

 * \param[in]    k         Order of the Exp-Golomb code

 * \param[in]    cmax      Maximum unary prefix length above which Exp-Golomb

 *                         is used instead of Golomb-Rice coding

 *

 */

static int read_truncated_rice(MACROBLOCKD *xd, avm_reader *r, int m, int k,

                               int cmax) {

  int q = avm_read_unary(r, cmax, ACCT_INFO("hr"));

  int rem = (q == cmax) ? read_exp_golomb(xd, r, k)

                        : avm_read_literal(r, m, ACCT_INFO("hr"));

  return rem + (q << m);

}

/*!\brief Read and decode from the bitstream the high range (HR) value of

 * the coefficient level using adaptive Truncated-Rice coding.

 *

 * \ingroup coefficient_coding

 *

 * This function reads and decodes from the bitstream the high range (HR)

 * value of the coefficient level following adaptive Truncated-Rice coding.

 * It first derives the Rice parameter m from the input context value ctx.

 * It then invokes the read_truncated_rice function to decode the HR

 * coefficient value using Rice parameter m, Exp-Golomb order k = m + 1,

 * and maximum unary prefix lenght cmax = AVMMIN(m + 4, 6).

 *

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

 *                         current macroblockd

 * \param[in]    r         Pointer to the reader of the bitstream

 * \param[in]    ctx       Context value

 *

 */

static int read_adaptive_hr(MACROBLOCKD *xd, avm_reader *r, int ctx) {

  int m = get_adaptive_param(ctx);

  return read_truncated_rice(xd, r, m, m + 1, AVMMIN(m + 4, 6));

}

// Read high range part of coeff

static INLINE int read_high_range(MACROBLOCKD *xd, avm_reader *r, int tcq_mode,

                                  int level, int lf, int *hr_avg, int plane) {

  int max_br = lf ? (plane == 0 ? LF_MAX_BASE_BR_RANGE  // 8

                                : LF_NUM_BASE_LEVELS + 1)

                  : MAX_BASE_BR_RANGE;

  int use_tcq_hr = tcq_mode && (level >= max_br - 1);

  int hr_level_avg = *hr_avg;

  int use_hr = use_tcq_hr || level >= max_br;

  if (use_hr) {

    int hr = read_adaptive_hr(xd, r, hr_level_avg);

    level += hr << (tcq_mode ? 1 : 0);

    *hr_avg = (hr_level_avg + hr) >> 1;

  }

  return level;

}

static INLINE int rec_eob_pos(const int eob_token, const int extra) {

  int eob = av2_eob_group_start[eob_token];

  if (eob > 2) {

    eob += extra;

  }

  return eob;

}

static INLINE int get_dqv(const int32_t *dequant, int coeff_idx,

                          const qm_val_t *iqmatrix) {

  int dqv = dequant[!!coeff_idx];

  if (iqmatrix != NULL)

    dqv =

        ((iqmatrix[coeff_idx] * dqv) + (1 << (AVM_QM_BITS - 1))) >> AVM_QM_BITS;

  return dqv;

}

static int read_low_range(avm_reader *r, avm_cdf_prob *cdf) {

  int br_level = 0;

  for (int idx = 0; idx < COEFF_BASE_RANGE; idx += BR_CDF_SIZE - 1) {

    const int k = avm_read_symbol(r, cdf, BR_CDF_SIZE,

                                  ACCT_INFO("k", "read_low_range cdf"));

    br_level += k;

    if (k < BR_CDF_SIZE - 1) break;

  }

  return br_level;

}

static INLINE void read_coeffs_reverse_2d(

    avm_reader *r, int start_si, int end_si, const int16_t *scan, int bwl,

    uint8_t *levels, base_lf_cdf_arr base_lf_cdf, br_cdf_arr br_lf_cdf,

    int plane, base_cdf_arr base_cdf, br_cdf_arr br_cdf,

    base_lf_uv_cdf_arr base_lf_uv_cdf, base_uv_cdf_arr base_uv_cdf,

    br_cdf_arr br_uv_cdf, int *state) {

  for (int c = end_si; c >= start_si; --c) {

    const int pos = scan[c];

    int level = 0;

    const int row = pos >> bwl;

    const int col = pos - (row << bwl);

    int limits = get_lf_limits(row, col, 0, plane);

    int q_i = tcq_quant(*state);

    if (plane > 0) {

      if (limits) {

        const int coeff_ctx =

            get_lower_levels_ctx_lf_2d_chroma(levels, pos, bwl, plane);

        level += avm_read_symbol(r, base_lf_uv_cdf[coeff_ctx], LF_BASE_SYMBOLS,

                                 ACCT_INFO("level", "base_lf_uv_cdf"));

      } else {

        const int coeff_ctx =

            get_lower_levels_ctx_2d_chroma(levels, pos, bwl, plane);

        level += avm_read_symbol(r, base_uv_cdf[coeff_ctx], 4,

                                 ACCT_INFO("level", "base_uv_cdf"));

        if (level > NUM_BASE_LEVELS) {

          const int br_ctx = get_br_ctx_2d_chroma(levels, pos, bwl);

          avm_cdf_prob *cdf = br_uv_cdf[br_ctx];

          level += read_low_range(r, cdf);

        }

      }

    } else {

      if (limits) {

        const int coeff_ctx = get_lower_levels_ctx_lf_2d(levels, pos, bwl);

        level +=

            avm_read_symbol(r, base_lf_cdf[coeff_ctx][q_i], LF_BASE_SYMBOLS,

                            ACCT_INFO("level", "base_lf_cdf"));

        if (level > LF_NUM_BASE_LEVELS) {

          const int br_ctx = get_br_lf_ctx_2d(levels, pos, bwl);

          avm_cdf_prob *cdf = br_lf_cdf[br_ctx];

          level += read_low_range(r, cdf);

        }

      } else {

        const int coeff_ctx = get_lower_levels_ctx_2d(levels, pos, bwl, plane);

        level += avm_read_symbol(r, base_cdf[coeff_ctx][q_i], 4,

                                 ACCT_INFO("level", "base_cdf"));

        if (level > NUM_BASE_LEVELS) {

          const int br_ctx = get_br_ctx_2d(levels, pos, bwl);

          avm_cdf_prob *cdf = br_cdf[br_ctx];

          level += read_low_range(r, cdf);

        }

      }

    }

    levels[get_padded_idx(pos, bwl)] = level;

    *state = tcq_next_state(*state, level);

  }

}

static INLINE void read_coeffs_reverse(

    avm_reader *r, TX_CLASS tx_class, int start_si, int end_si,

    const int16_t *scan, int bwl, uint8_t *levels, base_lf_cdf_arr base_lf_cdf,

    br_cdf_arr br_lf_cdf, int plane, base_cdf_arr base_cdf, br_cdf_arr br_cdf,

    base_lf_uv_cdf_arr base_lf_uv_cdf, base_uv_cdf_arr base_uv_cdf,

    br_cdf_arr br_uv_cdf, int *state) {

  for (int c = end_si; c >= start_si; --c) {

    const int pos = scan[c];

    int level = 0;

    int q_i = tcq_quant(*state);

    const int row = pos >> bwl;

    const int col = pos - (row << bwl);

    int limits = get_lf_limits(row, col, tx_class, plane);

    if (plane > 0) {

      if (limits) {

        const int coeff_ctx =

            get_lower_levels_lf_ctx_chroma(levels, pos, bwl, tx_class, plane);

        level += avm_read_symbol(r, base_lf_uv_cdf[coeff_ctx], LF_BASE_SYMBOLS,

                                 ACCT_INFO("level", "base_lf_uv_cdf"));

      } else {

        const int coeff_ctx =

            get_lower_levels_ctx_chroma(levels, pos, bwl, tx_class, plane);

        level += avm_read_symbol(r, base_uv_cdf[coeff_ctx], 4,

                                 ACCT_INFO("level", "base_uv_cdf"));

        if (level > NUM_BASE_LEVELS) {

          const int br_ctx = get_br_ctx_chroma(levels, pos, bwl, tx_class);

          avm_cdf_prob *cdf = br_uv_cdf[br_ctx];

          level += read_low_range(r, cdf);

        }

      }

    } else {

      if (limits) {

        const int coeff_ctx =

            get_lower_levels_lf_ctx(levels, pos, bwl, tx_class);

        level +=

            avm_read_symbol(r, base_lf_cdf[coeff_ctx][q_i], LF_BASE_SYMBOLS,

                            ACCT_INFO("level", "base_lf_cdf"));

        if (level > LF_NUM_BASE_LEVELS) {

          const int br_ctx = get_br_lf_ctx(levels, pos, bwl, tx_class);

          avm_cdf_prob *cdf = br_lf_cdf[br_ctx];

          level += read_low_range(r, cdf);

        }

      } else {

        const int coeff_ctx =

            get_lower_levels_ctx(levels, pos, bwl, tx_class, plane);

        level += avm_read_symbol(r, base_cdf[coeff_ctx][q_i], 4,

                                 ACCT_INFO("level", "base_cdf"));

        if (level > NUM_BASE_LEVELS) {

          const int br_ctx = get_br_ctx(levels, pos, bwl, tx_class);

          avm_cdf_prob *cdf = br_cdf[br_ctx];

          level += read_low_range(r, cdf);

        }

      }

    }

    levels[get_padded_idx(pos, bwl)] = level;

    *state = tcq_next_state(*state, level);

  }

}

static INLINE void read_coeffs_forward_2d(avm_reader *r, int start_si,

                                          int end_si, const int16_t *scan,

                                          int bwl, uint8_t *levels,

                                          base_fsc_cdf_arr base_cdf,

                                          br_cdf_arr br_cdf) {

  for (int c = start_si; c <= end_si; c++) {

    const int pos = scan[c];

    const int coeff_ctx = get_upper_levels_ctx_2d(levels, pos, bwl);

    const int nsymbs = 4;

    int level = avm_read_symbol(r, base_cdf[coeff_ctx], nsymbs,

                                ACCT_INFO("level", "base_cdf"));

    if (level > NUM_BASE_LEVELS) {

      const int br_ctx = get_br_ctx_skip(levels, pos, bwl);

      avm_cdf_prob *cdf = br_cdf[br_ctx];

      level += read_low_range(r, cdf);

    }

    levels[get_padded_idx_left(pos, bwl)] = level;

  }

}

// Decode the end-of-block syntax.

static INLINE void decode_eob(DecoderCodingBlock *dcb, avm_reader *const r,

                              const int plane, const TX_SIZE tx_size

#if CONFIG_PARAKIT_COLLECT_DATA

                              ,

                              const AV2_COMMON *const cm

#endif

) {

  MACROBLOCKD *const xd = &dcb->xd;

  FRAME_CONTEXT *const ec_ctx = xd->tile_ctx;

  const int is_inter = is_inter_block(xd->mi[0], xd->tree_type);

  const int pl_ctx = get_eob_plane_ctx(plane, is_inter);

  eob_info *eob_data = dcb->eob_data[plane] + dcb->txb_offset[plane];

  uint16_t *const eob = &(eob_data->eob);

  eob_info *bob_data = dcb->bob_data[plane] + dcb->txb_offset[plane];

  uint16_t *const bob = &(bob_data->eob);

  int eob_extra = 0;

  int eob_pt = 1;

  const int eob_multi_size = txsize_log2_minus4[tx_size];

#if CONFIG_PARAKIT_COLLECT_DATA

  const int qp_index = get_q_ctx(cm->quant_params.base_qindex);

  int idxlist[MAX_CTX_DIM];

  idxlist[0] = qp_index;

  idxlist[1] = pl_ctx;

  idxlist[2] = -1;

  idxlist[3] = -1;

#endif

  switch (eob_multi_size) {

    case 0:

#if CONFIG_PARAKIT_COLLECT_DATA

    {

      eob_pt =

          avm_read_symbol_probdata(r, ec_ctx->eob_flag_cdf16[pl_ctx], idxlist,

                                   cm->prob_models[EOB_FLAG_CDF16]) +

          1;

      break;

    }

#else

      eob_pt =

          avm_read_symbol(r, ec_ctx->eob_flag_cdf16[pl_ctx], EOB_MAX_SYMS - 6,

                          ACCT_INFO("eob_pt", "eob_multi_size:0")) +

          1;

      break;

#endif

    case 1:

#if CONFIG_PARAKIT_COLLECT_DATA

    {

      eob_pt =

          avm_read_symbol_probdata(r, ec_ctx->eob_flag_cdf32[pl_ctx], idxlist,

                                   cm->prob_models[EOB_FLAG_CDF32]) +

          1;

      break;

    }

#else

      eob_pt =

          avm_read_symbol(r, ec_ctx->eob_flag_cdf32[pl_ctx], EOB_MAX_SYMS - 5,

                          ACCT_INFO("eob_pt", "eob_multi_size:1")) +

          1;

      break;

#endif

    case 2:

      eob_pt =

          avm_read_symbol(r, ec_ctx->eob_flag_cdf64[pl_ctx], EOB_MAX_SYMS - 4,

                          ACCT_INFO("eob_pt", "eob_multi_size:2")) +

          1;

      break;

    case 3:

      eob_pt =

          avm_read_symbol(r, ec_ctx->eob_flag_cdf128[pl_ctx], EOB_MAX_SYMS - 3,

                          ACCT_INFO("eob_pt", "eob_multi_size:3")) +

          1;

      break;

    case 4:

      eob_pt =

          avm_read_symbol(r, ec_ctx->eob_flag_cdf256[pl_ctx], EOB_MAX_SYMS - 3,

                          ACCT_INFO("eob_pt", "eob_multi_size:4"));

      if (eob_pt == EOB_PT_INDEX_COUNT - 1)

        eob_pt +=

            avm_read_literal(r, 1, ACCT_INFO("eob_pt", "eob_multi_size:4"));

      eob_pt += 1;

      break;

    case 5:

      eob_pt =

          avm_read_symbol(r, ec_ctx->eob_flag_cdf512[pl_ctx], EOB_MAX_SYMS - 3,

                          ACCT_INFO("eob_pt", "eob_multi_size:5"));

      if (eob_pt == EOB_PT_INDEX_COUNT - 1)

        eob_pt +=

            avm_read_literal(r, 2, ACCT_INFO("eob_pt", "eob_multi_size:5"));

      if (eob_pt >= EOB_MAX_SYMS - 1) {

        avm_internal_error(xd->error_info, AVM_CODEC_CORRUPT_FRAME,

                           "Invalid value for EOB position token");

      }

      eob_pt += 1;

      break;

    case 6:

    default:

      eob_pt =

          avm_read_symbol(r, ec_ctx->eob_flag_cdf1024[pl_ctx], EOB_MAX_SYMS - 3,

                          ACCT_INFO("eob_pt", "eob_multi_size:6"));

      if (eob_pt == EOB_PT_INDEX_COUNT - 1)

        eob_pt +=

            avm_read_literal(r, 2, ACCT_INFO("eob_pt", "eob_multi_size:6"));

      eob_pt += 1;

      break;

  }

  const int eob_offset_bits = av2_eob_offset_bits[eob_pt];

  if (eob_offset_bits > 0) {

    int bit = avm_read_symbol(r, ec_ctx->eob_extra_cdf, 2,

                              ACCT_INFO("eob_extra_cdf"));

    if (bit) {

      eob_extra += (1 << (eob_offset_bits - 1));

    }

    eob_extra +=

        avm_read_literal(r, eob_offset_bits - 1, ACCT_INFO("eob_extra"));

  }

  *eob = rec_eob_pos(eob_pt, eob_extra);

  *bob = *eob;  // escape character

}

uint8_t av2_read_sig_txtype(const AV2_COMMON *const cm, DecoderCodingBlock *dcb,

                            avm_reader *const r, const int blk_row,

                            const int blk_col, const int plane,

                            const TXB_CTX *const txb_ctx,

                            const TX_SIZE tx_size) {

  MACROBLOCKD *const xd = &dcb->xd;

  FRAME_CONTEXT *const ec_ctx = xd->tile_ctx;

  const TX_SIZE txs_ctx = get_txsize_entropy_ctx(tx_size);

  const int is_inter = is_inter_block(xd->mi[0], xd->tree_type);

  eob_info *eob_data = dcb->eob_data[plane] + dcb->txb_offset[plane];

  uint16_t *const eob = &(eob_data->eob);

  uint16_t *const max_scan_line = &(eob_data->max_scan_line);

  *max_scan_line = 0;

  *eob = 0;

  int txb_skip_ctx = txb_ctx->txb_skip_ctx;

  int all_zero;

  if (plane == AVM_PLANE_Y || plane == AVM_PLANE_U) {

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

    const int pred_mode_ctx =

        (is_inter || mbmi->fsc_mode[xd->tree_type == CHROMA_PART]) ? 1 : 0;

    all_zero = avm_read_symbol(

        r, ec_ctx->txb_skip_cdf[pred_mode_ctx][txs_ctx][txb_skip_ctx], 2,

        ACCT_INFO("all_zero", "plane_y_or_u"));

  } else {

    txb_skip_ctx += (xd->eob_u_flag ? V_TXB_SKIP_CONTEXT_OFFSET : 0);

    all_zero = avm_read_symbol(r, ec_ctx->v_txb_skip_cdf[txb_skip_ctx], 2,

                               ACCT_INFO("all_zero", "plane_v"));

  }

#if CONFIG_INSPECTION

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

  if (plane == 0) {

    const int txk_type_idx =

        av2_get_txk_type_index(mbmi->sb_type[0], blk_row, blk_col);

    mbmi->tx_skip[txk_type_idx] = all_zero;

  }

#endif  // CONFIG_INSPECTION

  if (plane == AVM_PLANE_U) {

    xd->eob_u_flag = all_zero ? 0 : 1;

  }

#if !CONFIG_INSPECTION

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

#endif  // !CONFIG_INSPECTION

  if (is_inter_block(mbmi, xd->tree_type)) {

    const int txb_idx = get_tx_partition_idx(mbmi, plane);

    mbmi->is_wide_angle[plane > 0][txb_idx] = 0;

    mbmi->mapped_intra_mode[plane > 0][txb_idx] = DC_PRED;

  } else {

    PREDICTION_MODE mode = (plane == PLANE_TYPE_Y ? mbmi->mode : mbmi->uv_mode);

    const int angle_delta =

        mbmi->angle_delta[plane != AVM_PLANE_Y] * ANGLE_STEP;

    wide_angle_mapping(mbmi, angle_delta, tx_size, mode, plane);

  }

  if (all_zero) {

    *max_scan_line = 0;

    if (plane == 0) {

      xd->tx_type_map[blk_row * xd->tx_type_map_stride + blk_col] = DCT_DCT;

    }

    return 0;

  }

  decode_eob(dcb, r, plane, tx_size

#if CONFIG_PARAKIT_COLLECT_DATA

             ,

             cm

#endif

  );

  av2_read_tx_type(cm, xd, blk_row, blk_col, tx_size, r, plane, *eob,

                   is_inter ? 0 : *eob);

  if (plane == AVM_PLANE_U && is_cctx_allowed(cm, xd)) {

    const int skip_cctx = is_inter ? 0 : (*eob == 1);

    if (!all_zero && !skip_cctx) {

      av2_read_cctx_type(cm, xd, blk_row, blk_col, tx_size, r);

    } else {

      int row_offset, col_offset;

      get_chroma_mi_offsets(xd, &row_offset, &col_offset);

      update_cctx_array(xd, blk_row, blk_col, row_offset, col_offset, tx_size,

                        CCTX_NONE);

    }

  }

  return 1;

}

uint8_t av2_read_coeffs_txb_skip(const AV2_COMMON *const cm,

                                 DecoderCodingBlock *dcb, avm_reader *const r,

                                 const int blk_row, const int blk_col,

                                 const int plane, const TX_SIZE tx_size) {

  MACROBLOCKD *const xd = &dcb->xd;

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

  FRAME_CONTEXT *const ec_ctx = xd->tile_ctx;

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

  const PLANE_TYPE plane_type = get_plane_type(plane);

  const int32_t max_value = (1 << (7 + xd->bd)) - 1;

  const int32_t min_value = -(1 << (7 + xd->bd));

  const int32_t *const dequant = pd->seg_dequant_QTX[mbmi->segment_id];

  tran_low_t *const tcoeffs = dcb->dqcoeff_block[plane] + dcb->cb_offset[plane];

  const int shift = av2_get_tx_scale(tx_size);

  const int bwl = get_txb_bwl(tx_size);

  const int width = get_txb_wide(tx_size);

  const int height = get_txb_high(tx_size);

#if CONFIG_INSPECTION

  tran_low_t *const tcoeffs_copy =

      dcb->dqcoeff_block_copy[plane] + dcb->cb_offset[plane];

  tran_low_t *const quant_coeffs =

      dcb->qcoeff_block[plane] + dcb->cb_offset[plane];

  tran_low_t *const dequant_values =

      dcb->dequant_values[plane] + dcb->cb_offset[plane];

  // For TX sizes > 32x32, all coeffs are zero except for top-left 32x32.

  const int coeff_width = AVMMIN(width, 32);

  const int coeff_height = AVMMIN(height, 32);

  memset(tcoeffs_copy, 0, sizeof(tran_low_t) * coeff_width * coeff_height);

  memset(quant_coeffs, 0, sizeof(tran_low_t) * coeff_width * coeff_height);

  memset(dequant_values, 0, sizeof(tran_low_t) * coeff_width * coeff_height);

#endif  // CONFIG_INSPECTION

  int cul_level = 0;

  int dc_val = 0;

  uint8_t levels_buf[TX_PAD_2D];

  uint8_t *const levels = set_levels(levels_buf, width);

  int8_t signs_buf[TX_PAD_2D];

  int8_t *const signs = set_signs(signs_buf, width);

  eob_info *eob_data = dcb->eob_data[plane] + dcb->txb_offset[plane];

  eob_data->max_scan_line = 0;

  eob_data->eob = av2_get_max_eob(tx_size);

  eob_info *bob_data = dcb->bob_data[plane] + dcb->txb_offset[plane];

  bob_data->max_scan_line = 0;

  const TX_TYPE tx_type =

      av2_get_tx_type(xd, plane_type, blk_row, blk_col, tx_size,

                      is_reduced_tx_set_used(cm, plane_type));

  const qm_val_t *iqmatrix =

      av2_get_iqmatrix(&cm->quant_params, xd, plane, tx_size, tx_type);

#if CONFIG_INSPECTION

  for (int c = 0; c < width * height; c++) {

    dequant_values[c] = get_dqv(dequant, c, iqmatrix);

  }

#endif  // CONFIG_INSPECTION

  const SCAN_ORDER *const scan_order = get_scan(tx_size, tx_type);

  const int16_t *const scan = scan_order->scan;

  const TX_SIZE txs_ctx = get_txsize_entropy_ctx(tx_size);

  const int size_ctx = AVMMIN(txs_ctx, TX_16X16);

  if (eob_data->eob > 1) {

    memset(levels_buf, 0, sizeof(*levels_buf) * TX_PAD_2D);

    memset(signs_buf, 0, sizeof(*signs_buf) * TX_PAD_2D);

    base_fsc_cdf_arr base_cdf = ec_ctx->coeff_base_cdf_idtx[size_ctx];

    br_cdf_arr br_cdf = ec_ctx->coeff_br_cdf_idtx[size_ctx];

    const int bob = av2_get_max_eob(tx_size) - bob_data->eob;

    {

      const int pos = scan[bob];

      const int coeff_ctx_bob = get_lower_levels_ctx_bob(bwl, height, bob);

      const int nsymbs_bob = 3;

      avm_cdf_prob *cdf_bob =

          ec_ctx->coeff_base_bob_cdf[size_ctx][coeff_ctx_bob];

      int level = avm_read_symbol(r, cdf_bob, nsymbs_bob,

                                  ACCT_INFO("level", "cdf_bob")) +

                  1;

      if (level > NUM_BASE_LEVELS) {

        const int br_ctx = get_br_ctx_skip(levels, pos, bwl);

        avm_cdf_prob *cdf = br_cdf[br_ctx];

        level += read_low_range(r, cdf);

      }

      levels[get_padded_idx_left(pos, bwl)] = level;

    }

    read_coeffs_forward_2d(r, bob + 1, eob_data->eob - 1, scan, bwl, levels,

                           base_cdf, br_cdf);

  }

  const int bob = av2_get_max_eob(tx_size) - bob_data->eob;

  int hr_level_avg = 0;

  for (int c = bob; c < eob_data->eob; c++) {

    const int pos = scan[c];

    const int sign_idx = get_padded_idx_left(pos, bwl);

    uint8_t sign;

    tran_low_t level = levels[get_padded_idx_left(pos, bwl)];

    if (level) {

      eob_data->max_scan_line = AVMMAX(eob_data->max_scan_line, pos);

      int idtx_sign_ctx = get_sign_ctx_skip(signs, levels, pos, bwl);

      sign = avm_read_symbol(r, ec_ctx->idtx_sign_cdf[size_ctx][idtx_sign_ctx],

                             2, ACCT_INFO("sign"));

      signs[sign_idx] = sign > 0 ? -1 : 1;

      level = read_high_range(xd, r, 0, level, 0, &hr_level_avg, plane);

      if (c == 0) dc_val = sign ? -level : level;

      // Bitmasking to clamp level to valid range:

      // The valid range for 8/10/12 bit video is at most 14/16/18 bit

      if ((level >> 20) != 0) {

        avm_internal_error(xd->error_info, AVM_CODEC_ERROR,

                           "Invalid FSC coeff level beyond 20 bits");

      }

      cul_level += level;

      tran_low_t dq_coeff;

      // Bitmasking to clamp dq_coeff to valid range:

      // The valid range for 8/10/12 bit video is at most 17/19/21 bits

      const int64_t dq_coeff_hp =

          (int64_t)level * get_dqv(dequant, scan[c], iqmatrix) & 0xffffff;

      dq_coeff =

          (tran_low_t)(ROUND_POWER_OF_TWO_64(dq_coeff_hp, QUANT_TABLE_BITS));

      dq_coeff = dq_coeff >> shift;

      if (sign) {

        dq_coeff = -dq_coeff;

      }

      tcoeffs[pos] = clamp(dq_coeff, min_value, max_value);

#if CONFIG_INSPECTION

      tcoeffs_copy[pos] = tcoeffs[pos];

      quant_coeffs[pos] = sign ? -level : level;

#endif  // CONFIG_INSPECTION

    }

  }

  cul_level = AVMMIN(COEFF_CONTEXT_MASK, cul_level);

  set_dc_sign(&cul_level, dc_val);

  return cul_level;

}

// This function returns the partial absolute level of the coefficient

// with hidden parity.

static INLINE tran_low_t read_coeff_hidden(avm_reader *r, TX_CLASS tx_class,

                                           const int16_t *scan, int bwl,

                                           uint8_t *levels, int parity,

                                           base_ph_cdf_arr base_cdf_ph) {

  int q_index;

  const int pos = scan[0];

  int ctx_idx = get_base_ctx_ph(levels, pos, bwl, tx_class);

  q_index = avm_read_symbol(r, base_cdf_ph[ctx_idx], 4, ACCT_INFO("q_index"));

  assert(q_index <= MAX_BASE_BR_RANGE);

  uint8_t level = (q_index << 1) + parity;

  levels[get_padded_idx(pos, bwl)] = level;

  return level;

}

uint8_t av2_read_coeffs_txb(const AV2_COMMON *const cm, DecoderCodingBlock *dcb,

                            avm_reader *const r, const int blk_row,

                            const int blk_col, const int plane,

                            const TXB_CTX *const txb_ctx,

                            const TX_SIZE tx_size) {

  MACROBLOCKD *const xd = &dcb->xd;

  FRAME_CONTEXT *const ec_ctx = xd->tile_ctx;

  const int32_t max_value = (1 << (7 + xd->bd)) - 1;

  const int32_t min_value = -(1 << (7 + xd->bd));

  const TX_SIZE txs_ctx = get_txsize_entropy_ctx(tx_size);

  const PLANE_TYPE plane_type = get_plane_type(plane);

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

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

  const int32_t *const dequant = pd->seg_dequant_QTX[mbmi->segment_id];

  tran_low_t *const tcoeffs = dcb->dqcoeff_block[plane] + dcb->cb_offset[plane];

  const int shift = av2_get_tx_scale(tx_size);

  const int bwl = get_txb_bwl(tx_size);

  const int width = get_txb_wide(tx_size);

  const int height = get_txb_high(tx_size);

#if CONFIG_INSPECTION

  tran_low_t *const tcoeffs_copy =

      dcb->dqcoeff_block_copy[plane] + dcb->cb_offset[plane];

  tran_low_t *const quant_coeffs =

      dcb->qcoeff_block[plane] + dcb->cb_offset[plane];

  tran_low_t *const dequant_values =

      dcb->dequant_values[plane] + dcb->cb_offset[plane];

  // For TX sizes > 32x32, all coeffs are zero except for top-left 32x32.

  const int coeff_width = AVMMIN(width, 32);

  const int coeff_height = AVMMIN(height, 32);

  memset(tcoeffs_copy, 0, sizeof(tran_low_t) * coeff_width * coeff_height);

  memset(quant_coeffs, 0, sizeof(tran_low_t) * coeff_width * coeff_height);

  memset(dequant_values, 0, sizeof(tran_low_t) * coeff_width * coeff_height);

#endif  // CONFIG_INSPECTION

  int cul_level = 0;

  int dc_val = 0;

  uint8_t levels_buf[TX_PAD_2D];

  uint8_t *const levels = set_levels(levels_buf, width);

  eob_info *eob_data = dcb->eob_data[plane] + dcb->txb_offset[plane];

  uint16_t *const eob = &(eob_data->eob);

  uint16_t *const max_scan_line = &(eob_data->max_scan_line);

  const TX_TYPE tx_type =

      av2_get_tx_type(xd, plane_type, blk_row, blk_col, tx_size,

                      is_reduced_tx_set_used(cm, plane_type));

  const TX_CLASS tx_class = tx_type_to_class[get_primary_tx_type(tx_type)];

  const qm_val_t *iqmatrix =

      av2_get_iqmatrix(&cm->quant_params, xd, plane, tx_size, tx_type);

#if CONFIG_INSPECTION

  for (int c = 0; c < width * height; c++) {

    dequant_values[c] = get_dqv(dequant, c, iqmatrix);

  }

#endif  // CONFIG_INSPECTION

  const SCAN_ORDER *const scan_order = get_scan(tx_size, tx_type);

  const int16_t *const scan = scan_order->scan;

  // read  sec_tx_type here

  // Only y plane's sec_tx_type is transmitted

  if ((plane == AVM_PLANE_Y) &&

      (is_inter_block(mbmi, xd->tree_type)

           ? (*eob > 3 && cm->seq_params.enable_inter_ist)

           : (*eob != 1 && cm->seq_params.enable_ist &&

              !xd->mi[0]->fsc_mode[xd->tree_type == CHROMA_PART]))) {

    av2_read_sec_tx_type(cm, xd, blk_row, blk_col, tx_size, eob, r);

  }

  if (*eob > 1) {

    memset(levels_buf, 0, sizeof(*levels_buf) * TX_PAD_2D);

  }

  int tcq_mode = tcq_enable(cm->features.tcq_mode,

                            xd->lossless[mbmi->segment_id], plane, tx_class);

  int state = tcq_init_state(tcq_mode);

  {

    // Read the non-zero coefficient with scan index eob-1

    // TODO(angiebird): Put this into a function

    const int c = *eob - 1;

    const int pos = scan[c];

    const int coeff_ctx = get_lower_levels_ctx_eob(bwl, height, c);

    int level = 0;

    const int row = pos >> bwl;

    const int col = pos - (row << bwl);

    int limits = get_lf_limits(row, col, tx_class, plane);

    if (plane > 0) {

      if (limits) {

        avm_cdf_prob *cdf = ec_ctx->coeff_base_lf_eob_uv_cdf[coeff_ctx];

        level +=

            avm_read_symbol(r, cdf, LF_BASE_SYMBOLS - 1,

                            ACCT_INFO("level", "coeff_base_lf_eob_uv_cdf")) +

            1;

      } else {

        avm_cdf_prob *cdf = ec_ctx->coeff_base_eob_uv_cdf[coeff_ctx];

        level += avm_read_symbol(r, cdf, 3,

                                 ACCT_INFO("level", "coeff_base_eob_uv_cdf")) +

                 1;

        if (level > NUM_BASE_LEVELS) {

          const int br_ctx = 0; /* get_lf_ctx_eob */

          cdf = ec_ctx->coeff_br_uv_cdf[br_ctx];

          level += read_low_range(r, cdf);

        }

      }

    } else {

      if (limits) {

        avm_cdf_prob *cdf = ec_ctx->coeff_base_lf_eob_cdf[txs_ctx][coeff_ctx];

        level += avm_read_symbol(r, cdf, LF_BASE_SYMBOLS - 1,

                                 ACCT_INFO("level", "coeff_base_lf_eob_cdf")) +

                 1;

        if (level > LF_NUM_BASE_LEVELS) {

          const int br_ctx = get_br_ctx_lf_eob(pos, tx_class);

          cdf = ec_ctx->coeff_br_lf_cdf[br_ctx];

          level += read_low_range(r, cdf);

        }

      } else {

        avm_cdf_prob *cdf = ec_ctx->coeff_base_eob_cdf[txs_ctx][coeff_ctx];

        level += avm_read_symbol(r, cdf, 3,

                                 ACCT_INFO("level", "coeff_base_eob_cdf")) +

                 1;

        if (level > NUM_BASE_LEVELS) {

          const int br_ctx = 0; /* get_lf_ctx_eob */

          cdf = ec_ctx->coeff_br_cdf[br_ctx];

          level += read_low_range(r, cdf);

        }

      }

    }

    levels[get_padded_idx(pos, bwl)] = level;

    state = tcq_next_state(state, level);

  }

  bool enable_parity_hiding =

      cm->features.allow_parity_hiding && !xd->lossless[mbmi->segment_id] &&

      plane == PLANE_TYPE_Y &&

      ph_allowed_tx_types[get_primary_tx_type(tx_type)] && (*eob > PHTHRESH);

  int num_nz = 0, sum_abs1 = 0;

  bool is_hidden = false;

  if (*eob > 1) {

    base_lf_cdf_arr base_lf_cdf = ec_ctx->coeff_base_lf_cdf[txs_ctx];

    br_cdf_arr br_lf_cdf = ec_ctx->coeff_br_lf_cdf;

    base_cdf_arr base_cdf = ec_ctx->coeff_base_cdf[txs_ctx];

    br_cdf_arr br_cdf = ec_ctx->coeff_br_cdf;

    base_lf_uv_cdf_arr base_lf_uv_cdf = ec_ctx->coeff_base_lf_uv_cdf;

    base_uv_cdf_arr base_uv_cdf = ec_ctx->coeff_base_uv_cdf;

    br_cdf_arr br_uv_cdf = ec_ctx->coeff_br_uv_cdf;

    if (tx_class == TX_CLASS_2D) {

      read_coeffs_reverse_2d(r, 1, *eob - 2, scan, bwl, levels, base_lf_cdf,

                             br_lf_cdf, plane, base_cdf, br_cdf, base_lf_uv_cdf,

                             base_uv_cdf, br_uv_cdf, &state);

      if (enable_parity_hiding) {

        for (int si = *eob - 1; si > 0; --si) {

          int pos = scan[si];

          int level =

              AVMMIN(levels[get_padded_idx(pos, bwl)], MAX_BASE_BR_RANGE);

          if (level) {

            ++num_nz;

            sum_abs1 += level;

          }

        }

        is_hidden = num_nz >= PHTHRESH;

      }

      if (is_hidden) {

        read_coeff_hidden(r, tx_class, scan, bwl, levels, (sum_abs1 & 1),

                          ec_ctx->coeff_base_ph_cdf);

      } else {

        read_coeffs_reverse(r, tx_class, 0, 0, scan, bwl, levels, base_lf_cdf,

                            br_lf_cdf, plane, base_cdf, br_cdf, base_lf_uv_cdf,

                            base_uv_cdf, br_uv_cdf, &state);

      }

    } else {

      read_coeffs_reverse(r, tx_class, 1, *eob - 2, scan, bwl, levels,

                          base_lf_cdf, br_lf_cdf, plane, base_cdf, br_cdf,

                          base_lf_uv_cdf, base_uv_cdf, br_uv_cdf, &state);

      if (enable_parity_hiding) {

        for (int si = *eob - 1; si > 0; --si) {

          int pos = scan[si];

          int level =

              AVMMIN(levels[get_padded_idx(pos, bwl)], MAX_BASE_BR_RANGE);

          if (level) {

            ++num_nz;

            sum_abs1 += level;

          }

        }

        is_hidden = num_nz >= PHTHRESH;

      }

      if (is_hidden) {

        read_coeff_hidden(r, tx_class, scan, bwl, levels, (sum_abs1 & 1),

                          ec_ctx->coeff_base_ph_cdf);

      } else {

        read_coeffs_reverse(r, tx_class, 0, 0, scan, bwl, levels, base_lf_cdf,

                            br_lf_cdf, plane, base_cdf, br_cdf, base_lf_uv_cdf,

                            base_uv_cdf, br_uv_cdf, &state);

      }

    }

  }

  int hr_level_avg = 0;

  for (int c = *eob - 1; c >= 0; --c) {

    const int pos = scan[c];

    uint8_t sign;

    tran_low_t level = levels[get_padded_idx(pos, bwl)];

    const int tmp_sign_idx = pos;

    if (plane == AVM_PLANE_U) {

      xd->tmp_sign[pos] = 0;

    }

    if (level) {

      *max_scan_line = AVMMAX(*max_scan_line, pos);

      const int row = pos >> bwl;

      const int col = pos - (row << bwl);

      const bool dc_2dtx = (c == 0);

      const bool dc_hor = (col == 0) && tx_class == TX_CLASS_HORIZ;

      const bool dc_ver = (row == 0) && tx_class == TX_CLASS_VERT;

      if (dc_2dtx || dc_hor || dc_ver) {

        const int dc_sign_ctx = dc_2dtx ? txb_ctx->dc_sign_ctx : 0;

        if (plane == AVM_PLANE_Y || plane == AVM_PLANE_U) {

          if (plane == AVM_PLANE_Y) {

            sign = avm_read_symbol(

                r,

                ec_ctx->dc_sign_cdf[plane_type][is_hidden ? 1 : 0][dc_sign_ctx],

                2, ACCT_INFO("sign", "dc_sign_cdf", "plane_y_or_u"));

          } else {

            sign = avm_read_literal(

                r, 1, ACCT_INFO("sign", "dc_sign_cdf", "plane_y_or_u"));

          }

        } else {

          sign = avm_read_literal(

              r, 1, ACCT_INFO("sign", "v_dc_sign_cdf", "plane_v"));

        }

        if (plane == AVM_PLANE_U) xd->tmp_sign[tmp_sign_idx] = (sign ? 2 : 1);

      } else {

        sign = avm_read_bit(r, ACCT_INFO("sign"));

      }

      if (is_hidden && c == 0) {

        if (level >= ((NUM_BASE_LEVELS + 1) << 1)) {

          int hr_level = (read_adaptive_hr(xd, r, hr_level_avg >> 1) << 1);

          level += hr_level;

          hr_level_avg = (hr_level_avg + hr_level) >> 1;

        }

      } else {

        int limits = get_lf_limits(row, col, tx_class, plane);

        level = read_high_range(xd, r, tcq_mode, level, limits, &hr_level_avg,

                                plane);

      }

      if (c == 0) dc_val = sign ? -level : level;

      // Bitmasking to clamp level to valid range:

      //   The valid range for 8/10/12 bit vdieo is at most 14/16/18 bit

      if ((level >> 20) != 0) {

        avm_internal_error(xd->error_info, AVM_CODEC_ERROR,

                           "Invalid coeff level beyond 20 bits");

      }

      cul_level += level;

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

        tcoeffs[pos] = sign ? -level : level;

      } else {

        tran_low_t dq_coeff;

        // Bitmasking to clamp dq_coeff to valid range:

        //   The valid range for 8/10/12 bit video is at most 17/19/21 bit

        const int64_t dq_coeff_hp =

            (int64_t)level * get_dqv(dequant, scan[c], iqmatrix) & 0xffffff;

        dq_coeff =

            (tran_low_t)(ROUND_POWER_OF_TWO_64(dq_coeff_hp, QUANT_TABLE_BITS));

        dq_coeff = dq_coeff >> shift;

        if (sign) {

          dq_coeff = -dq_coeff;

        }

        tcoeffs[pos] = clamp(dq_coeff, min_value, max_value);

#if CONFIG_INSPECTION

        tcoeffs_copy[pos] = tcoeffs[pos];

        quant_coeffs[pos] = sign ? -level : level;

#endif  // CONFIG_INSPECTION

      }

    }

  }

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

    state = tcq_init_state(tcq_mode);

    for (int c = *eob - 1; c >= 0; --c) {

      const int pos = scan[c];

      int Qx = tcq_quant(state);

      int level = abs(tcoeffs[pos]);

      if (abs(tcoeffs[pos])) {

        tran_low_t dq_coeff;

        tran_low_t qIdx = AVMMAX(0, (abs(tcoeffs[pos]) << 1) - Qx);

        if ((qIdx >> 20) != 0) {

          avm_internal_error(xd->error_info, AVM_CODEC_ERROR,

                             "Invalid TCQ coeff level beyond 20 bits");

        }

        // Bitmasking to clamp dq_coeff to valid range:

        //   The valid range for 8/10/12 bit video is at most 17/19/21 bit

        int tempdqv = get_dqv(dequant, scan[c], iqmatrix);

        const int64_t dq_coeff_hp = (int64_t)qIdx * tempdqv & 0xffffff;

        dq_coeff =

            (tran_low_t)(ROUND_POWER_OF_TWO_64(dq_coeff_hp, QUANT_TABLE_BITS));

        dq_coeff = dq_coeff >> (shift + 1);

        // dq_coeff = dq_coeff >> (shift);

        // dq_coeff = (tran_low_t)ROUND_POWER_OF_TWO_64((tran_high_t)qIdx *

        // get_dqv(dequant, scan[c], iqmatrix), QUANT_TABLE_BITS) >> (shift +

        // 1);

        if (tcoeffs[pos] < 0) {

          dq_coeff = -dq_coeff;

        }

        tcoeffs[pos] = clamp(dq_coeff, min_value, max_value);

      }

      // state transition

      state = tcq_next_state(state, level);

    }

  }

  cul_level = AVMMIN(COEFF_CONTEXT_MASK, cul_level);

  // DC value

  set_dc_sign(&cul_level, dc_val);

  return cul_level;

}

void av2_read_coeffs_txb_facade(const AV2_COMMON *const cm,

                                DecoderCodingBlock *dcb, avm_reader *const r,

                                const int plane, const int row, const int col,

                                const TX_SIZE tx_size) {

#if TXCOEFF_TIMER

  struct avm_usec_timer timer;

  avm_usec_timer_start(&timer);

#endif

  MACROBLOCKD *const xd = &dcb->xd;

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

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

  const BLOCK_SIZE plane_bsize = get_mb_plane_block_size(

      xd, mbmi, plane, pd->subsampling_x, pd->subsampling_y);

  TXB_CTX txb_ctx;

  get_txb_ctx(plane_bsize, tx_size, plane, pd->above_entropy_context + col,

              pd->left_entropy_context + row, &txb_ctx,

              mbmi->fsc_mode[xd->tree_type == CHROMA_PART] &&

                  cm->seq_params.enable_fsc);

  const uint8_t decode_rest =

      av2_read_sig_txtype(cm, dcb, r, row, col, plane, &txb_ctx, tx_size);

  const PLANE_TYPE plane_type = get_plane_type(plane);

  const TX_TYPE tx_type =

      av2_get_tx_type(xd, plane_type, row, col, tx_size,

                      is_reduced_tx_set_used(cm, plane_type));

  const int is_inter = is_inter_block(mbmi, xd->tree_type);

  uint8_t cul_level = 0;

  if (decode_rest) {

    if (((cm->seq_params.enable_fsc &&

          mbmi->fsc_mode[xd->tree_type == CHROMA_PART] &&

          get_primary_tx_type(tx_type) == IDTX && plane == PLANE_TYPE_Y) ||

         use_inter_fsc(cm, plane, tx_type, is_inter))) {

      cul_level =

          av2_read_coeffs_txb_skip(cm, dcb, r, row, col, plane, tx_size);

    } else {

      cul_level =

          av2_read_coeffs_txb(cm, dcb, r, row, col, plane, &txb_ctx, tx_size);

    }

  } else {

    av2_update_txk_skip_array(cm, xd->mi_row, xd->mi_col, xd->tree_type,

                              &mbmi->chroma_ref_info, plane, row, col, tx_size);

  }

  av2_set_entropy_contexts(xd, pd, plane, plane_bsize, tx_size, cul_level, col,

                           row);

  if (is_inter_block(mbmi, xd->tree_type) && (plane == 0)) {

    const TX_TYPE tx_type_inter =

        av2_get_tx_type(xd, plane_type, row, col, tx_size,

                        is_reduced_tx_set_used(cm, plane_type));

    update_txk_array(xd, row, col, tx_size, tx_type_inter);

  }

#if TXCOEFF_TIMER

  avm_usec_timer_mark(&timer);

  const int64_t elapsed_time = avm_usec_timer_elapsed(&timer);