/*
* 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 "EbSegmentation.h"
#include "EbSegmentationParams.h"
#include "EbMotionEstimationContext.h"
#include "common_dsp_rtcd.h"
#if DEBUG_SEGMENT_QP
#include "EbLog.h"
#endif

static uint16_t get_variance_for_cu(const BlockGeom *blk_geom, uint16_t *variance_ptr) {
    int index0, index1;
    //Assumes max CU size is 64
    switch (blk_geom->bsize) {
    case BLOCK_4X4:
    case BLOCK_4X8:
    case BLOCK_8X4:
    case BLOCK_8X8:
        index0 = index1 = ME_TIER_ZERO_PU_8x8_0 + ((blk_geom->org_x >> 3) + blk_geom->org_y);
        break;

    case BLOCK_8X16:
        index0 = ME_TIER_ZERO_PU_8x8_0 + ((blk_geom->org_x >> 3) + blk_geom->org_y);
        index1 = index0 + 1;
        break;

    case BLOCK_16X8:
        index0 = ME_TIER_ZERO_PU_8x8_0 + ((blk_geom->org_x >> 3) + blk_geom->org_y);
        index1 = index0 + blk_geom->org_y;
        break;

    case BLOCK_4X16:
    case BLOCK_16X4:
    case BLOCK_16X16:
        index0 = index1 = ME_TIER_ZERO_PU_16x16_0 +
            ((blk_geom->org_x >> 4) + (blk_geom->org_y >> 2));
        break;

    case BLOCK_16X32:
        index0 = ME_TIER_ZERO_PU_16x16_0 + ((blk_geom->org_x >> 4) + (blk_geom->org_y >> 2));
        index1 = index0 + 1;
        break;

    case BLOCK_32X16:
        index0 = ME_TIER_ZERO_PU_16x16_0 + ((blk_geom->org_x >> 4) + (blk_geom->org_y >> 2));
        index1 = index0 + (blk_geom->org_y >> 2);
        break;

    case BLOCK_8X32:
    case BLOCK_32X8:
    case BLOCK_32X32:
        index0 = index1 = ME_TIER_ZERO_PU_32x32_0 +
            ((blk_geom->org_x >> 5) + (blk_geom->org_y >> 4));
        break;

    case BLOCK_32X64:
        index0 = ME_TIER_ZERO_PU_32x32_0 + ((blk_geom->org_x >> 5) + (blk_geom->org_y >> 4));
        index1 = index0 + 1;
        break;

    case BLOCK_64X32:
        index0 = ME_TIER_ZERO_PU_32x32_0 + ((blk_geom->org_x >> 5) + (blk_geom->org_y >> 4));
        index1 = index0 + (blk_geom->org_y >> 4);
        break;

    case BLOCK_64X64:
    case BLOCK_16X64:
    case BLOCK_64X16:
    default: index0 = index1 = 0; break;
    }
    return (variance_ptr[index0] + variance_ptr[index1]) >> 1;
}

void svt_aom_apply_segmentation_based_quantization(const BlockGeom   *blk_geom,
                                                   PictureControlSet *pcs, SuperBlock *sb_ptr,
                                                   BlkStruct *blk_ptr) {
    uint16_t           *variance_ptr        = pcs->ppcs->variance[sb_ptr->index];
    SegmentationParams *segmentation_params = &pcs->ppcs->frm_hdr.segmentation_params;
    uint16_t            variance            = get_variance_for_cu(blk_geom, variance_ptr);
    blk_ptr->segment_id                     = 0;
    for (int i = MAX_SEGMENTS - 1; i >= 0; i--) {
        if (variance <= segmentation_params->variance_bin_edge[i]) {
            int32_t q_index = pcs->ppcs->frm_hdr.quantization_params.base_q_idx +
                segmentation_params->feature_data[i][SEG_LVL_ALT_Q];
            // Avoid lossless since SVT-AV1 doesn't support it.
            // Spec: Uncompressed header syntax and get_qindex(1, segmentID). And spec 5.11.34. Residual syntax. force TX_4X4 when lossless.
            if (q_index > 0) {
                blk_ptr->segment_id = i;
                break;
            }
        }
    }
}

void svt_aom_setup_segmentation(PictureControlSet *pcs, SequenceControlSet *scs) {
    SegmentationParams *segmentation_params = &pcs->ppcs->frm_hdr.segmentation_params;
    segmentation_params->segmentation_enabled =
        (Bool)(scs->static_config.enable_adaptive_quantization == 1);
    if (segmentation_params->segmentation_enabled) {
        segmentation_params->segmentation_update_data =
            1; //always updating for now. Need to set this based on actual deltas
        segmentation_params->segmentation_update_map = 1;
        segmentation_params->segmentation_temporal_update =
            FALSE; //!(pcs->ppcs->av1FrameType == KEY_FRAME || pcs->ppcs->av1FrameType == INTRA_ONLY_FRAME);
        find_segment_qps(segmentation_params, pcs);
        for (int i = 0; i < MAX_SEGMENTS; i++)
            segmentation_params->feature_enabled[i][SEG_LVL_ALT_Q] = 1;

        calculate_segmentation_data(segmentation_params);
    }
}

void calculate_segmentation_data(SegmentationParams *segmentation_params) {
    for (int i = 0; i < MAX_SEGMENTS; i++) {
        for (int j = 0; j < SEG_LVL_MAX; j++) {
            if (segmentation_params->feature_enabled[i][j]) {
                segmentation_params->last_active_seg_id = i;
                if (j >= SEG_LVL_REF_FRAME) {
                    segmentation_params->seg_id_pre_skip = 1;
                }
            }
        }
    }
}

void find_segment_qps(SegmentationParams *segmentation_params,
                      PictureControlSet  *pcs) { //QP needs to be specified as qpindex, not qp.
    uint16_t    min_var = UINT16_MAX, max_var = MIN_UNSIGNED_VALUE, avg_var = 0;
    const float strength = 2; //to tune

    // get range of variance
    for (uint32_t sb_idx = 0; sb_idx < pcs->b64_total_count; ++sb_idx) {
        uint16_t *variance_ptr = pcs->ppcs->variance[sb_idx];
        uint32_t  var_index, local_avg = 0;
        // Loop over all 8x8s in a 64x64
        for (var_index = ME_TIER_ZERO_PU_8x8_0; var_index <= ME_TIER_ZERO_PU_8x8_63; var_index++) {
            max_var = MAX(max_var, variance_ptr[var_index]);
            min_var = MIN(min_var, variance_ptr[var_index]);
            local_avg += variance_ptr[var_index];
        }
        avg_var += (local_avg >> 6);
    }
    avg_var /= pcs->b64_total_count;
    avg_var = svt_log2f(avg_var);

    //get variance bin edges & QPs
    uint16_t min_var_log = svt_log2f(MAX(1, min_var));
    uint16_t max_var_log = svt_log2f(MAX(1, max_var));
    uint16_t step_size   = (uint16_t)(max_var_log - min_var_log) <= MAX_SEGMENTS
          ? 1
          : ROUND(((max_var_log - min_var_log)) / MAX_SEGMENTS);
    uint16_t bin_edge    = min_var_log + step_size;
    uint16_t bin_center  = bin_edge >> 1;

    for (int i = MAX_SEGMENTS - 1; i >= 0; i--) {
        segmentation_params->variance_bin_edge[i]           = POW2(bin_edge);
        segmentation_params->feature_data[i][SEG_LVL_ALT_Q] = ROUND((uint16_t)strength *
                                                                    (MAX(1, bin_center) - avg_var));
        bin_edge += step_size;
        bin_center += step_size;
    }
    if (segmentation_params->feature_data[0][SEG_LVL_ALT_Q] < 0) {
        // avoid lossless block
        segmentation_params->feature_data[0][SEG_LVL_ALT_Q] = 0;
    }
#if DEBUG_SEGMENT_QP
    SVT_LOG("frame %d, base_qindex %d, seg qp offset: %d %d %d %d %d %d %d %d\n",
            (int)pcs->picture_number,
            pcs->ppcs->frm_hdr.quantization_params.base_q_idx,
            segmentation_params->feature_data[0][SEG_LVL_ALT_Q],
            segmentation_params->feature_data[1][SEG_LVL_ALT_Q],
            segmentation_params->feature_data[2][SEG_LVL_ALT_Q],
            segmentation_params->feature_data[3][SEG_LVL_ALT_Q],
            segmentation_params->feature_data[4][SEG_LVL_ALT_Q],
            segmentation_params->feature_data[5][SEG_LVL_ALT_Q],
            segmentation_params->feature_data[6][SEG_LVL_ALT_Q],
            segmentation_params->feature_data[7][SEG_LVL_ALT_Q]);
#endif
}
