/*

* Copyright(c) 2019 Intel Corporation

* 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 https://www.aomedia.org/license/software-license. 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 https://www.aomedia.org/license/patent-license.

*/

/***************************************

* Includes

***************************************/

#include "rd_cost.h"

#include "mcomp.h" // svt_mv_cost (shared MV-rate helper; replaces local duplicate)

#include "common_utils.h"

#include "aom_dsp_rtcd.h"

#include "svt_log.h"

#include "enc_inter_prediction.h"

#include "full_loop.h"

#include "entropy_coding.h"

#include <assert.h>

#define MV_COST_WEIGHT 108

int svt_aom_get_reference_mode_context_new(const MacroBlockD* xd);

int svt_av1_get_pred_context_uni_comp_ref_p(const MacroBlockD* xd);

int svt_av1_get_pred_context_uni_comp_ref_p1(const MacroBlockD* xd);

int svt_av1_get_pred_context_uni_comp_ref_p2(const MacroBlockD* xd);

int svt_aom_get_comp_reference_type_context_new(const MacroBlockD* xd);

int  svt_aom_get_palette_bsize_ctx(BlockSize bsize);

int  svt_aom_get_palette_mode_ctx(const MacroBlockD* xd);

int  svt_aom_write_uniform_cost(int n, int v);

int  svt_get_palette_cache_y(const MacroBlockD* const xd, uint16_t* cache);

int  svt_av1_palette_color_cost_y(const PaletteModeInfo* const pmi, uint16_t* color_cache, const int palette_size,

                                  int n_cache, int bit_depth);

int  svt_av1_cost_color_map(ModeDecisionCandidate* cand, MdRateEstimationContext* rate_table,

                            BlkStruct* blk_ptr, int plane, BlockSize bsize, COLOR_MAP_TYPE type);

void svt_aom_get_block_dimensions(BlockSize bsize, int plane, const MacroBlockD* xd, int* width, int* height,

                                  int* rows_within_bounds, int* cols_within_bounds);

int  svt_aom_allow_palette(int allow_screen_content_tools, BlockSize bsize);

int  svt_aom_allow_intrabc(const FrameHeader* frm_hdr, SliceType slice_type);

int32_t svt_av1_mv_bit_cost_light(const Mv mv, const Mv ref) {

    const uint32_t factor     = 50;

    const uint32_t absmvdiffx = ABS(mv.x - ref.x);

    const uint32_t absmvdiffy = ABS(mv.y - ref.y);

    const uint32_t mv_rate    = 1296 + (factor * (absmvdiffx + absmvdiffy));

    return mv_rate;

}

int32_t svt_av1_mv_bit_cost(const Mv mv, const Mv ref, const int32_t* mvjcost, const int32_t* const mvcost[2],

                            int32_t weight) {

    // Restrict the size of the MV diff to be within the max AV1 range.  If the MV diff

    // is outside this range, the diff will index beyond the cost array, causing a seg fault.

    // Both the MVs and the MV diffs should be within the allowable range for accessing the MV cost

    // infrastructure.

    const int16_t x         = MIN(MAX(mv.x - ref.x, MV_LOW), MV_UPP);

    const int16_t y         = MIN(MAX(mv.y - ref.y, MV_LOW), MV_UPP);

    Mv            temp_diff = {{x, y}};

    return ROUND_POWER_OF_TWO(svt_mv_cost(temp_diff, mvjcost, mvcost) * weight, 7);

}

/////////////////////////////COEFFICIENT CALCULATION //////////////////////////////////////////////

static INLINE int32_t get_golomb_cost(int32_t abs_qc) {

    if (abs_qc >= 1 + NUM_BASE_LEVELS + COEFF_BASE_RANGE) {

        const int32_t r      = abs_qc - COEFF_BASE_RANGE - NUM_BASE_LEVELS;

        const int32_t length = get_msb(r) + 1;

        return av1_cost_literal(2 * length - 1);

    }

    return 0;

}

void svt_av1_txb_init_levels_c(const TranLow* const coeff, const int32_t width, const int32_t height,

                               uint8_t* const levels) {

    uint8_t* ls = levels;

    for (int32_t i = 0; i < height; i++) {

        for (int32_t j = 0; j < width; j++) {

            *ls++ = (uint8_t)clamp(abs(coeff[i * width + j]), 0, INT8_MAX);

        }

        for (int32_t j = 0; j < TX_PAD_HOR; j++) {

            *ls++ = 0;

        }

    }

}

static int32_t av1_transform_type_rate_estimation(ModeDecisionContext* ctx, uint8_t allow_update_cdf, FRAME_CONTEXT* fc,

                                                  ModeDecisionCandidateBuffer* cand_bf, bool is_inter,

                                                  TxSize transform_size, TxType transform_type,

                                                  bool reduced_tx_set_used) {

    // const MbModeInfo *mbmi = &xd->mi[0]->mbmi;

    // const int32_t is_inter = is_inter_block(mbmi);

    if (get_ext_tx_types(transform_size, is_inter, reduced_tx_set_used) >

        1 /*&&    !xd->lossless[xd->mi[0]->mbmi.segment_id]  WE ARE NOT LOSSLESS*/) {

        const TxSize square_tx_size = txsize_sqr_map[transform_size];

        assert(square_tx_size < EXT_TX_SIZES);

        const int32_t ext_tx_set = get_ext_tx_set(transform_size, is_inter, reduced_tx_set_used);

        if (is_inter) {

            if (ext_tx_set > 0) {

                if (allow_update_cdf) {

                    const TxSetType tx_set_type = get_ext_tx_set_type(transform_size, is_inter, reduced_tx_set_used);

                    update_cdf(fc->inter_ext_tx_cdf[ext_tx_set][square_tx_size],

                               av1_ext_tx_ind[tx_set_type][transform_type],

                               av1_num_ext_tx_set[tx_set_type]);

                }

                return ctx->md_rate_est_ctx->inter_tx_type_fac_bits[ext_tx_set][square_tx_size][transform_type];

            }

        } else {

            if (ext_tx_set > 0) {

                PredictionMode intra_dir;

                if (cand_bf->cand->block_mi.filter_intra_mode != FILTER_INTRA_MODES) {

                    intra_dir = fimode_to_intradir[cand_bf->cand->block_mi.filter_intra_mode];

                } else {

                    intra_dir = cand_bf->cand->block_mi.mode;

                }

                assert(intra_dir < INTRA_MODES);

                const TxSetType tx_set_type = get_ext_tx_set_type(transform_size, is_inter, reduced_tx_set_used);

                if (allow_update_cdf) {

                    update_cdf(fc->intra_ext_tx_cdf[ext_tx_set][square_tx_size][intra_dir],

                               av1_ext_tx_ind[tx_set_type][transform_type],

                               av1_num_ext_tx_set[tx_set_type]);

                }

                return ctx->md_rate_est_ctx

                    ->intra_tx_type_fac_bits[ext_tx_set][square_tx_size][intra_dir][transform_type];

            }

        }

    }

    return 0;

}

// Update the eob-related CDFs. Function assumes allow_update_cdf is true

// as the only action of the function is to update the CDFs.

static void update_eob_context(int eob, TxSize tx_size, TxClass tx_class, PlaneType plane, FRAME_CONTEXT* ec_ctx) {

    int          eob_extra;

    const int    eob_pt  = get_eob_pos_token(eob, &eob_extra);

    const TxSize txs_ctx = (TxSize)((txsize_sqr_map[tx_size] + txsize_sqr_up_map[tx_size] + 1) >> 1);

    assert(txs_ctx < TX_SIZES);

    const int eob_multi_size = txsize_log2_minus4[tx_size];

    const int eob_multi_ctx  = (tx_class == TX_CLASS_2D) ? 0 : 1;

    switch (eob_multi_size) {

    case 0:

        update_cdf(ec_ctx->eob_flag_cdf16[plane][eob_multi_ctx], eob_pt - 1, 5);

        break;

    case 1:

        update_cdf(ec_ctx->eob_flag_cdf32[plane][eob_multi_ctx], eob_pt - 1, 6);

        break;

    case 2:

        update_cdf(ec_ctx->eob_flag_cdf64[plane][eob_multi_ctx], eob_pt - 1, 7);

        break;

    case 3:

        update_cdf(ec_ctx->eob_flag_cdf128[plane][eob_multi_ctx], eob_pt - 1, 8);

        break;

    case 4:

        update_cdf(ec_ctx->eob_flag_cdf256[plane][eob_multi_ctx], eob_pt - 1, 9);

        break;

    case 5:

        update_cdf(ec_ctx->eob_flag_cdf512[plane][eob_multi_ctx], eob_pt - 1, 10);

        break;

    case 6:

    default:

        update_cdf(ec_ctx->eob_flag_cdf1024[plane][eob_multi_ctx], eob_pt - 1, 11);

        break;

    }

    if (eob_pt > 2) {

        const int cnt = eob_pt - 3;

        const int bit = (eob_extra >> cnt) & 1;

        update_cdf(ec_ctx->eob_extra_cdf[txs_ctx][plane][cnt], bit, 2);

    }

}

// Transform end of block bit estimation

int get_eob_cost(int eob, const LvMapEobCost* txb_eob_costs, const LvMapCoeffCost* txb_costs, TxClass tx_class) {

    int       eob_extra;

    const int eob_pt        = get_eob_pos_token(eob, &eob_extra);

    const int eob_multi_ctx = (tx_class == TX_CLASS_2D) ? 0 : 1;

    int       eob_cost      = txb_eob_costs->eob_cost[eob_multi_ctx][eob_pt - 1];

    if (eob_pt > 2) {

        const int cnt = eob_pt - 3;

        const int bit = (eob_extra >> cnt) & 1;

        eob_cost += txb_costs->eob_extra_cost[cnt][bit];

        eob_cost += av1_cost_literal(cnt);

    }

    return eob_cost;

}

static INLINE int32_t av1_cost_skip_txb(ModeDecisionContext* ctx, uint8_t allow_update_cdf, FRAME_CONTEXT* ec_ctx,

                                        TxSize transform_size, PlaneType plane_type, int16_t txb_skip_ctx) {

    const TxSize txs_ctx = (TxSize)((txsize_sqr_map[transform_size] + txsize_sqr_up_map[transform_size] + 1) >> 1);

    assert(txs_ctx < TX_SIZES);

    const LvMapCoeffCost* const coeff_costs = &ctx->md_rate_est_ctx->coeff_fac_bits[txs_ctx][plane_type];

    if (allow_update_cdf) {

        update_cdf(ec_ctx->txb_skip_cdf[txs_ctx][txb_skip_ctx], 1, 2);

    }

    return coeff_costs->txb_skip_cost[txb_skip_ctx][1];

}

static INLINE int32_t av1_cost_coeffs_txb_loop_cost_one_eob(const TranLow* const qcoeff, int8_t* const coeff_contexts,

                                                            const LvMapCoeffCost* coeff_costs, int16_t dc_sign_ctx) {

    const TranLow v         = qcoeff[0];

    const int32_t level     = abs(v);

    const int32_t coeff_ctx = coeff_contexts[0];

    assert((AOMMIN(level, 3) - 1) >= 0);

    int32_t cost = coeff_costs->base_eob_cost[coeff_ctx][AOMMIN(level, 3) - 1];

    if (v != 0) {

        const int32_t sign = (v < 0) ? 1 : 0;

        // sign bit cost

        cost += coeff_costs->dc_sign_cost[dc_sign_ctx][sign];

        if (level > NUM_BASE_LEVELS) {

            const int32_t base_range = level - 1 - NUM_BASE_LEVELS;

            if (base_range < COEFF_BASE_RANGE) {

                cost += coeff_costs->lps_cost[0][base_range];

            } else {

                cost += coeff_costs->lps_cost[0][COEFF_BASE_RANGE];

            }

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

                cost += get_golomb_cost(level);

            }

        }

    }

    return cost;

}

static INLINE int32_t av1_cost_coeffs_txb_loop_cost_eob(ModeDecisionContext* md_ctx, uint16_t eob,

                                                        const int16_t* const scan, const TranLow* const qcoeff,

                                                        int8_t* const coeff_contexts, const LvMapCoeffCost* coeff_costs,

                                                        int16_t dc_sign_ctx, uint8_t* const levels, const int32_t bwl,

                                                        TxType transform_type) {

    const uint32_t cost_literal = av1_cost_literal(1);

    int32_t        cost         = 0;

    //Optimized/simplified function when eob is 1

    if (eob == 1) {

        return av1_cost_coeffs_txb_loop_cost_one_eob(qcoeff, coeff_contexts, coeff_costs, dc_sign_ctx);

    }

    //  first (eob - 1) index

    {

        const int32_t pos       = scan[eob - 1];

        const TranLow v         = qcoeff[pos];

        const int32_t level     = abs(v);

        const int32_t coeff_ctx = coeff_contexts[pos];

        assert((AOMMIN(level, 3) - 1) >= 0);

        cost += coeff_costs->base_eob_cost[coeff_ctx][AOMMIN(level, 3) - 1];

        if (v != 0) {

            cost += cost_literal;

            if (level > NUM_BASE_LEVELS) {

                int32_t       ctx        = get_br_ctx(levels, pos, bwl, tx_type_to_class[transform_type]);

                const int32_t base_range = level - 1 - NUM_BASE_LEVELS;

                if (base_range < COEFF_BASE_RANGE) {

                    cost += coeff_costs->lps_cost[ctx][base_range];

                } else {

                    cost += coeff_costs->lps_cost[ctx][COEFF_BASE_RANGE];

                }

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

                    cost += get_golomb_cost(level);

                }

            }

        }

    }

    // last (0) index

    {

        const TranLow v         = qcoeff[0];

        const int32_t level     = abs(v);

        const int32_t coeff_ctx = coeff_contexts[0];

        cost += coeff_costs->base_cost[coeff_ctx][AOMMIN(level, 3)];

        if (v != 0) {

            const int32_t sign = (v < 0) ? 1 : 0;

            // sign bit cost

            cost += coeff_costs->dc_sign_cost[dc_sign_ctx][sign];

            if (level > NUM_BASE_LEVELS) {

                int32_t       ctx        = get_br_ctx(levels, 0, bwl, tx_type_to_class[transform_type]);

                const int32_t base_range = level - 1 - NUM_BASE_LEVELS;

                if (base_range < COEFF_BASE_RANGE) {

                    cost += coeff_costs->lps_cost[ctx][base_range];

                } else {

                    cost += coeff_costs->lps_cost[ctx][COEFF_BASE_RANGE];

                }

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

                    cost += get_golomb_cost(level);

                }

            }

        }

    }

    int32_t c;

    /* Optimized Loop, omitted first (eob - 1) and last (0) index */

    // Estimate the rate of the first(eob / fast_coeff_est_level) coeff(s), DC and last coeff only

    int32_t  c_start = MIN(eob - 2, eob / MAX(1, (int)(md_ctx->mds_fast_coeff_est_level - md_ctx->mds_subres_step)));

    uint32_t cost_literal_cnt = 0;

    for (c = c_start; c >= 1; --c) {

        const int32_t pos = scan[c];

        cost_literal_cnt += !!(qcoeff[pos]);

        const int32_t level = abs(qcoeff[pos]);

        if (level > NUM_BASE_LEVELS) {

            int32_t       ctx        = get_br_ctx(levels, pos, bwl, tx_type_to_class[transform_type]);

            const int32_t base_range = level - 1 - NUM_BASE_LEVELS;

            cost += coeff_costs->base_cost[coeff_contexts[pos]][3];

            if (base_range < COEFF_BASE_RANGE) {

                cost += coeff_costs->lps_cost[ctx][base_range];

            } else {

                cost += get_golomb_cost(level) + coeff_costs->lps_cost[ctx][COEFF_BASE_RANGE];

            }

        } else {

            cost += coeff_costs->base_cost[coeff_contexts[pos]][level];

        }

    }

    cost += cost_literal_cnt * cost_literal;

    return cost;

}

// Note: don't call this function when eob is 0.

uint64_t svt_av1_cost_coeffs_txb(ModeDecisionContext* ctx, uint8_t allow_update_cdf, FRAME_CONTEXT* ec_ctx,

                                 ModeDecisionCandidateBuffer* cand_bf, const TranLow* const qcoeff, uint16_t eob,

                                 PlaneType plane_type, TxSize transform_size, TxType transform_type,

                                 int16_t txb_skip_ctx, int16_t dc_sign_ctx, bool reduced_transform_set_flag)

{

    //Note: there is a different version of this function in AOM that seems to be efficient as its name is:

    //warehouse_efficients_txb

    const TxSize  txs_ctx  = get_txsize_entropy_ctx(transform_size);

    const TxClass tx_class = tx_type_to_class[transform_type];

    int32_t       cost;

    const int32_t bwl    = get_txb_bwl(transform_size);

    const int32_t width  = get_txb_wide(transform_size);

    const int32_t height = get_txb_high(transform_size);

    const ScanOrder* const scan_order     = get_scan_order(transform_size, transform_type);

    const int16_t* const   scan           = scan_order->scan;

    uint8_t* const         levels         = set_levels(ctx->md_levels_buf, width, height);

    int8_t* const          coeff_contexts = ctx->md_coeff_contexts;

    assert(txs_ctx < TX_SIZES);

    const LvMapCoeffCost* const coeff_costs = &ctx->md_rate_est_ctx->coeff_fac_bits[txs_ctx][plane_type];

    const int32_t             eob_multi_size = txsize_log2_minus4[transform_size];

    const LvMapEobCost* const eob_bits       = &ctx->md_rate_est_ctx->eob_frac_bits[eob_multi_size][plane_type];

    // eob must be greater than 0 here.

    assert(eob > 0);

    cost = coeff_costs->txb_skip_cost[txb_skip_ctx][0];

    if (allow_update_cdf) {

        update_cdf(ec_ctx->txb_skip_cdf[txs_ctx][txb_skip_ctx], eob == 0, 2);

    }

    if (eob > 1) {

        svt_av1_txb_init_levels(qcoeff,

                                width,

                                height,

                                levels); // NM - Needs to be optimized - to be combined with the quantisation.

    }

    const bool is_inter = is_inter_mode(cand_bf->cand->block_mi.mode);

    // Transform type bit estimation

    cost += plane_type > PLANE_TYPE_Y ? 0

                                      : av1_transform_type_rate_estimation(ctx,

                                                                           allow_update_cdf,

                                                                           ec_ctx,

                                                                           cand_bf,

                                                                           is_inter,

                                                                           transform_size,

                                                                           transform_type,

                                                                           reduced_transform_set_flag);

    // Transform eob bit estimation

    cost += get_eob_cost(eob, eob_bits, coeff_costs, tx_class);

    if (allow_update_cdf) {

        update_eob_context(eob, transform_size, tx_class, plane_type, ec_ctx);

    }

    // Transform non-zero coeff bit estimation

    svt_av1_get_nz_map_contexts(levels,

                                scan,

                                eob,

                                transform_size,

                                tx_class,

                                coeff_contexts); // NM - Assembly version is available in AOM

    assert(eob <= width * height);

    if (allow_update_cdf) {

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

            const int     pos       = scan[c];

            const int     coeff_ctx = coeff_contexts[pos];

            const TranLow v         = qcoeff[pos];

            const TranLow level     = abs(v);

            if (c == eob - 1) {

                assert(coeff_ctx < 4);

                update_cdf(ec_ctx->coeff_base_eob_cdf[txs_ctx][plane_type][coeff_ctx], AOMMIN(level, 3) - 1, 3);

            } else {

                update_cdf(ec_ctx->coeff_base_cdf[txs_ctx][plane_type][coeff_ctx], AOMMIN(level, 3), 4);

            }

            {

                if (c == eob - 1) {

                    assert(coeff_ctx < 4);

                }

            }

            if (level > NUM_BASE_LEVELS) {

                const int base_range = level - 1 - NUM_BASE_LEVELS;

                int       br_ctx;

                if (eob == 1) {

                    br_ctx = 0;

                } else {

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

                }

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

                    const int k = AOMMIN(base_range - idx, BR_CDF_SIZE - 1);

                    update_cdf(ec_ctx->coeff_br_cdf[AOMMIN(txs_ctx, TX_32X32)][plane_type][br_ctx], k, BR_CDF_SIZE);

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

                        if (lps == k) {

                            break;

                        }

                    }

                    if (k < BR_CDF_SIZE - 1) {

                        break;

                    }

                }

            }

        }

        if (qcoeff[0] != 0) {

            update_cdf(ec_ctx->dc_sign_cdf[plane_type][dc_sign_ctx], qcoeff[0] < 0, 2);

        }

        //TODO: CHKN  for 128x128 where we need more than one TXb, we need to update the txb_context(dc_sign+skip_ctx) in a Txb basis.

        return 0;

    }

    cost += av1_cost_coeffs_txb_loop_cost_eob(

        ctx, eob, scan, qcoeff, coeff_contexts, coeff_costs, dc_sign_ctx, levels, bwl, transform_type);

    return cost;

}

uint64_t svt_aom_get_intra_uv_fast_rate(PictureControlSet* pcs, ModeDecisionContext* ctx,

                                        ModeDecisionCandidateBuffer* cand_bf, bool use_accurate_cfl) {

    const BlockGeom* const blk_geom = ctx->blk_geom;

    ModeDecisionCandidate* cand     = cand_bf->cand;

    assert(ctx->has_uv);

    assert(!(svt_aom_allow_intrabc(&pcs->ppcs->frm_hdr, pcs->ppcs->slice_type) && cand->block_mi.use_intrabc));

    MdRateEstimationContext* md_rate_est_ctx = ctx->md_rate_est_ctx;

    const uint8_t            is_cfl_allowed  = (blk_geom->bwidth <= 32 && blk_geom->bheight <= 32) ? 1 : 0;

    PredictionMode           intra_mode      = (PredictionMode)cand->block_mi.mode;

    // If CFL alphas are not known yet, calculate the chroma mode bits based on DC Mode. If CFL is selected the chroma mode bits must be updated later

    const UvPredictionMode chroma_mode = cand->block_mi.uv_mode == UV_CFL_PRED && !use_accurate_cfl

        ? UV_DC_PRED

        : cand->block_mi.uv_mode;

    const uint32_t         mi_row      = ctx->blk_org_y >> MI_SIZE_LOG2;

    const uint32_t         mi_col      = ctx->blk_org_x >> MI_SIZE_LOG2;

    // Subsampling assumes YUV 420 content

    const uint8_t ss_x = 1;

    const uint8_t ss_y = 1;

    uint64_t chroma_rate = 0;

    // Estimate chroma nominal intra mode bits

    chroma_rate += (uint64_t)md_rate_est_ctx->intra_uv_mode_fac_bits[is_cfl_allowed][intra_mode][chroma_mode];

    // Estimate chroma angular mode bits; angular offset only allow for bsize >= 8x8

    if (blk_geom->bsize >= BLOCK_8X8 && av1_is_directional_mode(get_uv_mode(chroma_mode))) {

        chroma_rate +=

            md_rate_est_ctx->angle_delta_fac_bits[chroma_mode - V_PRED]

                                                 [MAX_ANGLE_DELTA + cand->block_mi.angle_delta[PLANE_TYPE_UV]];

    }

    // Estimate CFL factor bits when CFL is used

    if (chroma_mode == UV_CFL_PRED) {

        chroma_rate += (uint64_t)md_rate_est_ctx->cfl_alpha_fac_bits[cand->block_mi.cfl_alpha_signs][CFL_PRED_U]

                                                                    [CFL_IDX_U(cand->block_mi.cfl_alpha_idx)] +

            (uint64_t)md_rate_est_ctx->cfl_alpha_fac_bits[cand->block_mi.cfl_alpha_signs][CFL_PRED_V]

                                                         [CFL_IDX_V(cand->block_mi.cfl_alpha_idx)];

    }

    // Estimate chroma palette mode bits (currently not supported, so just cost of signalling off)

    if (chroma_mode == UV_DC_PRED &&

        svt_aom_allow_palette(pcs->ppcs->frm_hdr.allow_screen_content_tools, blk_geom->bsize) &&

        is_chroma_reference(mi_row, mi_col, blk_geom->bsize, ss_x, ss_y)) {

        const int use_palette_y  = cand->palette_info && (cand->palette_size[0] > 0);

        const int use_palette_uv = cand->palette_info && (cand->palette_size[1] > 0);

        chroma_rate += ctx->md_rate_est_ctx->palette_uv_mode_fac_bits[use_palette_y][use_palette_uv];

    }

    return chroma_rate;

}

uint64_t svt_aom_intra_fast_cost(PictureControlSet* pcs, ModeDecisionContext* ctx, ModeDecisionCandidateBuffer* cand_bf,

                                 uint64_t lambda, uint64_t luma_distortion) {

    const BlockGeom*       blk_geom = ctx->blk_geom;

    BlkStruct*             blk_ptr  = ctx->blk_ptr;

    ModeDecisionCandidate* cand     = cand_bf->cand;

    if (svt_aom_allow_intrabc(&pcs->ppcs->frm_hdr, pcs->ppcs->slice_type) && cand->block_mi.use_intrabc) {

        uint64_t rate = 0;

        Mv         mv        = {.as_int = cand->block_mi.mv[0].as_int};

        Mv         ref_mv    = {.as_int = cand->pred_mv[0].as_int};

        const int* dvcost[2] = {(int*)&ctx->md_rate_est_ctx->dv_cost[0][MV_MAX],

                                (int*)&ctx->md_rate_est_ctx->dv_cost[1][MV_MAX]};

        int32_t    mv_rate   = svt_av1_mv_bit_cost(

            mv, ref_mv, ctx->md_rate_est_ctx->dv_joint_cost, dvcost, MV_COST_WEIGHT_SUB);

        rate                      = mv_rate + ctx->md_rate_est_ctx->intrabc_fac_bits[cand->block_mi.use_intrabc];

        cand_bf->fast_luma_rate   = rate;

        cand_bf->fast_chroma_rate = 0;

        return (RDCOST(lambda, rate, luma_distortion));

    } else {

        // Number of bits for each synatax element

        uint64_t       intra_mode_bits_num          = 0;

        uint64_t       intra_luma_mode_bits_num     = 0;

        uint64_t       intra_luma_ang_mode_bits_num = 0;

        uint64_t       intra_filter_mode_bits_num   = 0;

        uint64_t       skip_mode_rate               = 0;

        const uint8_t  skip_mode_ctx                = ctx->skip_mode_ctx;

        PredictionMode intra_mode                   = (PredictionMode)cand->block_mi.mode;

        // Luma and chroma rate

        uint32_t rate;

        uint32_t luma_rate   = 0;

        uint32_t chroma_rate = 0;

        intra_mode_bits_num  = pcs->slice_type != I_SLICE

             ? (uint64_t)ctx->md_rate_est_ctx->mb_mode_fac_bits[eb_size_group_lookup[blk_geom->bsize]][intra_mode]

             : ZERO_COST;

        skip_mode_rate = pcs->slice_type != I_SLICE && pcs->ppcs->frm_hdr.skip_mode_params.skip_mode_flag &&

                is_comp_ref_allowed(blk_geom->bsize)

            ? (uint64_t)ctx->md_rate_est_ctx->skip_mode_fac_bits[skip_mode_ctx][0]

            : ZERO_COST;

        // Estimate luma nominal intra mode bits for key frame

        intra_luma_mode_bits_num = pcs->slice_type == I_SLICE

            ? (uint64_t)

                  ctx->md_rate_est_ctx->y_mode_fac_bits[ctx->intra_luma_top_ctx][ctx->intra_luma_left_ctx][intra_mode]

            : ZERO_COST;

        // Estimate luma angular mode bits

        if (blk_geom->bsize >= BLOCK_8X8 && av1_is_directional_mode(cand->block_mi.mode)) {

            assert((intra_mode - V_PRED) < 8);

            assert((intra_mode - V_PRED) >= 0);

            intra_luma_ang_mode_bits_num =

                ctx->md_rate_est_ctx->angle_delta_fac_bits[intra_mode - V_PRED]

                                                          [MAX_ANGLE_DELTA + cand->block_mi.angle_delta[PLANE_TYPE_Y]];

        }

        if (svt_aom_allow_palette(pcs->ppcs->frm_hdr.allow_screen_content_tools, blk_geom->bsize) &&

            intra_mode == DC_PRED) {

            const int use_palette = cand->palette_info ? (cand->palette_size[0] > 0) : 0;

            const int bsize_ctx   = svt_aom_get_palette_bsize_ctx(blk_geom->bsize);

            const int mode_ctx    = svt_aom_get_palette_mode_ctx(blk_ptr->av1xd);

            intra_luma_mode_bits_num += ctx->md_rate_est_ctx->palette_ymode_fac_bits[bsize_ctx][mode_ctx][use_palette];

            if (use_palette) {

                const uint8_t* const color_map = cand->palette_info->color_idx_map;

                int                  block_width, block_height, rows, cols;

                svt_aom_get_block_dimensions(

                    blk_geom->bsize, 0, blk_ptr->av1xd, &block_width, &block_height, &rows, &cols);

                const int plt_size = cand->palette_size[0];

                int       palette_mode_cost =

                    ctx->md_rate_est_ctx->palette_ysize_fac_bits[bsize_ctx][plt_size - PALETTE_MIN_SIZE] +

                    svt_aom_write_uniform_cost(plt_size, color_map[0]);

                uint16_t  color_cache[2 * PALETTE_MAX_SIZE];

                const int n_cache = svt_get_palette_cache_y(blk_ptr->av1xd, color_cache);

                palette_mode_cost += svt_av1_palette_color_cost_y(&cand->palette_info->pmi,

                                                                  color_cache,

                                                                  cand->palette_size[0],

                                                                  n_cache,

                                                                  pcs->ppcs->scs->encoder_bit_depth);

                palette_mode_cost += svt_av1_cost_color_map(

                    cand, ctx->md_rate_est_ctx, blk_ptr, 0, blk_geom->bsize, PALETTE_MAP);

                intra_luma_mode_bits_num += palette_mode_cost;

            }

        }

        if (svt_aom_filter_intra_allowed(pcs->ppcs->scs->seq_header.filter_intra_level,

                                         blk_geom->bsize,

                                         cand->palette_info ? cand->palette_size[0] : 0,

                                         intra_mode)) {

            intra_filter_mode_bits_num =

                ctx->md_rate_est_ctx

                    ->filter_intra_fac_bits[blk_geom->bsize][cand->block_mi.filter_intra_mode != FILTER_INTRA_MODES];

            if (cand->block_mi.filter_intra_mode != FILTER_INTRA_MODES) {

                intra_filter_mode_bits_num +=

                    ctx->md_rate_est_ctx->filter_intra_mode_fac_bits[cand->block_mi.filter_intra_mode];

            }

        }

        if (ctx->has_uv) {

            // CFL info not known in fasta loop, so assume DC mode when CFL is allowed

            chroma_rate = (uint32_t)svt_aom_get_intra_uv_fast_rate(pcs, ctx, cand_bf, 0);

        }

        uint32_t is_inter_rate = pcs->slice_type != I_SLICE

            ? ctx->md_rate_est_ctx->intra_inter_fac_bits[ctx->is_inter_ctx][0]

            : 0;

        luma_rate              = (uint32_t)(intra_mode_bits_num + skip_mode_rate + intra_luma_mode_bits_num +

                               intra_luma_ang_mode_bits_num + is_inter_rate + intra_filter_mode_bits_num);

        if (svt_aom_allow_intrabc(&pcs->ppcs->frm_hdr, pcs->ppcs->slice_type)) {

            svt_aom_assert_err(cand->block_mi.use_intrabc == 0, "this block ibc should be off\n");

            luma_rate += ctx->md_rate_est_ctx->intrabc_fac_bits[cand->block_mi.use_intrabc];

        }

        // Keep the Fast Luma and Chroma rate for future use

        cand_bf->fast_luma_rate   = luma_rate;

        cand_bf->fast_chroma_rate = chroma_rate;

        rate                      = luma_rate + chroma_rate;

        // Assign fast cost

        return (RDCOST(lambda, rate, luma_distortion));

    }

}

// This function encodes the reference frame

uint64_t estimate_ref_frame_type_bits(ModeDecisionContext* ctx, BlkStruct* blk_ptr, uint8_t ref_frame_type,

                                      bool is_compound) {

    uint64_t ref_rate_bits = 0;

    MbModeInfo* const mbmi = blk_ptr->av1xd->mi[0];

    MvReferenceFrame  ref_type[2];

    av1_set_ref_frame(ref_type, ref_frame_type);

    mbmi->block_mi.ref_frame[0] = ref_type[0];

    mbmi->block_mi.ref_frame[1] = ref_type[1];

    //const int is_compound = svt_aom_has_second_ref(mbmi);

    {

        if (is_compound) {

            const CompReferenceType comp_ref_type = has_uni_comp_refs(&mbmi->block_mi) ? UNIDIR_COMP_REFERENCE

                                                                                       : BIDIR_COMP_REFERENCE;

            ref_rate_bits += ctx->md_rate_est_ctx->comp_ref_type_fac_bits[svt_aom_get_comp_reference_type_context_new(

                blk_ptr->av1xd)][comp_ref_type];

            /*aom_write_symbol(w, comp_ref_type,

               svt_aom_get_comp_reference_type_cdf(blk_ptr->av1xd), 2);*/

            if (comp_ref_type == UNIDIR_COMP_REFERENCE) {

                // SVT_LOG("ERROR[AN]: UNIDIR_COMP_REFERENCE not supported\n");

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

                ref_rate_bits += ctx->md_rate_est_ctx->uni_comp_ref_fac_bits[svt_av1_get_pred_context_uni_comp_ref_p(

                    blk_ptr->av1xd)][0][bit];

                // blk_ptr->av1xd->tile_ctx->uni_comp_ref_cdf[pred_context][0];

                // WRITE_REF_BIT(bit, uni_comp_ref_p);

                if (!bit) {

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

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

                        mbmi->block_mi.ref_frame[1] == GOLDEN_FRAME;

                    ref_rate_bits +=

                        ctx->md_rate_est_ctx

                            ->uni_comp_ref_fac_bits[svt_av1_get_pred_context_uni_comp_ref_p1(blk_ptr->av1xd)][1][bit1];

                    // ref_rate_d = blk_ptr->av1xd->tile_ctx->uni_comp_ref_cdf[pred_context][1];

                    // WRITE_REF_BIT(bit1, uni_comp_ref_p1);

                    if (bit1) {

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

                        ref_rate_bits +=

                            ctx->md_rate_est_ctx->uni_comp_ref_fac_bits[svt_av1_get_pred_context_uni_comp_ref_p2(

                                blk_ptr->av1xd)][2][bit2];

                        // ref_rate_e = blk_ptr->av1xd->tile_ctx->uni_comp_ref_cdf[pred_context][2];

                        //WRITE_REF_BIT(bit2, uni_comp_ref_p2);

                    }

                }

                return ref_rate_bits;

            }

            assert(comp_ref_type == BIDIR_COMP_REFERENCE);

            const int bit = (mbmi->block_mi.ref_frame[0] == GOLDEN_FRAME || mbmi->block_mi.ref_frame[0] == LAST3_FRAME);

            const int pred_ctx = svt_av1_get_pred_context_comp_ref_p(blk_ptr->av1xd);

            ref_rate_bits += ctx->md_rate_est_ctx->comp_ref_fac_bits[pred_ctx][0][bit];

            // ref_rate_f = blk_ptr->av1xd->tile_ctx->comp_ref_cdf[pred_ctx][0];

            // WRITE_REF_BIT(bit, comp_ref_p);

            if (!bit) {

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

                ref_rate_bits += ctx->md_rate_est_ctx

                                     ->comp_ref_fac_bits[svt_av1_get_pred_context_comp_ref_p1(blk_ptr->av1xd)][1][bit1];

                // ref_rate_g = blk_ptr->av1xd->tile_ctx->comp_ref_cdf[pred_context][1];

                // WRITE_REF_BIT(bit1, comp_ref_p1);

            } else {

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

                ref_rate_bits += ctx->md_rate_est_ctx

                                     ->comp_ref_fac_bits[svt_av1_get_pred_context_comp_ref_p2(blk_ptr->av1xd)][2][bit2];

                // ref_rate_h = blk_ptr->av1xd->tile_ctx->comp_ref_cdf[pred_context][2];

                // WRITE_REF_BIT(bit2, comp_ref_p2);

            }

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

            const int pred_ctx_2 = svt_av1_get_pred_context_comp_bwdref_p(blk_ptr->av1xd);

            ref_rate_bits += ctx->md_rate_est_ctx->comp_bwd_ref_fac_bits[pred_ctx_2][0][bit_bwd];

            // ref_rate_i = blk_ptr->av1xd->tile_ctx->comp_bwdref_cdf[pred_ctx_2][0];

            // WRITE_REF_BIT(bit_bwd, comp_bwdref_p);

            if (!bit_bwd) {

                ref_rate_bits += ctx->md_rate_est_ctx->comp_bwd_ref_fac_bits[svt_av1_get_pred_context_comp_bwdref_p1(

                    blk_ptr->av1xd)][1][ref_type[1] == ALTREF2_FRAME];

                // ref_rate_j = blk_ptr->av1xd->tile_ctx->comp_bwdref_cdf[pred_context][1];

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

            }

        } else {

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

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

            ref_rate_bits += ctx->md_rate_est_ctx

                                 ->single_ref_fac_bits[svt_av1_get_pred_context_single_ref_p1(blk_ptr->av1xd)][0][bit0];

            // ref_rate_k =

            // blk_ptr->av1xd->tile_ctx->single_ref_cdf[svt_av1_get_pred_context_single_ref_p1(blk_ptr->av1xd)][0];

            // WRITE_REF_BIT(bit0, single_ref_p1);

            if (bit0) {

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

                ref_rate_bits += ctx->md_rate_est_ctx->single_ref_fac_bits[svt_av1_get_pred_context_single_ref_p2(

                    blk_ptr->av1xd)][1][bit1];

                // ref_rate_l =

                // blk_ptr->av1xd->tile_ctx->single_ref_cdf[svt_av1_get_pred_context_single_ref_p2(blk_ptr->av1xd)][1];

                // WRITE_REF_BIT(bit1, single_ref_p2);

                if (!bit1) {

                    ref_rate_bits += ctx->md_rate_est_ctx->single_ref_fac_bits[svt_av1_get_pred_context_single_ref_p6(

                        blk_ptr->av1xd)][5][ref_frame_type == ALTREF2_FRAME];

                    // ref_rate_m =

                    // blk_ptr->av1xd->tile_ctx->single_ref_cdf[svt_av1_get_pred_context_single_ref_p6(blk_ptr->av1xd)][5];

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

                }

            } else {

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

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

                ref_rate_bits += ctx->md_rate_est_ctx->single_ref_fac_bits[svt_av1_get_pred_context_single_ref_p3(

                    blk_ptr->av1xd)][2][bit2];

                // ref_rate_n =

                // blk_ptr->av1xd->tile_ctx->single_ref_cdf[svt_av1_get_pred_context_single_ref_p3(blk_ptr->av1xd)][2];

                // WRITE_REF_BIT(bit2, single_ref_p3);

                if (!bit2) {

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

                    ref_rate_bits += ctx->md_rate_est_ctx->single_ref_fac_bits[svt_av1_get_pred_context_single_ref_p4(

                        blk_ptr->av1xd)][3][bit3];

                    // ref_rate_o =

                    // blk_ptr->av1xd->tile_ctx->single_ref_cdf[svt_av1_get_pred_context_single_ref_p4(blk_ptr->av1xd)][3];

                    // WRITE_REF_BIT(bit3, single_ref_p4);

                } else {

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

                    ref_rate_bits += ctx->md_rate_est_ctx->single_ref_fac_bits[svt_av1_get_pred_context_single_ref_p5(

                        blk_ptr->av1xd)][4][bit4];

                    // ref_rate_p =

                    // blk_ptr->av1xd->tile_ctx->single_ref_cdf[svt_av1_get_pred_context_single_ref_p5(blk_ptr->av1xd)][4];

                    // WRITE_REF_BIT(bit4, single_ref_p5);

                }

            }

        }

    }

    return ref_rate_bits;

}

int svt_aom_get_comp_group_idx_context_enc(const MacroBlockD* xd);

int is_any_masked_compound_used(BlockSize bsize);

static INLINE uint32_t get_compound_mode_rate(PictureControlSet* pcs, ModeDecisionContext* ctx,

                                              ModeDecisionCandidate* cand, BlockSize bsize) {

    BlkStruct*          blk_ptr   = ctx->blk_ptr;

    SequenceControlSet* scs       = pcs->ppcs->scs;

    uint32_t            comp_rate = 0;

    MbModeInfo* const   mbmi      = blk_ptr->av1xd->mi[0];

    MvReferenceFrame    rf[2]     = {cand->block_mi.ref_frame[0], cand->block_mi.ref_frame[1]};

    mbmi->block_mi.ref_frame[0]   = rf[0];

    mbmi->block_mi.ref_frame[1]   = rf[1];

    //NOTE  :  Make sure, any cuPtr data is already set before   usage

    if (has_second_ref(&mbmi->block_mi)) {

        const int masked_compound_used = is_any_masked_compound_used(bsize) && scs->seq_header.enable_masked_compound;

        if (masked_compound_used) {

            const int ctx_comp_group_idx = svt_aom_get_comp_group_idx_context_enc(blk_ptr->av1xd);

            comp_rate =

                ctx->md_rate_est_ctx->comp_group_idx_fac_bits[ctx_comp_group_idx][cand->block_mi.comp_group_idx];

        } else {

            assert(cand->block_mi.comp_group_idx == 0);

        }

        if (cand->block_mi.comp_group_idx == 0) {

            if (cand->block_mi.compound_idx) {

                assert(cand->block_mi.interinter_comp.type == COMPOUND_AVERAGE);

            }

            if (scs->seq_header.order_hint_info.enable_jnt_comp) {

                const int comp_index_ctx = svt_aom_get_comp_index_context_enc(pcs->ppcs,

                                                                              pcs->ppcs->cur_order_hint,

                                                                              pcs->ppcs->ref_order_hint[rf[0] - 1],

                                                                              pcs->ppcs->ref_order_hint[rf[1] - 1],

                                                                              blk_ptr->av1xd);

                comp_rate += ctx->md_rate_est_ctx->comp_idx_fac_bits[comp_index_ctx][cand->block_mi.compound_idx];

            } else {

                assert(cand->block_mi.compound_idx == 1);

            }

        } else {

            assert(pcs->ppcs->frm_hdr.reference_mode != SINGLE_REFERENCE &&

                   is_inter_compound_mode(cand->block_mi.mode));

            assert(masked_compound_used);

            // compound_diffwtd, wedge

            assert(cand->block_mi.interinter_comp.type == COMPOUND_WEDGE ||

                   cand->block_mi.interinter_comp.type == COMPOUND_DIFFWTD);

            if (is_interinter_compound_used(COMPOUND_WEDGE, bsize)) {

                comp_rate += ctx->md_rate_est_ctx

                                 ->compound_type_fac_bits[bsize][cand->block_mi.interinter_comp.type - COMPOUND_WEDGE];

            }

            if (cand->block_mi.interinter_comp.type == COMPOUND_WEDGE) {

                assert(is_interinter_compound_used(COMPOUND_WEDGE, bsize));

                comp_rate +=

                    ctx->md_rate_est_ctx->wedge_idx_fac_bits[bsize][cand->block_mi.interinter_comp.wedge_index];

                comp_rate += av1_cost_literal(1);

            } else {

                assert(cand->block_mi.interinter_comp.type == COMPOUND_DIFFWTD);

                comp_rate += av1_cost_literal(1);

            }

        }

    }

    return comp_rate;

}

int32_t svt_aom_get_switchable_rate(BlockModeInfo* block_mi, const FrameHeader* const frm_hdr, ModeDecisionContext* ctx,

                                    const bool enable_dual_filter) {

    if (frm_hdr->interpolation_filter != SWITCHABLE) {

        return 0;

    }

    int32_t   inter_filter_cost = 0;

    const int max_dir           = enable_dual_filter ? 2 : 1;

    for (int dir = 0; dir < max_dir; ++dir) {

        const int32_t pred_ctx = svt_aom_get_pred_context_switchable_interp(

            block_mi->ref_frame[0], block_mi->ref_frame[1], ctx->blk_ptr->av1xd, dir);

        const InterpFilter filter = av1_extract_interp_filter(block_mi->interp_filters, dir);

        assert(pred_ctx < SWITCHABLE_FILTER_CONTEXTS);

        assert(filter < SWITCHABLE_FILTERS);

        inter_filter_cost += ctx->md_rate_est_ctx->switchable_interp_fac_bitss[pred_ctx][filter];

    }

    return inter_filter_cost;

}

int svt_aom_is_interintra_wedge_used(BlockSize bsize);

static uint64_t av1_inter_fast_cost_light(ModeDecisionContext* ctx, BlkStruct* blk_ptr,

                                          ModeDecisionCandidateBuffer* cand_bf, uint64_t luma_distortion,

                                          uint64_t lambda, PictureControlSet* pcs, CandidateMv* ref_mv_stack) {

    ModeDecisionCandidate* cand = cand_bf->cand;

    // NM - fast inter cost estimation

    MdRateEstimationContext* r = ctx->md_rate_est_ctx;

    //_mm_prefetch(p, _MM_HINT_T2);

    // Luma rate

    uint32_t             luma_rate           = 0;

    uint64_t             mv_rate             = 0;

    const PredictionMode inter_mode          = (PredictionMode)cand->block_mi.mode;

    const uint8_t        have_nearmv         = have_nearmv_in_inter_mode(inter_mode);

    uint64_t             inter_mode_bits_num = 0;

    const uint8_t        skip_mode_ctx       = ctx->skip_mode_ctx;

    MvReferenceFrame     rf[2]               = {cand->block_mi.ref_frame[0], cand->block_mi.ref_frame[1]};

    const int8_t         ref_frame_type      = av1_ref_frame_type(rf);

    const uint8_t        is_compound         = is_inter_compound_mode(cand->block_mi.mode);

    const uint32_t       mode_context        = svt_aom_mode_context_analyzer(ctx->inter_mode_ctx[ref_frame_type], rf);

    uint64_t             reference_picture_bits_num = 0;

    if (ctx->approx_inter_rate < 2) {

        reference_picture_bits_num = ctx->estimate_ref_frames_num_bits[ref_frame_type];

    }

    if (is_compound) {

        assert(INTER_COMPOUND_OFFSET(inter_mode) < INTER_COMPOUND_MODES);

        inter_mode_bits_num += r->inter_compound_mode_fac_bits[mode_context][INTER_COMPOUND_OFFSET(inter_mode)];

    } else {

        int16_t newmv_ctx = mode_context & NEWMV_CTX_MASK;

        //aom_write_symbol(ec_writer, mode != NEWMV, frame_context->newmv_cdf[newmv_ctx], 2);

        inter_mode_bits_num += r->new_mv_mode_fac_bits[newmv_ctx][inter_mode != NEWMV];

        if (inter_mode != NEWMV) {