/*

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

 *

 * This source code is subject to the terms of the BSD 2 Clause License and

 * the Alliance for Open Media Patent License 1.0. If the BSD 2 Clause License

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

 * obtain it at www.aomedia.org/license/software. If the Alliance for Open

 * Media Patent License 1.0 was not distributed with this source code in the

 * PATENTS file, you can obtain it at www.aomedia.org/license/patent.

 */

#include <assert.h>

#include <limits.h>

#include <stdbool.h>

#include <stdint.h>

#include <stdio.h>

#include <string.h>

#include "aom/aom_encoder.h"

#include "aom_dsp/aom_dsp_common.h"

#include "aom_dsp/binary_codes_writer.h"

#include "aom_dsp/bitwriter_buffer.h"

#include "aom_mem/aom_mem.h"

#include "aom_ports/bitops.h"

#include "aom_ports/mem_ops.h"

#if CONFIG_BITSTREAM_DEBUG

#include "aom_util/debug_util.h"

#endif  // CONFIG_BITSTREAM_DEBUG

#include "av1/common/cdef.h"

#include "av1/common/cfl.h"

#include "av1/common/debugmodes.h"

#include "av1/common/entropy.h"

#include "av1/common/entropymode.h"

#include "av1/common/entropymv.h"

#include "av1/common/mvref_common.h"

#include "av1/common/pred_common.h"

#include "av1/common/reconinter.h"

#include "av1/common/reconintra.h"

#include "av1/common/seg_common.h"

#include "av1/common/tile_common.h"

#include "av1/encoder/bitstream.h"

#include "av1/encoder/cost.h"

#include "av1/encoder/encodemv.h"

#include "av1/encoder/encodetxb.h"

#include "av1/encoder/ethread.h"

#include "av1/encoder/mcomp.h"

#include "av1/encoder/palette.h"

#include "av1/encoder/pickrst.h"

#include "av1/encoder/segmentation.h"

#include "av1/encoder/tokenize.h"

#define ENC_MISMATCH_DEBUG 0

#define SETUP_TIME_OH_CONST 5     // Setup time overhead constant per worker

#define JOB_DISP_TIME_OH_CONST 1  // Job dispatch time overhead per tile

static inline void write_uniform(aom_writer *w, int n, int v) {

  const int l = get_unsigned_bits(n);

  const int m = (1 << l) - n;

  if (l == 0) return;

  if (v < m) {

    aom_write_literal(w, v, l - 1);

  } else {

    aom_write_literal(w, m + ((v - m) >> 1), l - 1);

    aom_write_literal(w, (v - m) & 1, 1);

  }

}

#if !CONFIG_REALTIME_ONLY

static inline void loop_restoration_write_sb_coeffs(

    const AV1_COMMON *const cm, MACROBLOCKD *xd, int runit_idx,

    aom_writer *const w, int plane, FRAME_COUNTS *counts);

#endif

static inline void write_intra_y_mode_kf(FRAME_CONTEXT *frame_ctx,

                                         const MB_MODE_INFO *mi,

                                         const MB_MODE_INFO *above_mi,

                                         const MB_MODE_INFO *left_mi,

                                         PREDICTION_MODE mode, aom_writer *w) {

  assert(!is_intrabc_block(mi));

  (void)mi;

  aom_write_symbol(w, mode, get_y_mode_cdf(frame_ctx, above_mi, left_mi),

                   INTRA_MODES);

}

static inline void write_inter_mode(aom_writer *w, PREDICTION_MODE mode,

                                    FRAME_CONTEXT *ec_ctx,

                                    const int16_t mode_ctx) {

  const int16_t newmv_ctx = mode_ctx & NEWMV_CTX_MASK;

  aom_write_symbol(w, mode != NEWMV, ec_ctx->newmv_cdf[newmv_ctx], 2);

  if (mode != NEWMV) {

    const int16_t zeromv_ctx =

        (mode_ctx >> GLOBALMV_OFFSET) & GLOBALMV_CTX_MASK;

    aom_write_symbol(w, mode != GLOBALMV, ec_ctx->zeromv_cdf[zeromv_ctx], 2);

    if (mode != GLOBALMV) {

      int16_t refmv_ctx = (mode_ctx >> REFMV_OFFSET) & REFMV_CTX_MASK;

      aom_write_symbol(w, mode != NEARESTMV, ec_ctx->refmv_cdf[refmv_ctx], 2);

    }

  }

}

static inline void write_drl_idx(FRAME_CONTEXT *ec_ctx,

                                 const MB_MODE_INFO *mbmi,

                                 const MB_MODE_INFO_EXT_FRAME *mbmi_ext_frame,

                                 aom_writer *w) {

  assert(mbmi->ref_mv_idx < 3);

  const int new_mv = mbmi->mode == NEWMV || mbmi->mode == NEW_NEWMV;

  if (new_mv) {

    int idx;

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

      if (mbmi_ext_frame->ref_mv_count > idx + 1) {

        uint8_t drl_ctx = av1_drl_ctx(mbmi_ext_frame->weight, idx);

        aom_write_symbol(w, mbmi->ref_mv_idx != idx, ec_ctx->drl_cdf[drl_ctx],

                         2);

        if (mbmi->ref_mv_idx == idx) return;

      }

    }

    return;

  }

  if (have_nearmv_in_inter_mode(mbmi->mode)) {

    int idx;

    // TODO(jingning): Temporary solution to compensate the NEARESTMV offset.

    for (idx = 1; idx < 3; ++idx) {

      if (mbmi_ext_frame->ref_mv_count > idx + 1) {

        uint8_t drl_ctx = av1_drl_ctx(mbmi_ext_frame->weight, idx);

        aom_write_symbol(w, mbmi->ref_mv_idx != (idx - 1),

                         ec_ctx->drl_cdf[drl_ctx], 2);

        if (mbmi->ref_mv_idx == (idx - 1)) return;

      }

    }

    return;

  }

}

static inline void write_inter_compound_mode(MACROBLOCKD *xd, aom_writer *w,

                                             PREDICTION_MODE mode,

                                             const int16_t mode_ctx) {

  assert(is_inter_compound_mode(mode));

  aom_write_symbol(w, INTER_COMPOUND_OFFSET(mode),

                   xd->tile_ctx->inter_compound_mode_cdf[mode_ctx],

                   INTER_COMPOUND_MODES);

}

static inline void write_tx_size_vartx(MACROBLOCKD *xd,

                                       const MB_MODE_INFO *mbmi,

                                       TX_SIZE tx_size, int depth, int blk_row,

                                       int blk_col, aom_writer *w) {

  FRAME_CONTEXT *const ec_ctx = xd->tile_ctx;

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

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

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

  if (depth == MAX_VARTX_DEPTH) {

    txfm_partition_update(xd->above_txfm_context + blk_col,

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

    return;

  }

  const 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(mbmi->bsize, blk_row, blk_col);

  const int write_txfm_partition =

      tx_size == mbmi->inter_tx_size[txb_size_index];

  if (write_txfm_partition) {

    aom_write_symbol(w, 0, ec_ctx->txfm_partition_cdf[ctx], 2);

    txfm_partition_update(xd->above_txfm_context + blk_col,

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

    // TODO(yuec): set correct txfm partition update for qttx

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

    aom_write_symbol(w, 1, ec_ctx->txfm_partition_cdf[ctx], 2);

    if (sub_txs == TX_4X4) {

      txfm_partition_update(xd->above_txfm_context + blk_col,

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

      return;

    }

    assert(bsw > 0 && bsh > 0);

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

      const int offsetr = blk_row + row;

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

        const int offsetc = blk_col + col;

        write_tx_size_vartx(xd, mbmi, sub_txs, depth + 1, offsetr, offsetc, w);

      }

    }

  }

}

static inline void write_selected_tx_size(const MACROBLOCKD *xd,

                                          aom_writer *w) {

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

  const BLOCK_SIZE bsize = mbmi->bsize;

  FRAME_CONTEXT *ec_ctx = xd->tile_ctx;

  if (block_signals_txsize(bsize)) {

    const TX_SIZE tx_size = mbmi->tx_size;

    const int tx_size_ctx = get_tx_size_context(xd);

    const int depth = tx_size_to_depth(tx_size, bsize);

    const int max_depths = bsize_to_max_depth(bsize);

    const int32_t tx_size_cat = bsize_to_tx_size_cat(bsize);

    assert(depth >= 0 && depth <= max_depths);

    assert(!is_inter_block(mbmi));

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

    aom_write_symbol(w, depth, ec_ctx->tx_size_cdf[tx_size_cat][tx_size_ctx],

                     max_depths + 1);

  }

}

static int write_skip(const AV1_COMMON *cm, const MACROBLOCKD *xd,

                      uint8_t segment_id, const MB_MODE_INFO *mi,

                      aom_writer *w) {

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

    return 1;

  } else {

    const int skip_txfm = mi->skip_txfm;

    const int ctx = av1_get_skip_txfm_context(xd);

    FRAME_CONTEXT *ec_ctx = xd->tile_ctx;

    aom_write_symbol(w, skip_txfm, ec_ctx->skip_txfm_cdfs[ctx], 2);

    return skip_txfm;

  }

}

static int write_skip_mode(const AV1_COMMON *cm, const MACROBLOCKD *xd,

                           uint8_t segment_id, const MB_MODE_INFO *mi,

                           aom_writer *w) {

  if (!cm->current_frame.skip_mode_info.skip_mode_flag) return 0;

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

    return 0;

  }

  const int skip_mode = mi->skip_mode;

  if (!is_comp_ref_allowed(mi->bsize)) {

    assert(!skip_mode);

    return 0;

  }

  if (segfeature_active(&cm->seg, segment_id, SEG_LVL_REF_FRAME) ||

      segfeature_active(&cm->seg, segment_id, SEG_LVL_GLOBALMV)) {

    // These features imply single-reference mode, while skip mode implies

    // compound reference. Hence, the two are mutually exclusive.

    // In other words, skip_mode is implicitly 0 here.

    assert(!skip_mode);

    return 0;

  }

  const int ctx = av1_get_skip_mode_context(xd);

  aom_write_symbol(w, skip_mode, xd->tile_ctx->skip_mode_cdfs[ctx], 2);

  return skip_mode;

}

static inline void write_is_inter(const AV1_COMMON *cm, const MACROBLOCKD *xd,

                                  uint8_t segment_id, aom_writer *w,

                                  const int is_inter) {

  if (!segfeature_active(&cm->seg, segment_id, SEG_LVL_REF_FRAME)) {

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

      assert(is_inter);

      return;

    }

    const int ctx = av1_get_intra_inter_context(xd);

    FRAME_CONTEXT *ec_ctx = xd->tile_ctx;

    aom_write_symbol(w, is_inter, ec_ctx->intra_inter_cdf[ctx], 2);

  }

}

static inline void write_motion_mode(const AV1_COMMON *cm, MACROBLOCKD *xd,

                                     const MB_MODE_INFO *mbmi, aom_writer *w) {

  MOTION_MODE last_motion_mode_allowed =

      cm->features.switchable_motion_mode

          ? motion_mode_allowed(cm->global_motion, xd, mbmi,

                                cm->features.allow_warped_motion)

          : SIMPLE_TRANSLATION;

  assert(mbmi->motion_mode <= last_motion_mode_allowed);

  switch (last_motion_mode_allowed) {

    case SIMPLE_TRANSLATION: break;

    case OBMC_CAUSAL:

      aom_write_symbol(w, mbmi->motion_mode == OBMC_CAUSAL,

                       xd->tile_ctx->obmc_cdf[mbmi->bsize], 2);

      break;

    default:

      aom_write_symbol(w, mbmi->motion_mode,

                       xd->tile_ctx->motion_mode_cdf[mbmi->bsize],

                       MOTION_MODES);

  }

}

static inline void write_delta_qindex(const MACROBLOCKD *xd, int delta_qindex,

                                      aom_writer *w) {

  int sign = delta_qindex < 0;

  int abs = sign ? -delta_qindex : delta_qindex;

  int rem_bits, thr;

  int smallval = abs < DELTA_Q_SMALL ? 1 : 0;

  FRAME_CONTEXT *ec_ctx = xd->tile_ctx;

  aom_write_symbol(w, AOMMIN(abs, DELTA_Q_SMALL), ec_ctx->delta_q_cdf,

                   DELTA_Q_PROBS + 1);

  if (!smallval) {

    rem_bits = get_msb(abs - 1);

    thr = (1 << rem_bits) + 1;

    aom_write_literal(w, rem_bits - 1, 3);

    aom_write_literal(w, abs - thr, rem_bits);

  }

  if (abs > 0) {

    aom_write_bit(w, sign);

  }

}

static inline void write_delta_lflevel(const AV1_COMMON *cm,

                                       const MACROBLOCKD *xd, int lf_id,

                                       int delta_lflevel, int delta_lf_multi,

                                       aom_writer *w) {

  int sign = delta_lflevel < 0;

  int abs = sign ? -delta_lflevel : delta_lflevel;

  int rem_bits, thr;

  int smallval = abs < DELTA_LF_SMALL ? 1 : 0;

  FRAME_CONTEXT *ec_ctx = xd->tile_ctx;

  (void)cm;

  if (delta_lf_multi) {

    assert(lf_id >= 0 && lf_id < (av1_num_planes(cm) > 1 ? FRAME_LF_COUNT

                                                         : FRAME_LF_COUNT - 2));

    aom_write_symbol(w, AOMMIN(abs, DELTA_LF_SMALL),

                     ec_ctx->delta_lf_multi_cdf[lf_id], DELTA_LF_PROBS + 1);

  } else {

    aom_write_symbol(w, AOMMIN(abs, DELTA_LF_SMALL), ec_ctx->delta_lf_cdf,

                     DELTA_LF_PROBS + 1);

  }

  if (!smallval) {

    rem_bits = get_msb(abs - 1);

    thr = (1 << rem_bits) + 1;

    aom_write_literal(w, rem_bits - 1, 3);

    aom_write_literal(w, abs - thr, rem_bits);

  }

  if (abs > 0) {

    aom_write_bit(w, sign);

  }

}

static inline void pack_map_tokens(aom_writer *w, const TokenExtra **tp, int n,

                                   int num, MapCdf map_pb_cdf) {

  const TokenExtra *p = *tp;

  const int palette_size_idx = n - PALETTE_MIN_SIZE;

  write_uniform(w, n, p->token);  // The first color index.

  ++p;

  --num;

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

    assert((p->color_ctx >= 0) &&

           (p->color_ctx < PALETTE_COLOR_INDEX_CONTEXTS));

    aom_cdf_prob *color_map_cdf = map_pb_cdf[palette_size_idx][p->color_ctx];

    aom_write_symbol(w, p->token, color_map_cdf, n);

    ++p;

  }

  *tp = p;

}

static inline void pack_txb_tokens(

    aom_writer *w, AV1_COMMON *cm, MACROBLOCK *const x, const TokenExtra **tp,

    const TokenExtra *const tok_end, MACROBLOCKD *xd, MB_MODE_INFO *mbmi,

    int plane, BLOCK_SIZE plane_bsize, aom_bit_depth_t bit_depth, int block,

    int blk_row, int blk_col, TX_SIZE tx_size, TOKEN_STATS *token_stats) {

  const int max_blocks_high = max_block_high(xd, plane_bsize, plane);

  const int max_blocks_wide = max_block_wide(xd, plane_bsize, plane);

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

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

  const TX_SIZE plane_tx_size =

      plane ? av1_get_max_uv_txsize(mbmi->bsize, pd->subsampling_x,

                                    pd->subsampling_y)

            : mbmi->inter_tx_size[av1_get_txb_size_index(plane_bsize, blk_row,

                                                         blk_col)];

  if (tx_size == plane_tx_size || plane) {

    av1_write_coeffs_txb(cm, x, w, blk_row, blk_col, plane, block, tx_size);

#if CONFIG_RD_DEBUG

    TOKEN_STATS tmp_token_stats;

    init_token_stats(&tmp_token_stats);

    token_stats->cost += tmp_token_stats.cost;

#endif

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

    const int step = bsh * bsw;

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

    assert(bsw > 0 && bsh > 0);

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

      const int offsetr = blk_row + r;

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

        const int offsetc = blk_col + c;

        pack_txb_tokens(w, cm, x, tp, tok_end, xd, mbmi, plane, plane_bsize,

                        bit_depth, block, offsetr, offsetc, sub_txs,

                        token_stats);

        block += step;

      }

    }

  }

}

static inline void set_spatial_segment_id(

    const CommonModeInfoParams *const mi_params, uint8_t *segment_ids,

    BLOCK_SIZE bsize, int mi_row, int mi_col, uint8_t segment_id) {

  const int mi_offset = mi_row * mi_params->mi_cols + mi_col;

  const int bw = mi_size_wide[bsize];

  const int bh = mi_size_high[bsize];

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

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

  const int mi_stride = mi_params->mi_cols;

  set_segment_id(segment_ids, mi_offset, xmis, ymis, mi_stride, segment_id);

}

int av1_neg_interleave(int x, int ref, int max) {

  assert(x < max);

  const int diff = x - ref;

  if (!ref) return x;

  if (ref >= (max - 1)) return -x + max - 1;

  if (2 * ref < max) {

    if (abs(diff) <= ref) {

      if (diff > 0)

        return (diff << 1) - 1;

      else

        return ((-diff) << 1);

    }

    return x;

  } else {

    if (abs(diff) < (max - ref)) {

      if (diff > 0)

        return (diff << 1) - 1;

      else

        return ((-diff) << 1);

    }

    return (max - x) - 1;

  }

}

static inline void write_segment_id(AV1_COMP *cpi, MACROBLOCKD *const xd,

                                    const MB_MODE_INFO *const mbmi,

                                    aom_writer *w,

                                    const struct segmentation *seg,

                                    struct segmentation_probs *segp,

                                    int skip_txfm) {

  if (!seg->enabled || !seg->update_map) return;

  AV1_COMMON *const cm = &cpi->common;

  int cdf_num;

  const uint8_t pred = av1_get_spatial_seg_pred(

      cm, xd, &cdf_num, cpi->cyclic_refresh->skip_over4x4);

  const int mi_row = xd->mi_row;

  const int mi_col = xd->mi_col;

  if (skip_txfm) {

    // Still need to transmit tx size for intra blocks even if skip_txfm is

    // true. Changing segment_id may make the tx size become invalid, e.g

    // changing from lossless to lossy.

    assert(is_inter_block(mbmi) || !cpi->enc_seg.has_lossless_segment);

    set_spatial_segment_id(&cm->mi_params, cm->cur_frame->seg_map, mbmi->bsize,

                           mi_row, mi_col, pred);

    set_spatial_segment_id(&cm->mi_params, cpi->enc_seg.map, mbmi->bsize,

                           mi_row, mi_col, pred);

    /* mbmi is read only but we need to update segment_id */

    ((MB_MODE_INFO *)mbmi)->segment_id = pred;

    return;

  }

  const int coded_id =

      av1_neg_interleave(mbmi->segment_id, pred, seg->last_active_segid + 1);

  aom_cdf_prob *pred_cdf = segp->spatial_pred_seg_cdf[cdf_num];

  aom_write_symbol(w, coded_id, pred_cdf, MAX_SEGMENTS);

  set_spatial_segment_id(&cm->mi_params, cm->cur_frame->seg_map, mbmi->bsize,

                         mi_row, mi_col, mbmi->segment_id);

}

#define WRITE_REF_BIT(bname, pname) \

  aom_write_symbol(w, bname, av1_get_pred_cdf_##pname(xd), 2)

// This function encodes the reference frame

static inline void write_ref_frames(const AV1_COMMON *cm, const MACROBLOCKD *xd,

                                    aom_writer *w) {

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

  const int is_compound = has_second_ref(mbmi);

  const uint8_t segment_id = mbmi->segment_id;

  // If segment level coding of this signal is disabled...

  // or the segment allows multiple reference frame options

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

    assert(!is_compound);

    assert(mbmi->ref_frame[0] ==

           get_segdata(&cm->seg, segment_id, SEG_LVL_REF_FRAME));

  } else if (segfeature_active(&cm->seg, segment_id, SEG_LVL_SKIP) ||

             segfeature_active(&cm->seg, segment_id, SEG_LVL_GLOBALMV)) {

    assert(!is_compound);

    assert(mbmi->ref_frame[0] == LAST_FRAME);

  } else {

    // does the feature use compound prediction or not

    // (if not specified at the frame/segment level)

    if (cm->current_frame.reference_mode == REFERENCE_MODE_SELECT) {

      if (is_comp_ref_allowed(mbmi->bsize))

        aom_write_symbol(w, is_compound, av1_get_reference_mode_cdf(xd), 2);

    } else {

      assert((!is_compound) ==

             (cm->current_frame.reference_mode == SINGLE_REFERENCE));

    }

    if (is_compound) {

      const COMP_REFERENCE_TYPE comp_ref_type = has_uni_comp_refs(mbmi)

                                                    ? UNIDIR_COMP_REFERENCE

                                                    : BIDIR_COMP_REFERENCE;

      aom_write_symbol(w, comp_ref_type, av1_get_comp_reference_type_cdf(xd),

                       2);

      if (comp_ref_type == UNIDIR_COMP_REFERENCE) {

        const int bit = mbmi->ref_frame[0] == BWDREF_FRAME;

        WRITE_REF_BIT(bit, uni_comp_ref_p);

        if (!bit) {

          assert(mbmi->ref_frame[0] == LAST_FRAME);

          const int bit1 = mbmi->ref_frame[1] == LAST3_FRAME ||

                           mbmi->ref_frame[1] == GOLDEN_FRAME;

          WRITE_REF_BIT(bit1, uni_comp_ref_p1);

          if (bit1) {

            const int bit2 = mbmi->ref_frame[1] == GOLDEN_FRAME;

            WRITE_REF_BIT(bit2, uni_comp_ref_p2);

          }

        } else {

          assert(mbmi->ref_frame[1] == ALTREF_FRAME);

        }

        return;

      }

      assert(comp_ref_type == BIDIR_COMP_REFERENCE);

      const int bit = (mbmi->ref_frame[0] == GOLDEN_FRAME ||

                       mbmi->ref_frame[0] == LAST3_FRAME);

      WRITE_REF_BIT(bit, comp_ref_p);

      if (!bit) {

        const int bit1 = mbmi->ref_frame[0] == LAST2_FRAME;

        WRITE_REF_BIT(bit1, comp_ref_p1);

      } else {

        const int bit2 = mbmi->ref_frame[0] == GOLDEN_FRAME;

        WRITE_REF_BIT(bit2, comp_ref_p2);

      }

      const int bit_bwd = mbmi->ref_frame[1] == ALTREF_FRAME;

      WRITE_REF_BIT(bit_bwd, comp_bwdref_p);

      if (!bit_bwd) {

        WRITE_REF_BIT(mbmi->ref_frame[1] == ALTREF2_FRAME, comp_bwdref_p1);

      }

    } else {

      const int bit0 = (mbmi->ref_frame[0] <= ALTREF_FRAME &&

                        mbmi->ref_frame[0] >= BWDREF_FRAME);

      WRITE_REF_BIT(bit0, single_ref_p1);

      if (bit0) {

        const int bit1 = mbmi->ref_frame[0] == ALTREF_FRAME;

        WRITE_REF_BIT(bit1, single_ref_p2);

        if (!bit1) {

          WRITE_REF_BIT(mbmi->ref_frame[0] == ALTREF2_FRAME, single_ref_p6);

        }

      } else {

        const int bit2 = (mbmi->ref_frame[0] == LAST3_FRAME ||

                          mbmi->ref_frame[0] == GOLDEN_FRAME);

        WRITE_REF_BIT(bit2, single_ref_p3);

        if (!bit2) {

          const int bit3 = mbmi->ref_frame[0] != LAST_FRAME;

          WRITE_REF_BIT(bit3, single_ref_p4);

        } else {

          const int bit4 = mbmi->ref_frame[0] != LAST3_FRAME;

          WRITE_REF_BIT(bit4, single_ref_p5);

        }

      }

    }

  }

}

static inline void write_filter_intra_mode_info(const AV1_COMMON *cm,

                                                const MACROBLOCKD *xd,

                                                const MB_MODE_INFO *const mbmi,

                                                aom_writer *w) {

  if (av1_filter_intra_allowed(cm, mbmi)) {

    aom_write_symbol(w, mbmi->filter_intra_mode_info.use_filter_intra,

                     xd->tile_ctx->filter_intra_cdfs[mbmi->bsize], 2);

    if (mbmi->filter_intra_mode_info.use_filter_intra) {

      const FILTER_INTRA_MODE mode =

          mbmi->filter_intra_mode_info.filter_intra_mode;

      aom_write_symbol(w, mode, xd->tile_ctx->filter_intra_mode_cdf,

                       FILTER_INTRA_MODES);

    }

  }

}

static inline void write_angle_delta(aom_writer *w, int angle_delta,

                                     aom_cdf_prob *cdf) {

  aom_write_symbol(w, angle_delta + MAX_ANGLE_DELTA, cdf,

                   2 * MAX_ANGLE_DELTA + 1);

}

static inline void write_mb_interp_filter(AV1_COMMON *const cm, ThreadData *td,

                                          aom_writer *w) {

  const MACROBLOCKD *xd = &td->mb.e_mbd;

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

  FRAME_CONTEXT *ec_ctx = xd->tile_ctx;

  if (!av1_is_interp_needed(xd)) {

    int_interpfilters filters = av1_broadcast_interp_filter(

        av1_unswitchable_filter(cm->features.interp_filter));

    assert(mbmi->interp_filters.as_int == filters.as_int);

    (void)filters;

    return;

  }

  if (cm->features.interp_filter == SWITCHABLE) {

    int dir;

    for (dir = 0; dir < 2; ++dir) {

      const int ctx = av1_get_pred_context_switchable_interp(xd, dir);

      InterpFilter filter =

          av1_extract_interp_filter(mbmi->interp_filters, dir);

      aom_write_symbol(w, filter, ec_ctx->switchable_interp_cdf[ctx],

                       SWITCHABLE_FILTERS);

      ++td->interp_filter_selected[filter];

      if (cm->seq_params->enable_dual_filter == 0) return;

    }

  }

}

// Transmit color values with delta encoding. Write the first value as

// literal, and the deltas between each value and the previous one. "min_val" is

// the smallest possible value of the deltas.

static inline void delta_encode_palette_colors(const int *colors, int num,

                                               int bit_depth, int min_val,

                                               aom_writer *w) {

  if (num <= 0) return;

  assert(colors[0] < (1 << bit_depth));

  aom_write_literal(w, colors[0], bit_depth);

  if (num == 1) return;

  int max_delta = 0;

  int deltas[PALETTE_MAX_SIZE];

  memset(deltas, 0, sizeof(deltas));

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

    assert(colors[i] < (1 << bit_depth));

    const int delta = colors[i] - colors[i - 1];

    deltas[i - 1] = delta;

    assert(delta >= min_val);

    if (delta > max_delta) max_delta = delta;

  }

  const int min_bits = bit_depth - 3;

  int bits = AOMMAX(av1_ceil_log2(max_delta + 1 - min_val), min_bits);

  assert(bits <= bit_depth);

  int range = (1 << bit_depth) - colors[0] - min_val;

  aom_write_literal(w, bits - min_bits, 2);

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

    aom_write_literal(w, deltas[i] - min_val, bits);

    range -= deltas[i];

    bits = AOMMIN(bits, av1_ceil_log2(range));

  }

}

// Transmit luma palette color values. First signal if each color in the color

// cache is used. Those colors that are not in the cache are transmitted with

// delta encoding.

static inline void write_palette_colors_y(const MACROBLOCKD *const xd,

                                          const PALETTE_MODE_INFO *const pmi,

                                          int bit_depth, aom_writer *w) {

  const int n = pmi->palette_size[0];

  uint16_t color_cache[2 * PALETTE_MAX_SIZE];

  const int n_cache = av1_get_palette_cache(xd, 0, color_cache);

  int out_cache_colors[PALETTE_MAX_SIZE];

  uint8_t cache_color_found[2 * PALETTE_MAX_SIZE];

  const int n_out_cache =

      av1_index_color_cache(color_cache, n_cache, pmi->palette_colors, n,

                            cache_color_found, out_cache_colors);

  int n_in_cache = 0;

  for (int i = 0; i < n_cache && n_in_cache < n; ++i) {

    const int found = cache_color_found[i];

    aom_write_bit(w, found);

    n_in_cache += found;

  }

  assert(n_in_cache + n_out_cache == n);

  delta_encode_palette_colors(out_cache_colors, n_out_cache, bit_depth, 1, w);

}

// Write chroma palette color values. U channel is handled similarly to the luma

// channel. For v channel, either use delta encoding or transmit raw values

// directly, whichever costs less.

static inline void write_palette_colors_uv(const MACROBLOCKD *const xd,

                                           const PALETTE_MODE_INFO *const pmi,

                                           int bit_depth, aom_writer *w) {

  const int n = pmi->palette_size[1];

  const uint16_t *colors_u = pmi->palette_colors + PALETTE_MAX_SIZE;

  const uint16_t *colors_v = pmi->palette_colors + 2 * PALETTE_MAX_SIZE;

  // U channel colors.

  uint16_t color_cache[2 * PALETTE_MAX_SIZE];

  const int n_cache = av1_get_palette_cache(xd, 1, color_cache);

  int out_cache_colors[PALETTE_MAX_SIZE];

  uint8_t cache_color_found[2 * PALETTE_MAX_SIZE];

  const int n_out_cache = av1_index_color_cache(

      color_cache, n_cache, colors_u, n, cache_color_found, out_cache_colors);

  int n_in_cache = 0;

  for (int i = 0; i < n_cache && n_in_cache < n; ++i) {

    const int found = cache_color_found[i];

    aom_write_bit(w, found);

    n_in_cache += found;

  }

  delta_encode_palette_colors(out_cache_colors, n_out_cache, bit_depth, 0, w);

  // V channel colors. Don't use color cache as the colors are not sorted.

  const int max_val = 1 << bit_depth;

  int zero_count = 0, min_bits_v = 0;

  int bits_v =

      av1_get_palette_delta_bits_v(pmi, bit_depth, &zero_count, &min_bits_v);

  const int rate_using_delta =

      2 + bit_depth + (bits_v + 1) * (n - 1) - zero_count;

  const int rate_using_raw = bit_depth * n;

  if (rate_using_delta < rate_using_raw) {  // delta encoding

    assert(colors_v[0] < (1 << bit_depth));

    aom_write_bit(w, 1);

    aom_write_literal(w, bits_v - min_bits_v, 2);

    aom_write_literal(w, colors_v[0], bit_depth);

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

      assert(colors_v[i] < (1 << bit_depth));

      if (colors_v[i] == colors_v[i - 1]) {  // No need to signal sign bit.

        aom_write_literal(w, 0, bits_v);

        continue;

      }

      const int delta = abs((int)colors_v[i] - colors_v[i - 1]);

      const int sign_bit = colors_v[i] < colors_v[i - 1];

      if (delta <= max_val - delta) {

        aom_write_literal(w, delta, bits_v);

        aom_write_bit(w, sign_bit);

      } else {

        aom_write_literal(w, max_val - delta, bits_v);

        aom_write_bit(w, !sign_bit);

      }

    }

  } else {  // Transmit raw values.

    aom_write_bit(w, 0);

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

      assert(colors_v[i] < (1 << bit_depth));

      aom_write_literal(w, colors_v[i], bit_depth);

    }

  }

}

static inline void write_palette_mode_info(const AV1_COMMON *cm,

                                           const MACROBLOCKD *xd,

                                           const MB_MODE_INFO *const mbmi,

                                           aom_writer *w) {

  const int num_planes = av1_num_planes(cm);

  const BLOCK_SIZE bsize = mbmi->bsize;

  assert(av1_allow_palette(cm->features.allow_screen_content_tools, bsize));

  const PALETTE_MODE_INFO *const pmi = &mbmi->palette_mode_info;

  const int bsize_ctx = av1_get_palette_bsize_ctx(bsize);

  if (mbmi->mode == DC_PRED) {

    const int n = pmi->palette_size[0];

    const int palette_y_mode_ctx = av1_get_palette_mode_ctx(xd);

    aom_write_symbol(

        w, n > 0,

        xd->tile_ctx->palette_y_mode_cdf[bsize_ctx][palette_y_mode_ctx], 2);

    if (n > 0) {

      aom_write_symbol(w, n - PALETTE_MIN_SIZE,

                       xd->tile_ctx->palette_y_size_cdf[bsize_ctx],

                       PALETTE_SIZES);

      write_palette_colors_y(xd, pmi, cm->seq_params->bit_depth, w);

    }

  }

  const int uv_dc_pred =

      num_planes > 1 && mbmi->uv_mode == UV_DC_PRED && xd->is_chroma_ref;

  if (uv_dc_pred) {

    const int n = pmi->palette_size[1];

    const int palette_uv_mode_ctx = (pmi->palette_size[0] > 0);

    aom_write_symbol(w, n > 0,

                     xd->tile_ctx->palette_uv_mode_cdf[palette_uv_mode_ctx], 2);

    if (n > 0) {

      aom_write_symbol(w, n - PALETTE_MIN_SIZE,

                       xd->tile_ctx->palette_uv_size_cdf[bsize_ctx],

                       PALETTE_SIZES);

      write_palette_colors_uv(xd, pmi, cm->seq_params->bit_depth, w);

    }

  }

}

void av1_write_tx_type(const AV1_COMMON *const cm, const MACROBLOCKD *xd,

                       TX_TYPE tx_type, TX_SIZE tx_size, aom_writer *w) {

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

  const FeatureFlags *const features = &cm->features;

  const int is_inter = is_inter_block(mbmi);

  if (get_ext_tx_types(tx_size, is_inter, features->reduced_tx_set_used) > 1 &&

      ((!cm->seg.enabled && cm->quant_params.base_qindex > 0) ||

       (cm->seg.enabled && xd->qindex[mbmi->segment_id] > 0)) &&

      !mbmi->skip_txfm &&

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

    FRAME_CONTEXT *ec_ctx = xd->tile_ctx;

    const TX_SIZE square_tx_size = txsize_sqr_map[tx_size];

    const TxSetType tx_set_type = av1_get_ext_tx_set_type(

        tx_size, is_inter, features->reduced_tx_set_used);

    const int eset =

        get_ext_tx_set(tx_size, is_inter, features->reduced_tx_set_used);

    // eset == 0 should correspond to a set with only DCT_DCT and there

    // is no need to send the tx_type

    assert(eset > 0);

    assert(av1_ext_tx_used[tx_set_type][tx_type]);

    if (is_inter) {

      aom_write_symbol(w, av1_ext_tx_ind[tx_set_type][tx_type],

                       ec_ctx->inter_ext_tx_cdf[eset][square_tx_size],

                       av1_num_ext_tx_set[tx_set_type]);

    } else {

      PREDICTION_MODE intra_dir;

      if (mbmi->filter_intra_mode_info.use_filter_intra)

        intra_dir =

            fimode_to_intradir[mbmi->filter_intra_mode_info.filter_intra_mode];

      else

        intra_dir = mbmi->mode;

      aom_write_symbol(

          w, av1_ext_tx_ind[tx_set_type][tx_type],

          ec_ctx->intra_ext_tx_cdf[eset][square_tx_size][intra_dir],

          av1_num_ext_tx_set[tx_set_type]);

    }

  }

}

static inline void write_intra_y_mode_nonkf(FRAME_CONTEXT *frame_ctx,

                                            BLOCK_SIZE bsize,

                                            PREDICTION_MODE mode,

                                            aom_writer *w) {

  aom_write_symbol(w, mode, frame_ctx->y_mode_cdf[size_group_lookup[bsize]],

                   INTRA_MODES);

}

static inline void write_intra_uv_mode(FRAME_CONTEXT *frame_ctx,

                                       UV_PREDICTION_MODE uv_mode,

                                       PREDICTION_MODE y_mode,

                                       CFL_ALLOWED_TYPE cfl_allowed,

                                       aom_writer *w) {

  aom_write_symbol(w, uv_mode, frame_ctx->uv_mode_cdf[cfl_allowed][y_mode],

                   UV_INTRA_MODES - !cfl_allowed);

}

static inline void write_cfl_alphas(FRAME_CONTEXT *const ec_ctx, uint8_t idx,

                                    int8_t joint_sign, aom_writer *w) {

  aom_write_symbol(w, joint_sign, ec_ctx->cfl_sign_cdf, CFL_JOINT_SIGNS);

  // Magnitudes are only signaled for nonzero codes.

  if (CFL_SIGN_U(joint_sign) != CFL_SIGN_ZERO) {

    aom_cdf_prob *cdf_u = ec_ctx->cfl_alpha_cdf[CFL_CONTEXT_U(joint_sign)];