/*

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

#include "pcs.h"

#include "utility.h"

#include "enc_mode_config.h"

void svt_aom_largest_coding_unit_dctor(EbPtr p) {

    // av1xd / final_blk_arr / ptree are borrowed from per-PCS pools (freed in the PCS dctor).

    (void)p;

}

static void setup_ptree(PARTITION_TREE* pc_tree, int index, BlockSize bsize, const int min_sq_size) {

    pc_tree->bsize = bsize;

    pc_tree->index = index;

    // 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->sub_tree[i] = pc_tree + i * blocks_to_skip + 1;

            setup_ptree(pc_tree->sub_tree[i], i, subsize, min_sq_size);

        }

    }

}

/*

Tasks & Questions

    -Need a GetEmptyChain function for testing sub partitions.  Tie it to an Itr?

    -How many empty CUs do we need?  We need to have enough for the max CU count AND

       enough for temp storage when calculating a subdivision.

    -Where do we store temp reconstructed picture data while deciding to subdivide or not?

    -Need a ReconPicture for each candidate.

    -I don't see a way around doing the copies in temp memory and then copying it in...

*/

EbErrorType svt_aom_largest_coding_unit_ctor(SuperBlock* larget_coding_unit_ptr, uint8_t sb_size_pix,

                                             uint16_t sb_origin_x, uint16_t sb_origin_y, uint16_t sb_index,

                                             EncMode enc_mode, bool rtc, bool allintra, PictureControlSet* pcs) {

    larget_coding_unit_ptr->dctor = svt_aom_largest_coding_unit_dctor;

    // ************ SB ***************

    // Which borderLargestCuSize is not a power of two

    // Which borderLargestCuSize is not a power of two

    larget_coding_unit_ptr->pcs = pcs;

    larget_coding_unit_ptr->org_x = sb_origin_x;

    larget_coding_unit_ptr->org_y = sb_origin_y;

    larget_coding_unit_ptr->index = sb_index;

    bool disallow_sub_8x8_nsq     = true;

    bool disallow_sub_16x16_nsq   = true;

    for (uint8_t coeff_lvl = 0; coeff_lvl <= HIGH_LVL + 1; coeff_lvl++) {

        uint8_t nsq_geom_lvl = allintra ? svt_aom_get_nsq_geom_level_allintra(enc_mode)

            : rtc                       ? svt_aom_get_nsq_geom_level_rtc()

                                        : svt_aom_get_nsq_geom_level_default(enc_mode, coeff_lvl);

        // nsq_geom_lvl level 0 means NSQ shapes are disallowed so don't adjust based on the level

        if (nsq_geom_lvl) {

            uint8_t allow_HVA_HVB, allow_HV4, min_nsq_bsize;

            svt_aom_set_nsq_geom_ctrls(NULL, nsq_geom_lvl, &allow_HVA_HVB, &allow_HV4, &min_nsq_bsize);

            if (min_nsq_bsize < 8) {

                disallow_sub_8x8_nsq = false;

            }

            if (min_nsq_bsize < 16) {

                disallow_sub_16x16_nsq = false;

            }

        }

    }

    bool     disallow_4x4 = allintra ? svt_aom_get_disallow_4x4_allintra(enc_mode)

            : rtc                    ? svt_aom_get_disallow_4x4_rtc()

                                     : svt_aom_get_disallow_4x4_default(enc_mode);

    bool     disallow_8x8 = allintra ? svt_aom_get_disallow_8x8_allintra()

            : rtc                    ? svt_aom_get_disallow_8x8_rtc(enc_mode, pcs->frame_width, pcs->frame_height)

                                     : svt_aom_get_disallow_8x8_default();

    uint32_t tot_blk_num;

    if (sb_size_pix == 128) {

        if (disallow_8x8 && disallow_sub_16x16_nsq) {

            tot_blk_num = 64;

        } else if (disallow_8x8 || (disallow_4x4 && disallow_sub_8x8_nsq)) {

            tot_blk_num = 256;

        } else if (disallow_4x4) {

            tot_blk_num = 512;

        } else {

            tot_blk_num = 1024;

        }

    } else if (disallow_8x8 && disallow_sub_16x16_nsq) {

        tot_blk_num = 16;

    } else if (disallow_8x8 || (disallow_4x4 && disallow_sub_8x8_nsq)) {

        tot_blk_num = 64;

    } else if (disallow_4x4) {

        tot_blk_num = 128;

    } else {

        tot_blk_num = 256;

    }

    // Do NOT initialize the final_blk_arr here

    // Malloc maximum but only initialize it only when actually used.

    // This will help to same actually memory usage

    // Alloc ptree, which is used to store final block data/mode info for the SB that is passed

    // from encdec to EC

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

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

    int     blocks_to_alloc  = 0;

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

        blocks_to_alloc += blocks_per_depth;

    }

    // All SBs in a PCS resolve to identical tot_blk_num / blocks_to_alloc, so back the whole

    // sb_ptr_array with one allocation each (created on the first SB) and hand each SB a slice,

    // instead of 3 allocations per SB.

    const uint16_t all_sb = pcs->sb_total_count;

    if (sb_index == 0) {

        // This ctor is re-run per frame on the resize/superres path (superres_setup_child_pcs),

        // so free any pool from a prior setup before reallocating (NULL-safe on first use).

        EB_FREE_ARRAY(pcs->sb_final_blk_arr_pool);

        EB_FREE_ARRAY(pcs->sb_av1xd_pool);

        EB_FREE_ARRAY(pcs->sb_ptree_pool);

        EB_MALLOC_ARRAY(pcs->sb_final_blk_arr_pool, (size_t)all_sb * tot_blk_num);

        EB_MALLOC_ARRAY(pcs->sb_av1xd_pool, all_sb);

        EB_CALLOC_ARRAY(pcs->sb_ptree_pool, (size_t)all_sb * blocks_to_alloc);

    }

    larget_coding_unit_ptr->final_blk_arr = pcs->sb_final_blk_arr_pool + (size_t)sb_index * tot_blk_num;

    larget_coding_unit_ptr->av1xd         = pcs->sb_av1xd_pool + sb_index;

    larget_coding_unit_ptr->ptree         = pcs->sb_ptree_pool + (size_t)sb_index * blocks_to_alloc;

    setup_ptree(larget_coding_unit_ptr->ptree, 0, sb_size_pix == 128 ? BLOCK_128X128 : BLOCK_64X64, min_bsize);

    return EB_ErrorNone;

}