/*

* Copyright(c) 2019 Intel Corporation

*

* 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.

*/

#include <stdlib.h>

#include "utility.h"

#include "md_process.h"

#include "lambda_rate_tables.h"

#include "rc_process.h"

#include "enc_mode_config.h"

const uint8_t quantizer_to_qindex[64] = {

    0,   4,   8,   12,  16,  20,  24,  28,  32,  36,  40,  44,  48,  52,  56,  60,  64,  68,  72,  76,  80,  84,

    88,  92,  96,  100, 104, 108, 112, 116, 120, 124, 128, 132, 136, 140, 144, 148, 152, 156, 160, 164, 168, 172,

    176, 180, 184, 188, 192, 196, 200, 204, 208, 212, 216, 220, 224, 228, 232, 236, 240, 244, 249, 255};

const int percents[2][FIXED_QP_OFFSET_COUNT] = {

    {75, 70, 60, 20, 15, 0}, {76, 60, 30, 15, 8, 4} // libaom offsets

};

const uint8_t uni_psy_bias[64] = {

    85, 85, 85, 85, 85,  85,  85,  85,  85,  85,  85,  85,  85,  85,  85,  85,  95,  95,  95,  95,  95, 95,

    95, 95, 95, 95, 95,  95,  95,  95,  95,  95,  95,  95,  95,  95,  95,  95,  95,  95,  95,  95,  95, 95,

    95, 95, 95, 95, 100, 100, 100, 100, 100, 100, 100, 100, 100, 100, 100, 100, 100, 100, 100, 100,

};

static void mode_decision_context_dctor(EbPtr p) {

    ModeDecisionContext* obj = (ModeDecisionContext*)p;

    uint32_t block_max_count_sb = obj->init_max_block_cnt;

    // MD palette search

    if (obj->palette_buffer) {

        EB_FREE(obj->palette_buffer);

    }

    if (obj->palette_cand_array) {

        // Free fields in palette_cand_array before freeing palette_cand_array

        for (int cd = 0; cd < MAX_PAL_CAND; cd++) {

            if (obj->palette_cand_array[cd].color_idx_map) {

                EB_FREE_ARRAY(obj->palette_cand_array[cd].color_idx_map);

            }

        }

        EB_FREE_ARRAY(obj->palette_cand_array);

    }

    if (obj->palette_size_array_0) {

        EB_FREE_ARRAY(obj->palette_size_array_0);

    }

    for (CandClass cand_class_it = CAND_CLASS_0; cand_class_it < CAND_CLASS_TOTAL; cand_class_it++) {

        EB_FREE_ARRAY(obj->cand_buff_indices[cand_class_it]);

    }

    EB_FREE_ARRAY(obj->best_candidate_index_array);

    EB_FREE_ARRAY(obj->above_txfm_context);

    EB_FREE_ARRAY(obj->left_txfm_context);

    for (uint32_t coded_leaf_index = 0; coded_leaf_index < block_max_count_sb; ++coded_leaf_index) {

        if (obj->md_blk_arr_nsq[coded_leaf_index].coeff_tmp) {

            EB_DELETE(obj->md_blk_arr_nsq[coded_leaf_index].coeff_tmp);

        }

        if (obj->md_blk_arr_nsq[coded_leaf_index].recon_tmp) {

            EB_DELETE(obj->md_blk_arr_nsq[coded_leaf_index].recon_tmp);

        }

    }

    EB_DELETE_PTR_ARRAY(obj->cand_bf_ptr_array, obj->max_nics_uv);

    EB_FREE_ARRAY(obj->cand_bf_tx_depth_1->cand);

    EB_DELETE(obj->cand_bf_tx_depth_1);

    EB_FREE_ARRAY(obj->cand_bf_tx_depth_2->cand);

    EB_DELETE(obj->cand_bf_tx_depth_2);

    EB_FREE_ALIGNED_ARRAY(obj->cfl_temp_luma_recon16bit);

    EB_FREE_ALIGNED_ARRAY(obj->cfl_temp_luma_recon);

    EB_FREE_ALIGNED_ARRAY(obj->pred_buf_q3);

    EB_FREE_ARRAY(obj->fast_cand_array);

    EB_FREE_2D(obj->injected_mvs);

    EB_FREE_ARRAY(obj->injected_ref_types);

    EB_FREE_ARRAY(obj->fast_cost_array);

    EB_FREE_ARRAY(obj->full_cost_array);

    if (obj->md_blk_arr_nsq) {

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

            EB_FREE_ARRAY(obj->md_blk_arr_nsq[0].neigh_left_recon_16bit[i]);

            EB_FREE_ARRAY(obj->md_blk_arr_nsq[0].neigh_top_recon_16bit[i]);

            EB_FREE_ARRAY(obj->md_blk_arr_nsq[0].neigh_left_recon[i]);

            EB_FREE_ARRAY(obj->md_blk_arr_nsq[0].neigh_top_recon[i]);

        }

    }

    if (obj->md_blk_arr_nsq) {

        EB_FREE_ARRAY(obj->md_blk_arr_nsq[0].av1xd);

    }

    EB_FREE_ARRAY(obj->mds);

    EB_FREE_ARRAY(obj->pc_tree);

    EB_FREE_ARRAY(obj->tested_blk);

    obj->blocks_to_alloc = 0;

    EB_FREE_ARRAY(obj->md_blk_arr_nsq);

    if (obj->rate_est_table) {

        EB_FREE_ARRAY(obj->rate_est_table);

    }

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

        if (obj->cmp_store.pred0_buf[i]) {

            EB_FREE(obj->cmp_store.pred0_buf[i]);

        }

        if (obj->cmp_store.pred1_buf[i]) {

            EB_FREE(obj->cmp_store.pred1_buf[i]);

        }

    }

    if (obj->residual1) {

        EB_FREE(obj->residual1);

    }

    if (obj->diff10) {

        EB_FREE(obj->diff10);

    }

    if (obj->intrapred_buf) {

        EB_FREE_2D(obj->intrapred_buf);

    }

    if (obj->obmc_buff_0) {

        EB_FREE(obj->obmc_buff_0);

    }

    if (obj->obmc_buff_1) {

        EB_FREE(obj->obmc_buff_1);

    }

    if (obj->wsrc_buf) {

        EB_FREE(obj->wsrc_buf);

    }

    if (obj->mask_buf) {

        EB_FREE(obj->mask_buf);

    }

    for (uint32_t txt_itr = 0; txt_itr < TX_TYPES; ++txt_itr) {

        EB_DELETE(obj->recon_coeff_ptr[txt_itr]);

        EB_DELETE(obj->recon_ptr[txt_itr]);

        EB_DELETE(obj->quant_coeff_ptr[txt_itr]);

    }

    EB_DELETE(obj->tx_coeffs);

    EB_DELETE(obj->scratch_prediction_ptr);

    EB_DELETE(obj->temp_residual);

    EB_DELETE(obj->temp_recon_ptr);

    EB_FREE_ARRAY(obj->full_cost_ssim_array);

}

void svt_aom_set_nics(SequenceControlSet* scs, NicScalingCtrls* scaling_ctrls, uint32_t mds1_count[CAND_CLASS_TOTAL],

                      uint32_t mds2_count[CAND_CLASS_TOTAL], uint32_t mds3_count[CAND_CLASS_TOTAL], uint8_t pic_type,

                      uint32_t qp);

static void setup_mds(SequenceControlSet* scs, MdScan* mds, uint32_t* mds_idx, int index, BlockSize bsize,

                      const int min_sq_size) {

    mds->mds_idx = *mds_idx;

    mds->bsize   = bsize;

    mds->index   = index;

    // If applicable, add split depths

    const BlockGeom* blk_geom = get_blk_geom_mds(scs->blk_geom_mds, *mds_idx);

    const int        sq_size  = block_size_wide[bsize];

    if (sq_size > min_sq_size) {

        const BlockSize subsize             = get_partition_subsize(bsize, PARTITION_SPLIT);

        const int       sq_subsize          = block_size_wide[subsize];

        int             blocks_per_subdepth = (sq_subsize / min_sq_size) * (sq_subsize / min_sq_size);

        int             blocks_to_skip      = 0;

        for (int i = min_sq_size; i <= sq_subsize; i <<= 1, blocks_per_subdepth >>= 2) {

            blocks_to_skip += blocks_per_subdepth;

        }

        *mds_idx += blk_geom->d1_depth_offset;

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

            mds->split[i] = mds + i * blocks_to_skip + 1;

            setup_mds(scs, mds->split[i], mds_idx, i, subsize, min_sq_size);

        }

    } else {

        *mds_idx += blk_geom->ns_depth_offset;

    }

}

static void setup_pc_tree(PC_TREE* pc_tree, bool (*test_blk_array)[PART_S][4], int index, BlockSize bsize,

                          const int min_sq_size) {

    pc_tree->bsize      = bsize;

    pc_tree->index      = index;

    pc_tree->tested_blk = test_blk_array[0];

    // If applicable, add split depths

    const int sq_size = block_size_wide[bsize];

    if (sq_size > min_sq_size) {

        const BlockSize subsize             = get_partition_subsize(bsize, PARTITION_SPLIT);

        const int       sq_subsize          = block_size_wide[subsize];

        int             blocks_per_subdepth = (sq_subsize / min_sq_size) * (sq_subsize / min_sq_size);

        int             blocks_to_skip      = 0;

        for (int i = min_sq_size; i <= sq_subsize; i <<= 1, blocks_per_subdepth >>= 2) {

            blocks_to_skip += blocks_per_subdepth;

        }

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

            pc_tree->split[i]         = pc_tree + i * blocks_to_skip + 1;

            pc_tree->split[i]->parent = pc_tree;

            setup_pc_tree(pc_tree->split[i], test_blk_array + i * blocks_to_skip + 1, i, subsize, min_sq_size);

        }

    }

}

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

 * Mode Decision Context Constructor

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

EbErrorType svt_aom_mode_decision_context_ctor(ModeDecisionContext* ctx, SequenceControlSet* scs,

                                               EbColorFormat color_format, uint8_t sb_size, EncMode enc_mode,

                                               uint16_t max_block_cnt, uint32_t encoder_bit_depth,

                                               EbFifo* mode_decision_configuration_input_fifo_ptr,

                                               EbFifo* mode_decision_output_fifo_ptr, uint8_t enable_hbd_mode_decision,

                                               uint8_t seq_qp_mod) {

    const bool allintra = scs->allintra;

    const bool rtc_tune = scs->static_config.rtc;

    uint32_t   buffer_index;

    uint32_t   cand_index;

    ctx->init_max_block_cnt     = max_block_cnt;

    uint32_t block_max_count_sb = max_block_cnt;

    ctx->sb_size = sb_size;

    (void)color_format;

    ctx->dctor  = mode_decision_context_dctor;

    ctx->hbd_md = enable_hbd_mode_decision;

    // Input/Output System Resource Manager FIFOs

    ctx->mode_decision_configuration_input_fifo_ptr = mode_decision_configuration_input_fifo_ptr;

    ctx->mode_decision_output_fifo_ptr              = mode_decision_output_fifo_ptr;

    // Maximum number of candidates MD can support

    // determine MAX_NICS for a given preset

    // get the min scaling level (the smallest scaling level is the most conservative)

    uint8_t min_nic_scaling_level = NICS_SCALING_LEVELS - 1;

    uint8_t stage1_scaling_num;

    if (allintra) {

        uint8_t nic_level  = svt_aom_get_nic_level_allintra(enc_mode);

        stage1_scaling_num = MD_STAGE_NICS_SCAL_NUM[svt_aom_set_nic_controls(NULL, nic_level)][MD_STAGE_1];

    } else if (rtc_tune) {

#if TUNE_SIMPLIFY_SETTINGS

        uint8_t nic_level = svt_aom_get_nic_level_rtc(enc_mode);

#else

        uint8_t nic_level = svt_aom_get_nic_level_rtc(enc_mode, scs->use_flat_ipp);

#endif

        stage1_scaling_num = MD_STAGE_NICS_SCAL_NUM[svt_aom_set_nic_controls(NULL, nic_level)][MD_STAGE_1];

    } else {

#if TUNE_SIMPLIFY_SETTINGS

        for (uint8_t is_base = 0; is_base < 2; is_base++) {

            uint8_t nic_level         = svt_aom_get_nic_level_default(enc_mode, is_base);

            uint8_t nic_scaling_level = svt_aom_set_nic_controls(NULL, nic_level);

            min_nic_scaling_level     = MIN(min_nic_scaling_level, nic_scaling_level);

        }

#else

        for (uint8_t sc_class1 = 0; sc_class1 < 2; sc_class1++) {

            for (uint8_t is_base = 0; is_base < 2; is_base++) {

                uint8_t nic_level         = svt_aom_get_nic_level_default(enc_mode, is_base, sc_class1);

                uint8_t nic_scaling_level = svt_aom_set_nic_controls(NULL, nic_level);

                min_nic_scaling_level     = MIN(min_nic_scaling_level, nic_scaling_level);

            }

        }

#endif

        stage1_scaling_num = MD_STAGE_NICS_SCAL_NUM[min_nic_scaling_level][MD_STAGE_1];

    }

    // scale max_nics

    uint32_t max_nics = 0;

    {

        NicScalingCtrls scaling_ctrls;

        scaling_ctrls.stage1_scaling_num = stage1_scaling_num;

        scaling_ctrls.stage2_scaling_num = stage1_scaling_num;

        scaling_ctrls.stage3_scaling_num = stage1_scaling_num;

        uint32_t mds1_count[CAND_CLASS_TOTAL];

        uint32_t mds2_count[CAND_CLASS_TOTAL];

        uint32_t mds3_count[CAND_CLASS_TOTAL];

        for (uint8_t pic_type = 0; pic_type < NICS_PIC_TYPE; pic_type++) {

            for (uint8_t qp = MIN_QP_VALUE; qp <= MAX_QP_VALUE; qp++) {

                svt_aom_set_nics(scs, &scaling_ctrls, mds1_count, mds2_count, mds3_count, pic_type, qp);

                uint32_t nics = 0;

                for (CandClass cidx = CAND_CLASS_0; cidx < CAND_CLASS_TOTAL; cidx++) {

                    nics += mds1_count[cidx];

                }

                max_nics = MAX(max_nics, nics);

            }

        }

    }

    // If independent chroma search is used, need to allocate additional 84 candidate buffers

    bool is_chroma_mode_0;

    if (allintra) {

        is_chroma_mode_0 = svt_aom_set_chroma_controls(NULL, svt_aom_get_chroma_level_allintra(enc_mode)) ==

            CHROMA_MODE_0;

    } else if (scs->static_config.rtc) {

#if TUNE_SIMPLIFY_SETTINGS

        is_chroma_mode_0 = svt_aom_set_chroma_controls(

                               NULL, svt_aom_get_chroma_level_rtc(enc_mode, scs->use_flat_ipp)) == CHROMA_MODE_0;

#else

        for (uint8_t is_i_slice = 0; is_i_slice < 2; is_i_slice++) {

            is_chroma_mode_0 = svt_aom_set_chroma_controls(NULL, svt_aom_get_chroma_level_rtc(enc_mode, is_i_slice)) ==

                CHROMA_MODE_0;

            if (is_chroma_mode_0) {

                break;

            }

        }

#endif

    } else {

        for (uint8_t is_i_slice = 0; is_i_slice < 2; is_i_slice++) {

            is_chroma_mode_0 = svt_aom_set_chroma_controls(

                                   NULL, svt_aom_get_chroma_level_default(enc_mode, is_i_slice)) == CHROMA_MODE_0;

            if (is_chroma_mode_0) {

                break;

            }

        }

    }

    const uint8_t ind_uv_cands = is_chroma_mode_0 ? 84 : 0;

    max_nics += CAND_CLASS_TOTAL; //need one extra temp buffer for each fast loop call

    ctx->max_nics    = max_nics;

    ctx->max_nics_uv = max_nics + ind_uv_cands;

    // Cfl scratch memory

    if (ctx->hbd_md > EB_8_BIT_MD) {

        EB_MALLOC_ALIGNED(ctx->cfl_temp_luma_recon16bit, sizeof(uint16_t) * sb_size * sb_size);

    }

    if (ctx->hbd_md != EB_10_BIT_MD) {

        EB_MALLOC_ALIGNED(ctx->cfl_temp_luma_recon, sizeof(uint8_t) * sb_size * sb_size);

    }

    EB_MALLOC_ALIGNED(ctx->pred_buf_q3, CFL_BUF_SQUARE);

    uint8_t use_update_cdf = 0;

    if (allintra) {

        use_update_cdf = svt_aom_get_update_cdf_level_allintra(enc_mode);

    } else if (rtc_tune) {

#if TUNE_SIMPLIFY_SETTINGS

        for (uint8_t is_islice = 0; is_islice < 2; is_islice++) {

            if (use_update_cdf) {

                break;

            }

            use_update_cdf |= svt_aom_get_update_cdf_level_rtc(enc_mode, is_islice);

        }

#else

        for (uint8_t sc_class1 = 0; sc_class1 < 2; sc_class1++) {

            for (uint8_t is_islice = 0; is_islice < 2; is_islice++) {

                for (uint8_t is_base = 0; is_base < 2; is_base++) {

                    if (use_update_cdf) {

                        break;

                    }

                    use_update_cdf |= svt_aom_get_update_cdf_level_rtc(enc_mode, is_islice, is_base, sc_class1);

                }

            }

        }

#endif

    } else {

#if TUNE_SIMPLIFY_SETTINGS

        for (uint8_t is_islice = 0; is_islice < 2; is_islice++) {

            for (uint8_t is_base = 0; is_base < 2; is_base++) {

                if (use_update_cdf) {

                    break;

                }

                use_update_cdf |= svt_aom_get_update_cdf_level_default(enc_mode, is_islice, is_base);

            }

        }

#else

        for (uint8_t sc_class1 = 0; sc_class1 < 2; sc_class1++) {

            for (uint8_t is_islice = 0; is_islice < 2; is_islice++) {

                for (uint8_t is_base = 0; is_base < 2; is_base++) {

                    if (use_update_cdf) {

                        break;

                    }

                    use_update_cdf |= svt_aom_get_update_cdf_level_default(enc_mode, is_islice, is_base, sc_class1);

                }

            }

        }

#endif

    }

    if (use_update_cdf) {

        EB_CALLOC_ARRAY(ctx->rate_est_table, 1);

    } else {

        ctx->rate_est_table = NULL;

    }

    // Allocate buffer for inter-inter compound prediction

    if (get_inter_compound_level(enc_mode)) {

        const uint8_t bits = ctx->hbd_md > EB_8_BIT_MD ? 2 : 1;

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

            EB_MALLOC(ctx->cmp_store.pred0_buf[i], sb_size * sb_size * bits * sizeof(uint8_t));

            EB_MALLOC(ctx->cmp_store.pred1_buf[i], sb_size * sb_size * bits * sizeof(uint8_t));

        }

        EB_MALLOC(ctx->residual1, sb_size * sb_size * sizeof(ctx->residual1[0]));

        EB_MALLOC(ctx->diff10, sb_size * sb_size * sizeof(ctx->diff10[0]));

    }

    // Allocate buffer for inter-intra prediction

    uint8_t ii_allowed = 0;

    for (uint8_t transition_present = 0; transition_present < 2; transition_present++) {

        if (ii_allowed) {

            break;

        }

        ii_allowed |= svt_aom_get_inter_intra_level(enc_mode, transition_present);

    }

    if (ii_allowed) {

        const uint8_t bits = ctx->hbd_md > EB_8_BIT_MD ? 2 : 1;

        // MAX block size for inter intra is 32x32

        EB_MALLOC_2D(ctx->intrapred_buf, INTERINTRA_MODES, 32 * 32 * bits * sizeof(ctx->intrapred_buf[0][0]));

    }

    // Allocate buffers for obmc prediction

    uint8_t obmc_allowed = 0;

    for (uint8_t is_base = 0; is_base < 2; is_base++) {

        for (uint8_t qp = MIN_QP_VALUE; qp <= MAX_QP_VALUE; qp++) {

            if (obmc_allowed) {

                break;

            }

#if TUNE_SHIFT_PRESETS_RTC

            obmc_allowed |= svt_aom_get_obmc_level(enc_mode, qp, seq_qp_mod, rtc_tune);

#else

            obmc_allowed |= svt_aom_get_obmc_level(enc_mode, qp, seq_qp_mod);

#endif

        }

    }

    if (obmc_allowed) {

        const uint8_t bits = ctx->hbd_md > EB_8_BIT_MD ? 2 : 1;

        EB_MALLOC(ctx->obmc_buff_0, sb_size * sb_size * bits * MAX_PLANES * sizeof(ctx->obmc_buff_0[0]));

        EB_MALLOC(ctx->obmc_buff_1, sb_size * sb_size * bits * MAX_PLANES * sizeof(ctx->obmc_buff_1[0]));

        EB_MALLOC(ctx->wsrc_buf, sb_size * sb_size * sizeof(ctx->wsrc_buf[0]));

        EB_MALLOC(ctx->mask_buf, sb_size * sb_size * sizeof(ctx->mask_buf[0]));

    }

    EB_MALLOC_ARRAY(ctx->md_blk_arr_nsq, block_max_count_sb);

    // Fast Candidate Array

    uint16_t max_can_count = svt_aom_get_max_can_count(enc_mode) + ind_uv_cands;

    EB_MALLOC_ARRAY(ctx->fast_cand_array, max_can_count);

    for (cand_index = 0; cand_index < max_can_count; ++cand_index) {

        ctx->fast_cand_array[cand_index].palette_info = NULL;

    }

    svt_aom_assert_err(max_can_count > ind_uv_cands, "Max. candidates is too low");

    EB_MALLOC_2D(ctx->injected_mvs, (uint16_t)(max_can_count - ind_uv_cands), 2);

    EB_MALLOC_ARRAY(ctx->injected_ref_types, (max_can_count - ind_uv_cands));

    // Set buffers for MD palette search to NULL; will be init'd at runtime if needed

    ctx->palette_buffer       = NULL;

    ctx->palette_cand_array   = NULL;

    ctx->palette_size_array_0 = NULL;

    // Cost Arrays

    EB_MALLOC_ARRAY(ctx->fast_cost_array, ctx->max_nics_uv);

    EB_MALLOC_ARRAY(ctx->full_cost_array, ctx->max_nics_uv);

    EB_MALLOC_ARRAY(ctx->full_cost_ssim_array, ctx->max_nics_uv);

    // Candidate Buffers

    EB_NEW(ctx->cand_bf_tx_depth_1,

           svt_aom_mode_decision_scratch_cand_bf_ctor,

           sb_size,

           ctx->hbd_md ? EB_TEN_BIT : EB_EIGHT_BIT);

    EB_ALLOC_PTR_ARRAY(ctx->cand_bf_tx_depth_1->cand, 1);

    EB_NEW(ctx->cand_bf_tx_depth_2,

           svt_aom_mode_decision_scratch_cand_bf_ctor,

           sb_size,

           ctx->hbd_md ? EB_TEN_BIT : EB_EIGHT_BIT);

    EB_ALLOC_PTR_ARRAY(ctx->cand_bf_tx_depth_2->cand, 1);

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

        ctx->md_blk_arr_nsq[0].neigh_left_recon[i]       = NULL;

        ctx->md_blk_arr_nsq[0].neigh_top_recon[i]        = NULL;

        ctx->md_blk_arr_nsq[0].neigh_left_recon_16bit[i] = NULL;

        ctx->md_blk_arr_nsq[0].neigh_top_recon_16bit[i]  = NULL;

    }

    uint32_t coded_leaf_index;

    uint16_t sz = sizeof(uint16_t);

    if (ctx->hbd_md > EB_8_BIT_MD) {

        EB_MALLOC_ARRAY(ctx->md_blk_arr_nsq[0].neigh_left_recon_16bit[0], block_max_count_sb * sb_size * sz);

        EB_MALLOC_ARRAY(ctx->md_blk_arr_nsq[0].neigh_top_recon_16bit[0], block_max_count_sb * sb_size * sz);

        EB_MALLOC_ARRAY(ctx->md_blk_arr_nsq[0].neigh_left_recon_16bit[1], block_max_count_sb * sb_size * sz >> 1);

        EB_MALLOC_ARRAY(ctx->md_blk_arr_nsq[0].neigh_top_recon_16bit[1], block_max_count_sb * sb_size * sz >> 1);

        EB_MALLOC_ARRAY(ctx->md_blk_arr_nsq[0].neigh_left_recon_16bit[2], block_max_count_sb * sb_size * sz >> 1);

        EB_MALLOC_ARRAY(ctx->md_blk_arr_nsq[0].neigh_top_recon_16bit[2], block_max_count_sb * sb_size * sz >> 1);

        for (coded_leaf_index = 0; coded_leaf_index < block_max_count_sb; ++coded_leaf_index) {

            size_t offset = coded_leaf_index * sb_size * sz;

            ctx->md_blk_arr_nsq[coded_leaf_index].neigh_left_recon_16bit[0] =

                ctx->md_blk_arr_nsq[0].neigh_left_recon_16bit[0] + offset;

            ctx->md_blk_arr_nsq[coded_leaf_index].neigh_top_recon_16bit[0] =

                ctx->md_blk_arr_nsq[0].neigh_top_recon_16bit[0] + offset;

            offset >>= 1;

            ctx->md_blk_arr_nsq[coded_leaf_index].neigh_left_recon_16bit[1] =

                ctx->md_blk_arr_nsq[0].neigh_left_recon_16bit[1] + offset;

            ctx->md_blk_arr_nsq[coded_leaf_index].neigh_top_recon_16bit[1] =

                ctx->md_blk_arr_nsq[0].neigh_top_recon_16bit[1] + offset;

            ctx->md_blk_arr_nsq[coded_leaf_index].neigh_left_recon_16bit[2] =

                ctx->md_blk_arr_nsq[0].neigh_left_recon_16bit[2] + offset;

            ctx->md_blk_arr_nsq[coded_leaf_index].neigh_top_recon_16bit[2] =

                ctx->md_blk_arr_nsq[0].neigh_top_recon_16bit[2] + offset;

        }

    }

    if (ctx->hbd_md != EB_10_BIT_MD) {

        EB_MALLOC_ARRAY(ctx->md_blk_arr_nsq[0].neigh_left_recon[0], block_max_count_sb * sb_size);

        EB_MALLOC_ARRAY(ctx->md_blk_arr_nsq[0].neigh_top_recon[0], block_max_count_sb * sb_size);

        EB_MALLOC_ARRAY(ctx->md_blk_arr_nsq[0].neigh_left_recon[1], block_max_count_sb * sb_size >> 1);

        EB_MALLOC_ARRAY(ctx->md_blk_arr_nsq[0].neigh_top_recon[1], block_max_count_sb * sb_size >> 1);

        EB_MALLOC_ARRAY(ctx->md_blk_arr_nsq[0].neigh_left_recon[2], block_max_count_sb * sb_size >> 1);

        EB_MALLOC_ARRAY(ctx->md_blk_arr_nsq[0].neigh_top_recon[2], block_max_count_sb * sb_size >> 1);

        for (coded_leaf_index = 0; coded_leaf_index < block_max_count_sb; ++coded_leaf_index) {

            size_t offset                                             = coded_leaf_index * sb_size;

            ctx->md_blk_arr_nsq[coded_leaf_index].neigh_left_recon[0] = ctx->md_blk_arr_nsq[0].neigh_left_recon[0] +

                offset;

            ctx->md_blk_arr_nsq[coded_leaf_index].neigh_top_recon[0] = ctx->md_blk_arr_nsq[0].neigh_top_recon[0] +

                offset;

            offset >>= 1;

            ctx->md_blk_arr_nsq[coded_leaf_index].neigh_left_recon[1] = ctx->md_blk_arr_nsq[0].neigh_left_recon[1] +

                offset;

            ctx->md_blk_arr_nsq[coded_leaf_index].neigh_top_recon[1] = ctx->md_blk_arr_nsq[0].neigh_top_recon[1] +

                offset;

            ctx->md_blk_arr_nsq[coded_leaf_index].neigh_left_recon[2] = ctx->md_blk_arr_nsq[0].neigh_left_recon[2] +

                offset;

            ctx->md_blk_arr_nsq[coded_leaf_index].neigh_top_recon[2] = ctx->md_blk_arr_nsq[0].neigh_top_recon[2] +

                offset;

        }

    }

    ctx->md_blk_arr_nsq[0].av1xd = NULL;

    EB_MALLOC_ARRAY(ctx->md_blk_arr_nsq[0].av1xd, block_max_count_sb);

    // Alloc mds and pc_tree, which are used to track tested blocks in MD

    bool disallow_4x4 = allintra ? svt_aom_get_disallow_4x4_allintra(enc_mode)

#if TUNE_SIMPLIFY_SETTINGS

        : rtc_tune ? svt_aom_get_disallow_4x4_rtc()

#else

        : rtc_tune ? svt_aom_get_disallow_4x4_rtc(enc_mode)

#endif

                   : svt_aom_get_disallow_4x4_default(enc_mode);

    bool    disallow_8x8     = allintra ? svt_aom_get_disallow_8x8_allintra()

               : rtc_tune ? svt_aom_get_disallow_8x8_rtc(enc_mode, scs->max_input_luma_width, scs->max_input_luma_height)

                          : svt_aom_get_disallow_8x8_default();

    uint8_t min_bsize        = disallow_8x8 ? 16 : disallow_4x4 ? 8 : 4;

    int     blocks_per_depth = (sb_size / min_bsize) * (sb_size / min_bsize);

    int     blocks_to_alloc  = 0;

    for (int i = min_bsize; i <= sb_size; i <<= 1, blocks_per_depth >>= 2) {

        blocks_to_alloc += blocks_per_depth;

    }

    EB_CALLOC_ARRAY(ctx->mds, blocks_to_alloc);

    uint32_t mds_idx = 0;

    setup_mds(scs, ctx->mds, &mds_idx, 0, scs->seq_header.sb_size, min_bsize);

    EB_CALLOC_ARRAY(ctx->pc_tree, blocks_to_alloc);

    EB_MALLOC_ARRAY(ctx->tested_blk, blocks_to_alloc);

    setup_pc_tree(ctx->pc_tree, ctx->tested_blk, 0, scs->seq_header.sb_size, min_bsize);

    ctx->blocks_to_alloc = blocks_to_alloc;

    bool bypass_encdec = allintra ? svt_aom_get_bypass_encdec_allintra(enc_mode)

        : rtc_tune                ? svt_aom_get_bypass_encdec_rtc(enc_mode, encoder_bit_depth)

                                  : svt_aom_get_bypass_encdec_default(enc_mode, encoder_bit_depth);

    for (coded_leaf_index = 0; coded_leaf_index < block_max_count_sb; ++coded_leaf_index) {

        ctx->md_blk_arr_nsq[coded_leaf_index].av1xd      = ctx->md_blk_arr_nsq[0].av1xd + coded_leaf_index;

        ctx->md_blk_arr_nsq[coded_leaf_index].segment_id = 0;

        const BlockGeom* blk_geom                        = get_blk_geom_mds(scs->blk_geom_mds, coded_leaf_index);

        if (bypass_encdec) {

            EbPictureBufferDescInitData init_data;

            init_data.buffer_enable_mask = PICTURE_BUFFER_DESC_FULL_MASK;

            init_data.max_width          = blk_geom->bwidth;

            init_data.max_height         = blk_geom->bheight;

            init_data.bit_depth          = EB_THIRTYTWO_BIT;

            init_data.color_format       = (blk_geom->bwidth > 4 && blk_geom->bheight > 4)

                      ? EB_YUV420

                      : EB_YUV444; // PW - must have at least 4x4 for chroma coeffs

            init_data.border             = 0;

            init_data.split_mode         = false;

            EB_NEW(ctx->md_blk_arr_nsq[coded_leaf_index].coeff_tmp, svt_picture_buffer_desc_ctor, (EbPtr)&init_data);

            init_data.buffer_enable_mask = PICTURE_BUFFER_DESC_FULL_MASK;

            init_data.max_width          = blk_geom->bwidth;

            init_data.max_height         = blk_geom->bheight;

            init_data.bit_depth          = ctx->hbd_md ? EB_TEN_BIT : EB_EIGHT_BIT;

            ;

            init_data.color_format = (blk_geom->bwidth > 4 && blk_geom->bheight > 4) ? EB_YUV420 : EB_YUV444;

            init_data.border       = 0;

            init_data.split_mode   = false;

            EB_NEW(ctx->md_blk_arr_nsq[coded_leaf_index].recon_tmp, svt_picture_buffer_desc_ctor, (EbPtr)&init_data);

        } else {

            ctx->md_blk_arr_nsq[coded_leaf_index].coeff_tmp = NULL;

            ctx->md_blk_arr_nsq[coded_leaf_index].recon_tmp = NULL;

        }

    }

    for (CandClass cand_class_it = CAND_CLASS_0; cand_class_it < CAND_CLASS_TOTAL; cand_class_it++) {

        EB_MALLOC_ARRAY(ctx->cand_buff_indices[cand_class_it], ctx->max_nics_uv);

    }

    EB_MALLOC_ARRAY(ctx->best_candidate_index_array, ctx->max_nics_uv);

    EB_MALLOC_ARRAY(ctx->above_txfm_context, (sb_size >> MI_SIZE_LOG2));

    EB_MALLOC_ARRAY(ctx->left_txfm_context, (sb_size >> MI_SIZE_LOG2));

    EbPictureBufferDescInitData thirty_two_width_picture_buffer_desc_init_data;

    EbPictureBufferDescInitData picture_buffer_desc_init_data;

    picture_buffer_desc_init_data.max_width          = sb_size;

    picture_buffer_desc_init_data.max_height         = sb_size;

    picture_buffer_desc_init_data.bit_depth          = ctx->hbd_md ? EB_TEN_BIT : EB_EIGHT_BIT;

    picture_buffer_desc_init_data.color_format       = EB_YUV420;

    picture_buffer_desc_init_data.buffer_enable_mask = PICTURE_BUFFER_DESC_FULL_MASK;

    picture_buffer_desc_init_data.border             = 0;

    picture_buffer_desc_init_data.split_mode         = false;

    picture_buffer_desc_init_data.is_16bit_pipeline  = false;

    thirty_two_width_picture_buffer_desc_init_data.max_width          = sb_size;

    thirty_two_width_picture_buffer_desc_init_data.max_height         = sb_size;

    thirty_two_width_picture_buffer_desc_init_data.bit_depth          = EB_THIRTYTWO_BIT;

    thirty_two_width_picture_buffer_desc_init_data.color_format       = EB_YUV420;

    thirty_two_width_picture_buffer_desc_init_data.buffer_enable_mask = PICTURE_BUFFER_DESC_FULL_MASK;

    thirty_two_width_picture_buffer_desc_init_data.border             = 0;

    thirty_two_width_picture_buffer_desc_init_data.split_mode         = false;

    thirty_two_width_picture_buffer_desc_init_data.is_16bit_pipeline  = false;

    for (uint32_t txt_itr = 0; txt_itr < TX_TYPES; ++txt_itr) {

        EB_NEW(ctx->recon_coeff_ptr[txt_itr],

               svt_picture_buffer_desc_ctor,

               (EbPtr)&thirty_two_width_picture_buffer_desc_init_data);

        EB_NEW(ctx->recon_ptr[txt_itr], svt_picture_buffer_desc_ctor, (EbPtr)&picture_buffer_desc_init_data);

        EB_NEW(ctx->quant_coeff_ptr[txt_itr],

               svt_picture_buffer_desc_ctor,

               (EbPtr)&thirty_two_width_picture_buffer_desc_init_data);

    }

    EB_NEW(ctx->tx_coeffs, svt_picture_buffer_desc_ctor, (EbPtr)&thirty_two_width_picture_buffer_desc_init_data);

    EB_NEW(ctx->scratch_prediction_ptr, svt_picture_buffer_desc_ctor, (EbPtr)&picture_buffer_desc_init_data);

    EbPictureBufferDescInitData double_width_picture_buffer_desc_init_data;

    double_width_picture_buffer_desc_init_data.max_width          = sb_size;

    double_width_picture_buffer_desc_init_data.max_height         = sb_size;

    double_width_picture_buffer_desc_init_data.bit_depth          = EB_SIXTEEN_BIT;

    double_width_picture_buffer_desc_init_data.color_format       = EB_YUV420;

    double_width_picture_buffer_desc_init_data.buffer_enable_mask = PICTURE_BUFFER_DESC_FULL_MASK;

    double_width_picture_buffer_desc_init_data.border             = 0;

    double_width_picture_buffer_desc_init_data.split_mode         = false;

    double_width_picture_buffer_desc_init_data.is_16bit_pipeline  = false;

    // The temp_recon_ptr and temp_residual will be shared by all candidates

    // If you want to do something with residual or recon, you need to create one

    EB_NEW(ctx->temp_recon_ptr, svt_picture_buffer_desc_ctor, (EbPtr)&picture_buffer_desc_init_data);

    EB_NEW(ctx->temp_residual, svt_picture_buffer_desc_ctor, (EbPtr)&double_width_picture_buffer_desc_init_data);

    // Candidate Buffers

    EB_ALLOC_PTR_ARRAY(ctx->cand_bf_ptr_array, ctx->max_nics_uv);

    for (buffer_index = 0; buffer_index < ctx->max_nics; ++buffer_index) {

        EB_NEW(ctx->cand_bf_ptr_array[buffer_index],

               svt_aom_mode_decision_cand_bf_ctor,

               ctx->hbd_md ? EB_TEN_BIT : EB_EIGHT_BIT,

               sb_size,

               PICTURE_BUFFER_DESC_FULL_MASK,

               ctx->temp_residual,

               ctx->temp_recon_ptr,

               &(ctx->fast_cost_array[buffer_index]),

               &(ctx->full_cost_array[buffer_index]),

               &(ctx->full_cost_ssim_array[buffer_index]));

    }

    for (buffer_index = max_nics; buffer_index < ctx->max_nics_uv; ++buffer_index) {

        EB_NEW(ctx->cand_bf_ptr_array[buffer_index],

               svt_aom_mode_decision_cand_bf_ctor,

               ctx->hbd_md ? EB_TEN_BIT : EB_EIGHT_BIT,

               sb_size,

               PICTURE_BUFFER_DESC_CHROMA_MASK,

               ctx->temp_residual,

               ctx->temp_recon_ptr,

               &(ctx->fast_cost_array[buffer_index]),

               &(ctx->full_cost_array[buffer_index]),

               &(ctx->full_cost_ssim_array[buffer_index]));

    }

    return EB_ErrorNone;

}

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

 * Reset Mode Decision Neighbor Arrays

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

void svt_aom_reset_mode_decision_neighbor_arrays(PictureControlSet* pcs, uint16_t tile_idx) {

    uint8_t depth;

    for (depth = 0; depth < NA_TOT_CNT; depth++) {

        svt_aom_neighbor_array_unit_reset(pcs->mdleaf_partition_na[depth][tile_idx]);

        if (pcs->hbd_md != EB_10_BIT_MD) {

            svt_aom_neighbor_array_unit_reset(pcs->md_luma_recon_na[depth][tile_idx]);

            svt_aom_neighbor_array_unit_reset(pcs->md_tx_depth_1_luma_recon_na[depth][tile_idx]);

            svt_aom_neighbor_array_unit_reset(pcs->md_tx_depth_2_luma_recon_na[depth][tile_idx]);

            svt_aom_neighbor_array_unit_reset(pcs->md_cb_recon_na[depth][tile_idx]);

            svt_aom_neighbor_array_unit_reset(pcs->md_cr_recon_na[depth][tile_idx]);

        }

        if (pcs->hbd_md > EB_8_BIT_MD || (pcs->scs->encoder_bit_depth > EB_EIGHT_BIT && pcs->pic_bypass_encdec)) {

            svt_aom_neighbor_array_unit_reset(pcs->md_luma_recon_na_16bit[depth][tile_idx]);

            svt_aom_neighbor_array_unit_reset(pcs->md_tx_depth_1_luma_recon_na_16bit[depth][tile_idx]);

            svt_aom_neighbor_array_unit_reset(pcs->md_tx_depth_2_luma_recon_na_16bit[depth][tile_idx]);

            svt_aom_neighbor_array_unit_reset(pcs->md_cb_recon_na_16bit[depth][tile_idx]);

            svt_aom_neighbor_array_unit_reset(pcs->md_cr_recon_na_16bit[depth][tile_idx]);

        }

        svt_aom_neighbor_array_unit_reset(pcs->md_y_dcs_na[depth][tile_idx]);

        svt_aom_neighbor_array_unit_reset(pcs->md_tx_depth_1_luma_dc_sign_level_coeff_na[depth][tile_idx]);

        svt_aom_neighbor_array_unit_reset(pcs->md_cb_dc_sign_level_coeff_na[depth][tile_idx]);

        svt_aom_neighbor_array_unit_reset(pcs->md_cr_dc_sign_level_coeff_na[depth][tile_idx]);

        svt_aom_neighbor_array_unit_reset(pcs->md_txfm_context_array[depth][tile_idx]);

    }

    return;

}

// If the ref intra percentage is below the TH, applying modulation to the MD lambda

#define LAMBDA_MOD_INTRA_TH 50

#define LAMBDA_MOD_INTRA_SCALING_FACTOR 138

// Set the lambda for each sb.

// When lambda tuning is on (blk_lambda_tuning), lambda of each block is set separately (full_lambda_md/fast_lambda_md)

// later in svt_aom_set_tuned_blk_lambda

// Testing showed that updating SAD lambda based on frame info was not helpful; therefore, the SAD lambda generation is not changed.

static void av1_lambda_assign_md(PictureControlSet* pcs, ModeDecisionContext* ctx) {

    ctx->full_lambda_md[0] = svt_aom_compute_rd_mult(pcs, ctx->qp_index, ctx->me_q_index, EB_EIGHT_BIT);

    ctx->fast_lambda_md[0] = svt_aom_compute_fast_lambda(pcs, ctx->qp_index, ctx->me_q_index, EB_EIGHT_BIT);

    ctx->full_lambda_md[1] = svt_aom_compute_rd_mult(pcs, ctx->qp_index, ctx->me_q_index, EB_TEN_BIT);

    ctx->fast_lambda_md[1] = svt_aom_compute_fast_lambda(pcs, ctx->qp_index, ctx->me_q_index, EB_TEN_BIT);

    if (!pcs->scs->static_config.rtc && pcs->scs->stats_based_sb_lambda_modulation) {

        if (pcs->temporal_layer_index > 0) {

            if (pcs->ref_intra_percentage < LAMBDA_MOD_INTRA_TH) {

                ctx->full_lambda_md[0] = (ctx->full_lambda_md[0] * LAMBDA_MOD_INTRA_SCALING_FACTOR) >> 7;

                ctx->fast_lambda_md[0] = (ctx->fast_lambda_md[0] * LAMBDA_MOD_INTRA_SCALING_FACTOR) >> 7;

                ctx->full_lambda_md[1] = (ctx->full_lambda_md[1] * LAMBDA_MOD_INTRA_SCALING_FACTOR) >> 7;

                ctx->fast_lambda_md[1] = (ctx->fast_lambda_md[1] * LAMBDA_MOD_INTRA_SCALING_FACTOR) >> 7;

            }

        }

    }

    if (pcs->lambda_weight) {

        ctx->full_lambda_md[0] = (uint32_t)((ctx->full_lambda_md[0] * (uint64_t)pcs->lambda_weight) >> 7);

        ctx->fast_lambda_md[0] = (uint32_t)((ctx->fast_lambda_md[0] * (uint64_t)pcs->lambda_weight) >> 7);

        ctx->full_lambda_md[1] = (uint32_t)((ctx->full_lambda_md[1] * (uint64_t)pcs->lambda_weight) >> 7);

        ctx->fast_lambda_md[1] = (uint32_t)((ctx->fast_lambda_md[1] * (uint64_t)pcs->lambda_weight) >> 7);

    }

    ctx->full_lambda_md[1] *= 16;

    ctx->fast_lambda_md[1] *= 4;

    SequenceControlSet* scs          = pcs->scs;

    uint64_t            scale_factor = scs->static_config.lambda_scale_factors[pcs->ppcs->update_type];

    ctx->full_lambda_md[0]           = (uint32_t)((ctx->full_lambda_md[0] * scale_factor) >> 7);

    ctx->full_lambda_md[1]           = (uint32_t)((ctx->full_lambda_md[1] * scale_factor) >> 7);

    ctx->fast_lambda_md[0]           = (uint32_t)((ctx->fast_lambda_md[0] * scale_factor) >> 7);

    ctx->fast_lambda_md[1]           = (uint32_t)((ctx->fast_lambda_md[1] * scale_factor) >> 7);

    ctx->full_sb_lambda_md[0] = ctx->full_lambda_md[0];

    ctx->full_sb_lambda_md[1] = ctx->full_lambda_md[1];

}

void svt_aom_reset_mode_decision(SequenceControlSet* scs, ModeDecisionContext* ctx, PictureControlSet* pcs,

                                 uint16_t tile_group_idx, uint32_t segment_index) {

    const bool rtc_tune = scs->static_config.rtc;

    ctx->hbd_md         = pcs->hbd_md;

    // Reset MD rate Estimation table to initial values by copying from md_rate_est_ctx

    ctx->md_rate_est_ctx = pcs->md_rate_est_ctx;

    // Reset CABAC Contexts

    // Reset Neighbor Arrays at start of new Segment / Picture

    if (segment_index == 0) {

        for (uint16_t r = pcs->ppcs->tile_group_info[tile_group_idx].tile_group_tile_start_y;

             r < pcs->ppcs->tile_group_info[tile_group_idx].tile_group_tile_end_y;

             r++) {

            for (uint16_t c = pcs->ppcs->tile_group_info[tile_group_idx].tile_group_tile_start_x;

                 c < pcs->ppcs->tile_group_info[tile_group_idx].tile_group_tile_end_x;

                 c++) {

                uint16_t tile_idx = c + r * pcs->ppcs->av1_cm->tiles_info.tile_cols;

                svt_aom_reset_mode_decision_neighbor_arrays(pcs, tile_idx);

            }

        }

        (void)scs;

    }

    //each segment enherits the bypass encdec from the picture level

    ctx->bypass_encdec = pcs->pic_bypass_encdec;

    if (!rtc_tune && (pcs->enc_mode <= ENC_M11 || pcs->temporal_layer_index != 0)) {

        ctx->rtc_use_N4_dct_dct_shortcut = 1;

    } else {

        ctx->rtc_use_N4_dct_dct_shortcut = 0;

    }

    return;

}

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

 * Mode Decision Configure SB

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

void svt_aom_mode_decision_configure_sb(ModeDecisionContext* ctx, PictureControlSet* pcs, uint8_t sb_qp,

                                        uint8_t me_sb_qp) {

    /* Note(CHKN) : when Qp modulation varies QP on a sub-SB(CU) basis,  Lamda has to change based on Cu->QP , and then this code has to move inside the CU loop in MD */

    // Lambda Assignement

    ctx->qp_index = pcs->ppcs->frm_hdr.delta_q_params.delta_q_present || pcs->ppcs->r0_delta_qp_md

        ? sb_qp

        : (uint8_t)pcs->ppcs->frm_hdr.quantization_params.base_q_idx;

    ctx->me_q_index = me_sb_qp;

    av1_lambda_assign_md(pcs, ctx);

    ctx->hbd_pack_done = 0;

    return;

}