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

#include <stdlib.h>
#include <stdio.h>
#include <string.h>
#include "enc_handle.h"
#include "entropy_coding.h"
#include "ec_process.h"
#include "enc_dec_results.h"
#include "ec_results.h"
#include "rc_tasks.h"
#include "cabac_context_model.h"
#include "svt_log.h"
#include "common_dsp_rtcd.h"

static void rest_context_dctor(EbPtr p) {
    EbThreadContext*      thread_ctx = (EbThreadContext*)p;
    EntropyCodingContext* obj        = (EntropyCodingContext*)thread_ctx->priv;
    EB_FREE_ARRAY(obj);
}

/******************************************************
 * Enc Dec Context Constructor
 ******************************************************/
EbErrorType svt_aom_entropy_coding_context_ctor(EbThreadContext* thread_ctx, const EbEncHandle* enc_handle_ptr,
                                                int index) {
    EntropyCodingContext* context_ptr;
    EB_CALLOC_ARRAY(context_ptr, 1);
    thread_ctx->priv  = context_ptr;
    thread_ctx->dctor = rest_context_dctor;

    context_ptr->is_16bit = enc_handle_ptr->scs_instance->scs->static_config.encoder_bit_depth > EB_EIGHT_BIT;

    // Zero levels_buf once; the tail (offset >= LEVELS_TAIL_OFFSET) serves as
    // bottom padding for all block sizes via offset-based placement in set_levels.
    memset(context_ptr->levels_buf + LEVELS_TAIL_OFFSET, 0, TX_PAD_2D - LEVELS_TAIL_OFFSET);

    // Input/Output System Resource Manager FIFOs
    context_ptr->enc_dec_input_fifo_ptr = svt_system_resource_get_consumer_fifo(
        enc_handle_ptr->rest_results_resource_ptr, index);
    context_ptr->entropy_coding_output_fifo_ptr = svt_system_resource_get_producer_fifo(
        enc_handle_ptr->entropy_coding_results_resource_ptr, index);

    return EB_ErrorNone;
}

/***********************************************
 * Entropy Coding Reset Neighbor Arrays
 ***********************************************/
static void entropy_coding_reset_neighbor_arrays(PictureControlSet* pcs, uint16_t tile_idx) {
    svt_aom_neighbor_array_unit_reset(pcs->partition_context_na[tile_idx]);

    svt_aom_neighbor_array_unit_reset(pcs->luma_dc_sign_level_coeff_na[tile_idx]);
    svt_aom_neighbor_array_unit_reset(pcs->cb_dc_sign_level_coeff_na[tile_idx]);
    svt_aom_neighbor_array_unit_reset(pcs->cr_dc_sign_level_coeff_na[tile_idx]);
    svt_aom_neighbor_array_unit_reset(pcs->txfm_context_array[tile_idx]);
}

/**************************************************
 * Reset Entropy Coding Picture
 **************************************************/
static void reset_entropy_coding_picture(EntropyCodingContext* ctx, PictureControlSet* pcs, SequenceControlSet* scs) {
    PictureParentControlSet* ppcs     = pcs->ppcs;
    const uint16_t           tile_cnt = ppcs->av1_cm->tiles_info.tile_rows * ppcs->av1_cm->tiles_info.tile_cols;
    ctx->is_16bit                     = scs->static_config.encoder_bit_depth > EB_EIGHT_BIT;
    const FrameHeader* frm_hdr        = &ppcs->frm_hdr;
    // Asuming cb and cr offset to be the same for chroma QP in both slice and pps for lambda computation
    const uint32_t entropy_coding_qp = frm_hdr->quantization_params.base_q_idx;

    for (uint16_t tile_idx = 0; tile_idx < tile_cnt; tile_idx++) {
        ppcs->prev_qindex[tile_idx] = entropy_coding_qp;
    }
    if (frm_hdr->allow_intrabc) {
        assert(frm_hdr->delta_lf_params.delta_lf_present == 0);
    }
    if (frm_hdr->delta_lf_params.delta_lf_present) {
        ppcs->prev_delta_lf_from_base = 0;

        const int frame_lf_count = ppcs->monochrome == 0 ? FRAME_LF_COUNT : FRAME_LF_COUNT - 2;
        for (int lf_id = 0; lf_id < frame_lf_count; ++lf_id) {
            ppcs->prev_delta_lf[lf_id] = 0;
        }
    }

    // pass the ent
    for (uint16_t tile_idx = 0; tile_idx < tile_cnt; tile_idx++) {
        EntropyCoder*        ec                   = pcs->ec_info[tile_idx]->ec;
        OutputBitstreamUnit* output_bitstream_ptr = ec->ec_output_bitstream_ptr;
        //****************************************************************//
        ec->ec_writer.allow_update_cdf = !ppcs->large_scale_tile && !frm_hdr->disable_cdf_update;
        aom_start_encode(&ec->ec_writer, output_bitstream_ptr);
        // ADD Reset here
        const uint8_t primary_ref_frame = frm_hdr->primary_ref_frame;
        if (primary_ref_frame != PRIMARY_REF_NONE) {
            // primary ref stored as REF_FRAME_MINUS1, while get_list_idx/get_ref_frame_idx take arg of ref frame
            // Therefore, add 1 to the primary ref frame (e.g. LAST --> LAST_FRAME)
            const uint8_t      list_idx = get_list_idx(primary_ref_frame + 1);
            const uint8_t      ref_idx  = get_ref_frame_idx(primary_ref_frame + 1);
            EbReferenceObject* ref      = (EbReferenceObject*)pcs->ref_pic_ptr_array[list_idx][ref_idx]->object_ptr;
            svt_memcpy(ec->fc, &ref->frame_context, sizeof(FRAME_CONTEXT));
        } else {
            svt_aom_reset_entropy_coder(scs->enc_ctx, ec, entropy_coding_qp, pcs->slice_type);
        }

        entropy_coding_reset_neighbor_arrays(pcs, tile_idx);
    }
}

/* Entropy Coding */

/*********************************************************************************
 *
 * @brief
 *  The Entropy Coding process is responsible for producing an AV1 conformant bitstream for each
 *frame.
 *
 * @par Description:
 *  The entropy coder is a frame-based process and is based on multi-symbol arithmetic range coding.
 *  It takes as input the coding decisions and information for each block and produces as output the
 *bitstream for each frame.
 *
 * @param[in] Coding Decisions
 *  Coding decisions and information for each block.
 *
 * @param[out] bitstream
 *  Bitstream for each block
 *
 ********************************************************************************/
EbErrorType svt_aom_entropy_coding_kernel_iter(void* context) {
    EntropyCodingContext* context_ptr = (EntropyCodingContext*)context;

    // Input
    EbObjectWrapper* rest_results_wrapper;

    // Output
    EbObjectWrapper* entropy_coding_results_wrapper_ptr;

    // Get Mode Decision Results
    EB_GET_FULL_OBJECT(context_ptr->enc_dec_input_fifo_ptr, &rest_results_wrapper);

    RestResults*        rest_results = (RestResults*)rest_results_wrapper->object_ptr;
    PictureControlSet*  pcs          = (PictureControlSet*)rest_results->pcs_wrapper->object_ptr;
    SequenceControlSet* scs          = pcs->scs;
    // SB Constants

    uint32_t sb_size = scs->sb_size;

    uint8_t          sb_size_log2    = (uint8_t)svt_log2f(sb_size);
    uint32_t         pic_width_in_sb = (pcs->ppcs->aligned_width + sb_size - 1) >> sb_size_log2;
    uint16_t         tile_idx        = rest_results->tile_index;
    Av1Common* const cm              = pcs->ppcs->av1_cm;
    const uint16_t   tile_cnt        = cm->tiles_info.tile_rows * cm->tiles_info.tile_cols;
    const uint16_t   tile_col        = tile_idx % cm->tiles_info.tile_cols;
    const uint16_t   tile_row        = tile_idx / cm->tiles_info.tile_cols;
    const uint16_t   tile_sb_start_x = cm->tiles_info.tile_col_start_mi[tile_col] >> scs->seq_header.sb_size_log2;
    const uint16_t   tile_sb_start_y = cm->tiles_info.tile_row_start_mi[tile_row] >> scs->seq_header.sb_size_log2;

    uint16_t tile_width_in_sb = (cm->tiles_info.tile_col_start_mi[tile_col + 1] -
                                 cm->tiles_info.tile_col_start_mi[tile_col]) >>
        scs->seq_header.sb_size_log2;
    uint16_t tile_height_in_sb = (cm->tiles_info.tile_row_start_mi[tile_row + 1] -
                                  cm->tiles_info.tile_row_start_mi[tile_row]) >>
        scs->seq_header.sb_size_log2;

    bool frame_entropy_done = false;

    svt_block_on_mutex(pcs->entropy_coding_pic_mutex);
    if (pcs->entropy_coding_pic_reset_flag) {
        pcs->entropy_coding_pic_reset_flag = false;

        reset_entropy_coding_picture(context_ptr, pcs, scs);
    }
    svt_release_mutex(pcs->entropy_coding_pic_mutex);

    if (!svt_aom_is_pic_skipped(pcs->ppcs)) {
        context_ptr->tot_qindex = 0;
        context_ptr->valid_area = 0;
        for (uint32_t y_sb_index = 0; y_sb_index < tile_height_in_sb; ++y_sb_index) {
            for (uint32_t x_sb_index = 0; x_sb_index < tile_width_in_sb; ++x_sb_index) {
                uint16_t    sb_index = (uint16_t)((x_sb_index + tile_sb_start_x) +
                                               (y_sb_index + tile_sb_start_y) * pic_width_in_sb);
                SuperBlock* sb_ptr   = pcs->sb_ptr_array[sb_index];

                const uint32_t sb_origin_x = (x_sb_index + tile_sb_start_x) << sb_size_log2;
                const uint32_t sb_origin_y = (y_sb_index + tile_sb_start_y) << sb_size_log2;
                if (x_sb_index == 0 && y_sb_index == 0) {
                    svt_av1_reset_loop_restoration(context_ptr);
                    context_ptr->tok = pcs->tile_tok[tile_row][tile_col];
                }

                EbPictureBufferDesc* coeff_picture_ptr = pcs->ppcs->enc_dec_ptr->quantized_coeff[sb_index];
                context_ptr->coded_area_sb             = 0;
                context_ptr->coded_area_sb_uv          = 0;
                // Ensure EC buffer has room for worst-case SB output (4 bytes/pixel)
                EbErrorType ret = svt_aom_ec_ensure_capacity(&pcs->ec_info[tile_idx]->ec->ec_writer,
                                                             sb_size * sb_size * 4);
                if (ret != EB_ErrorNone) {
                    return ret;
                }
                svt_aom_write_modes_sb(context_ptr,
                                       sb_ptr,
                                       pcs,
                                       tile_idx,
                                       pcs->ec_info[tile_idx]->ec,
                                       coeff_picture_ptr,
                                       sb_ptr->ptree,
                                       sb_origin_y >> MI_SIZE_LOG2,
                                       sb_origin_x >> MI_SIZE_LOG2);
            }
        }
    }
    bool pic_ready = true;

    // Current tile ready
    svt_aom_encode_slice_finish(pcs->ec_info[tile_idx]->ec);

    svt_block_on_mutex(pcs->entropy_coding_pic_mutex);
    // Flush locally-accumulated qindex stats (avoids per-block mutex)
    pcs->ppcs->tot_qindex += context_ptr->tot_qindex;
    pcs->ppcs->valid_qindex_area += context_ptr->valid_area;
    pcs->ec_info[tile_idx]->entropy_coding_tile_done = true;
    for (uint16_t i = 0; i < tile_cnt; i++) {
        if (pcs->ec_info[i]->entropy_coding_tile_done == false) {
            pic_ready = false;
            break;
        }
    }
    svt_release_mutex(pcs->entropy_coding_pic_mutex);
    if (pic_ready) {
        if (pcs->ppcs->superres_total_recode_loop == 0) {
            // Release the reference Pictures from both lists
            for (REF_FRAME_MINUS1 ref = LAST; ref < ALT + 1; ref++) {
                const uint8_t list_idx = get_list_idx(ref + 1);
                const uint8_t ref_idx  = get_ref_frame_idx(ref + 1);
                if (pcs->ref_pic_ptr_array[list_idx][ref_idx] != NULL) {
                    svt_release_object(pcs->ref_pic_ptr_array[list_idx][ref_idx]);
                }
            }

            //free palette data
            if (pcs->tile_tok[0][0]) {
                EB_FREE_ARRAY(pcs->tile_tok[0][0]);
            }
        }
        frame_entropy_done = true;
    }

    if (frame_entropy_done) {
        if (pcs->ppcs->valid_qindex_area) {
            pcs->ppcs->avg_qp = ((pcs->ppcs->tot_qindex / pcs->ppcs->valid_qindex_area) + 2) >> 2;
        }
        // Get Empty Entropy Coding Results
        svt_get_empty_object(context_ptr->entropy_coding_output_fifo_ptr, &entropy_coding_results_wrapper_ptr);
        EntropyCodingResults* entropy_coding_results_ptr = (EntropyCodingResults*)
                                                               entropy_coding_results_wrapper_ptr->object_ptr;
        entropy_coding_results_ptr->pcs_wrapper = rest_results->pcs_wrapper;

        // Post EntropyCoding Results
        svt_post_full_object(entropy_coding_results_wrapper_ptr);
    }

    // Release Mode Decision Results
    svt_release_object(rest_results_wrapper);
    return EB_ErrorNone;
}

void* svt_aom_entropy_coding_kernel(void* input_ptr) {
    EbThreadContext* thread_ctx = (EbThreadContext*)input_ptr;
    for (;;) {
        EbErrorType err = svt_aom_entropy_coding_kernel_iter(thread_ctx->priv);
        if (err == EB_NoErrorFifoShutdown) {
            return NULL;
        }
    }
    return NULL;
}
