/*

* 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 "enc_handle.h"

#include "utility.h"

#include "pcs.h"

#include "pd_results.h"

#include "me_process.h"

#include "me_results.h"

#include "reference_object.h"

#include "motion_estimation.h"

#include "lambda_rate_tables.h"

#include "compute_sad.h"

#ifdef ARCH_X86_64

#include <emmintrin.h>

#endif

#include "temporal_filtering.h"

#include "global_me.h"

#include "resize.h"

#include "pic_demux_results.h"

#include "rc_tasks.h"

#include "firstpass.h"

#include "initial_rc_process.h"

#include "enc_mode_config.h"

/* --32x32-

|00||01|

|02||03|

--------*/

/* ------16x16-----

|00||01||04||05|

|02||03||06||07|

|08||09||12||13|

|10||11||14||15|

----------------*/

/* ------8x8----------------------------

|00||01||04||05|     |16||17||20||21|

|02||03||06||07|     |18||19||22||23|

|08||09||12||13|     |24||25||28||29|

|10||11||14||15|     |26||27||30||31|

|32||33||36||37|     |48||49||52||53|

|34||35||38||39|     |50||51||54||55|

|40||41||44||45|     |56||57||60||61|

|42||43||46||47|     |58||59||62||63|

-------------------------------------*/

void dg_detector_hme_level0(PictureParentControlSet* ppcs, uint32_t seg_idx);

static void motion_estimation_context_dctor(EbPtr p) {

    EbThreadContext*           thread_ctx = (EbThreadContext*)p;

    MotionEstimationContext_t* obj        = (MotionEstimationContext_t*)thread_ctx->priv;

    EB_DELETE(obj->me_ctx);

    EB_FREE_ARRAY(obj);

}

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

 * Motion Analysis Context Constructor

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

EbErrorType svt_aom_motion_estimation_context_ctor(EbThreadContext* thread_ctx, const EbEncHandle* enc_handle_ptr,

                                                   int index) {

    MotionEstimationContext_t* me_context_ptr;

    EB_CALLOC_ARRAY(me_context_ptr, 1);

    thread_ctx->priv                                        = me_context_ptr;

    thread_ctx->dctor                                       = motion_estimation_context_dctor;

    me_context_ptr->picture_decision_results_input_fifo_ptr = svt_system_resource_get_consumer_fifo(

        enc_handle_ptr->picture_decision_results_resource_ptr, index);

    me_context_ptr->motion_estimation_results_output_fifo_ptr = svt_system_resource_get_producer_fifo(

        enc_handle_ptr->motion_estimation_results_resource_ptr, index);

    EB_NEW(me_context_ptr->me_ctx, svt_aom_me_context_ctor);

    return EB_ErrorNone;

}

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

 * Motion Analysis Kernel

 * The Motion Analysis performs  Motion Estimation

 * This process has access to the current input picture as well as

 * the input pictures, which the current picture references according

 * to the prediction structure pattern.  The Motion Analysis process is multithreaded,

 * so pictures can be processed out of order as long as all inputs are available.

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

EbErrorType svt_aom_motion_estimation_kernel_iter(void* context) {

    MotionEstimationContext_t* me_context_ptr = (MotionEstimationContext_t*)context;

    EbObjectWrapper* in_results_wrapper_ptr;

    EbObjectWrapper* out_results_wrapper;

    // Get Input Full Object

    EB_GET_FULL_OBJECT(me_context_ptr->picture_decision_results_input_fifo_ptr, &in_results_wrapper_ptr);

    PictureDecisionResults*  in_results_ptr = (PictureDecisionResults*)in_results_wrapper_ptr->object_ptr;

    PictureParentControlSet* pcs            = (PictureParentControlSet*)in_results_ptr->pcs_wrapper->object_ptr;

    SequenceControlSet*      scs            = pcs->scs;

    if (in_results_ptr->task_type == TASK_TFME) {

        me_context_ptr->me_ctx->me_type = ME_MCTF;

    } else if (in_results_ptr->task_type == TASK_PAME || in_results_ptr->task_type == TASK_SUPERRES_RE_ME) {

        me_context_ptr->me_ctx->me_type = ME_OPEN_LOOP;

    } else if (in_results_ptr->task_type == TASK_DG_DETECTOR_HME) {

        me_context_ptr->me_ctx->me_type = ME_DG_DETECTOR;

    }

    // ME Kernel Signal(s) derivation

    if ((in_results_ptr->task_type == TASK_PAME) || (in_results_ptr->task_type == TASK_SUPERRES_RE_ME)) {

        svt_aom_sig_deriv_me(scs, pcs, me_context_ptr->me_ctx);

    }

    else if (in_results_ptr->task_type == TASK_TFME) {

        svt_aom_sig_deriv_me_tf(pcs, me_context_ptr->me_ctx);

    }

    if ((in_results_ptr->task_type == TASK_PAME) || (in_results_ptr->task_type == TASK_SUPERRES_RE_ME)) {

        EbPictureBufferDesc* sixteenth_picture_ptr;

        EbPictureBufferDesc* quarter_picture_ptr;

        EbPictureBufferDesc* input_padded_pic;

        EbPictureBufferDesc* input_pic;

        EbPaReferenceObject* pa_ref_obj_;

        //assert((int)pcs->pa_ref_pic_wrapper->live_count > 0);

        pa_ref_obj_ = (EbPaReferenceObject*)pcs->pa_ref_pic_wrapper->object_ptr;

        // Set 1/4 and 1/16 ME input buffer(s); filtered or decimated

        quarter_picture_ptr   = pa_ref_obj_->quarter_downsampled_picture_ptr;

        sixteenth_picture_ptr = pa_ref_obj_->sixteenth_downsampled_picture_ptr;

        input_padded_pic      = pa_ref_obj_->input_padded_pic;

        input_pic = pcs->enhanced_pic;

        // Segments

        uint32_t segment_index         = in_results_ptr->segment_index;

        uint32_t pic_width_in_b64      = (pcs->aligned_width + scs->b64_size - 1) / scs->b64_size;

        uint32_t picture_height_in_b64 = (pcs->aligned_height + scs->b64_size - 1) / scs->b64_size;

        uint32_t y_segment_index;

        uint32_t x_segment_index;

        SEGMENT_CONVERT_IDX_TO_XY(segment_index, x_segment_index, y_segment_index, pcs->me_segments_column_count);

        uint32_t x_b64_start_index = SEGMENT_START_IDX(

            x_segment_index, pic_width_in_b64, pcs->me_segments_column_count);

        uint32_t x_b64_end_index   = SEGMENT_END_IDX(x_segment_index, pic_width_in_b64, pcs->me_segments_column_count);

        uint32_t y_b64_start_index = SEGMENT_START_IDX(

            y_segment_index, picture_height_in_b64, pcs->me_segments_row_count);

        uint32_t y_b64_end_index = SEGMENT_END_IDX(y_segment_index, picture_height_in_b64, pcs->me_segments_row_count);

        bool skip_me = false;

        if (svt_aom_is_pic_skipped(pcs)) {

            skip_me = true;

        }

        // skip me for the first pass. ME is already performed

        if (!skip_me) {

            if (pcs->slice_type != I_SLICE) {

                // Use scaled source references if resolution of the reference is different that of the input

                svt_aom_use_scaled_source_refs_if_needed(

                    pcs, input_pic, pa_ref_obj_, &input_padded_pic, &quarter_picture_ptr, &sixteenth_picture_ptr);

                // 64x64 Block Loop

                for (uint32_t y_b64_index = y_b64_start_index; y_b64_index < y_b64_end_index; ++y_b64_index) {

                    for (uint32_t x_b64_index = x_b64_start_index; x_b64_index < x_b64_end_index; ++x_b64_index) {

                        uint32_t b64_index = (uint16_t)(x_b64_index + y_b64_index * pic_width_in_b64);

                        uint32_t b64_origin_x = x_b64_index * scs->b64_size;

                        uint32_t b64_origin_y = y_b64_index * scs->b64_size;

                        // Load the 64x64 Block from the input to the intermediate block buffer

                        uint32_t buffer_index = (b64_origin_y)*input_pic->y_stride + b64_origin_x;

#ifdef ARCH_X86_64

                        uint8_t* src_ptr    = &input_padded_pic->y_buffer[buffer_index];

                        uint32_t b64_height = (pcs->aligned_height - b64_origin_y) < BLOCK_SIZE_64

                            ? pcs->aligned_height - b64_origin_y

                            : BLOCK_SIZE_64;

                        //_MM_HINT_T0     //_MM_HINT_T1    //_MM_HINT_T2//_MM_HINT_NTA

                        for (uint32_t i = 0; i < b64_height; i++) {

                            char const* p = (char const*)(src_ptr + i * input_padded_pic->y_stride);

                            _mm_prefetch(p, _MM_HINT_T2);

                        }

#endif

                        me_context_ptr->me_ctx->b64_src_ptr    = &input_padded_pic->y_buffer[buffer_index];

                        me_context_ptr->me_ctx->b64_src_stride = input_padded_pic->y_stride;

                        // Load the 1/4 decimated SB from the 1/4 decimated input to the 1/4 intermediate SB buffer

                        if (me_context_ptr->me_ctx->enable_hme_level1_flag) {

                            buffer_index = ((b64_origin_y >> 1)) * quarter_picture_ptr->y_stride + (b64_origin_x >> 1);

                            me_context_ptr->me_ctx->quarter_b64_buffer = &quarter_picture_ptr->y_buffer[buffer_index];

                            me_context_ptr->me_ctx->quarter_b64_buffer_stride = quarter_picture_ptr->y_stride;

                        }

                        // Load the 1/16 decimated SB from the 1/16 decimated input to the 1/16 intermediate SB buffer

                        if (me_context_ptr->me_ctx->enable_hme_level0_flag) {

                            buffer_index = ((b64_origin_y >> 2)) * sixteenth_picture_ptr->y_stride +

                                (b64_origin_x >> 2);

                            me_context_ptr->me_ctx->sixteenth_b64_buffer =

                                &sixteenth_picture_ptr->y_buffer[buffer_index];

                            me_context_ptr->me_ctx->sixteenth_b64_buffer_stride = sixteenth_picture_ptr->y_stride;

                        }

                        me_context_ptr->me_ctx->me_type = ME_OPEN_LOOP;

                        if ((in_results_ptr->task_type == TASK_PAME) ||

                            (in_results_ptr->task_type == TASK_SUPERRES_RE_ME)) {

                            me_context_ptr->me_ctx->num_of_list_to_search = MAX_NUM_OF_REF_PIC_LIST;

                            me_context_ptr->me_ctx->num_of_ref_pic_to_search[0] = pcs->ref_list0_count_try;

                            me_context_ptr->me_ctx->num_of_ref_pic_to_search[1] = pcs->ref_list1_count_try;

                            me_context_ptr->me_ctx->temporal_layer_index        = pcs->temporal_layer_index;

                            me_context_ptr->me_ctx->is_ref                      = pcs->is_ref;

                            if (pcs->frame_superres_enabled || pcs->frame_resize_enabled) {

                                for (int i = 0; i < me_context_ptr->me_ctx->num_of_list_to_search; i++) {

                                    for (int j = 0; j < me_context_ptr->me_ctx->num_of_ref_pic_to_search[i]; j++) {

                                        //assert((int)pcs->ref_pa_pic_ptr_array[i][j]->live_count > 0);

                                        uint8_t sr_denom_idx     = svt_aom_get_denom_idx(pcs->superres_denom);

                                        uint8_t resize_denom_idx = svt_aom_get_denom_idx(pcs->resize_denom);

                                        EbPaReferenceObject* ref_object =

                                            (EbPaReferenceObject*)pcs->ref_pa_pic_ptr_array[i][j]->object_ptr;

                                        me_context_ptr->me_ctx->me_ds_ref_array[i][j].picture_ptr =

                                            ref_object

                                                ->downscaled_input_padded_picture_ptr[sr_denom_idx][resize_denom_idx];

                                        me_context_ptr->me_ctx->me_ds_ref_array[i][j].quarter_picture_ptr =

                                            ref_object->downscaled_quarter_downsampled_picture_ptr[sr_denom_idx]

                                                                                                  [resize_denom_idx];

                                        me_context_ptr->me_ctx->me_ds_ref_array[i][j].sixteenth_picture_ptr =

                                            ref_object->downscaled_sixteenth_downsampled_picture_ptr[sr_denom_idx]

                                                                                                    [resize_denom_idx];

                                        me_context_ptr->me_ctx->me_ds_ref_array[i][j].picture_number =

                                            ref_object->picture_number;

                                    }

                                }

                            } else {

                                for (int i = 0; i < me_context_ptr->me_ctx->num_of_list_to_search; i++) {

                                    for (int j = 0; j < me_context_ptr->me_ctx->num_of_ref_pic_to_search[i]; j++) {

                                        //assert((int)pcs->ref_pa_pic_ptr_array[i][j]->live_count > 0);

                                        EbPaReferenceObject* ref_object =

                                            (EbPaReferenceObject*)pcs->ref_pa_pic_ptr_array[i][j]->object_ptr;

                                        me_context_ptr->me_ctx->me_ds_ref_array[i][j].picture_ptr =

                                            ref_object->input_padded_pic;

                                        me_context_ptr->me_ctx->me_ds_ref_array[i][j].quarter_picture_ptr =

                                            ref_object->quarter_downsampled_picture_ptr;

                                        me_context_ptr->me_ctx->me_ds_ref_array[i][j].sixteenth_picture_ptr =

                                            ref_object->sixteenth_downsampled_picture_ptr;

                                        me_context_ptr->me_ctx->me_ds_ref_array[i][j].picture_number =

                                            ref_object->picture_number;

                                    }

                                }

                            }

                        }

                        svt_aom_motion_estimation_b64(

                            pcs, b64_index, b64_origin_x, b64_origin_y, me_context_ptr->me_ctx, input_pic);

                        if ((in_results_ptr->task_type == TASK_PAME) ||

                            (in_results_ptr->task_type == TASK_SUPERRES_RE_ME)) {

                            svt_block_on_mutex(pcs->me_processed_b64_mutex);

                            pcs->me_processed_b64_count++;

                            // We need to finish ME for all SBs to do GM

                            if (pcs->me_processed_b64_count == pcs->b64_total_count) {

                                if (pcs->gm_ctrls.enabled && (!pcs->gm_ctrls.pp_enabled || pcs->gm_pp_detected)) {

                                    svt_aom_global_motion_estimation(pcs, input_pic);

                                } else {

                                    // Initilize global motion to be OFF when GM is OFF

                                    memset(pcs->is_global_motion, false, MAX_NUM_OF_REF_PIC_LIST * REF_LIST_MAX_DEPTH);

                                }

                            }

                            svt_release_mutex(pcs->me_processed_b64_mutex);

                        }

                    }

                }

            }

        }

        // Get Empty Results Object

        svt_get_empty_object(me_context_ptr->motion_estimation_results_output_fifo_ptr, &out_results_wrapper);

        MotionEstimationResults* out_results = (MotionEstimationResults*)out_results_wrapper->object_ptr;

        out_results->pcs_wrapper             = in_results_ptr->pcs_wrapper;

        out_results->segment_index           = segment_index;

        out_results->task_type               = in_results_ptr->task_type;

        // Release the Input Results

        svt_release_object(in_results_wrapper_ptr);

        // Post the Full Results Object

        svt_post_full_object(out_results_wrapper);

    } else if (in_results_ptr->task_type == TASK_TFME) {

        //gm pre-processing for only base B

        if (pcs->gm_ctrls.pp_enabled && pcs->gm_pp_enabled && in_results_ptr->segment_index == 0) {

            svt_aom_gm_pre_processor(pcs, pcs->temp_filt_pcs_list);

        }

        // temporal filtering start

        me_context_ptr->me_ctx->me_type = ME_MCTF;

        svt_av1_init_temporal_filtering(pcs->temp_filt_pcs_list, pcs, me_context_ptr, in_results_ptr->segment_index);

        // Release the Input Results

        svt_release_object(in_results_wrapper_ptr);

    } else if (in_results_ptr->task_type == TASK_DG_DETECTOR_HME) {

        // dynamic gop detector

        dg_detector_hme_level0(pcs, in_results_ptr->segment_index);

        // Release the Input Results

        svt_release_object(in_results_wrapper_ptr);

    }

    return EB_ErrorNone;

}

void* svt_aom_motion_estimation_kernel(void* input_ptr) {

    EbThreadContext* thread_ctx = (EbThreadContext*)input_ptr;

    for (;;) {

        EbErrorType err = svt_aom_motion_estimation_kernel_iter(thread_ctx->priv);

        if (err == EB_NoErrorFifoShutdown) {

            return NULL;

        }

    }

    return NULL;

}