/*
 * Copyright (c) 2020, 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 www.aomedia.org/license/software. 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 www.aomedia.org/license/patent.
 */

#ifndef AOM_AV1_ENCODER_PARTITION_SEARCH_UTILS_H_
#define AOM_AV1_ENCODER_PARTITION_SEARCH_UTILS_H_

#include "av1/encoder/encodemb.h"
#include "av1/encoder/encoder.h"
#include "av1/encoder/encodeframe_utils.h"

#ifdef __cplusplus
extern "C" {
#endif

#if !CONFIG_REALTIME_ONLY
// Rectangular partition types.
enum { HORZ = 0, VERT, NUM_RECT_PARTS } UENUM1BYTE(RECT_PART_TYPE);

enum { PART_A = 0, PART_B, NUM_AB_TYPES } UENUM1BYTE(AB_PART_TYPE);
#define SUB_PARTITIONS_SPLIT 4
#define SUB_PARTITIONS_RECT 2

// This structure contains block size related
// variables for use in rd_pick_partition().
typedef struct {
  // Half of block width/height to determine block edge.
  int mi_step_w;
  int mi_step_h;

  // Block row and column indices.
  int mi_row;
  int mi_col;

  // Block width of current partition block.
  int width;

  BLOCK_SIZE max_sq_part;
  BLOCK_SIZE min_sq_part;

  // Block width of maximum partition size allowed.
  int max_partition_size_1d;
  // Block width of minimum partition size allowed.
  int min_partition_size_1d;

  int is_le_min_sq_part;
  int is_gt_max_sq_part;

  // Indicates edge blocks in frame.
  int has_rows;
  int has_cols;

  // Block size of current partition.
  BLOCK_SIZE bsize;

  // Chroma subsampling in x and y directions.
  int ss_x;
  int ss_y;
} PartitionBlkParams;

// Structure holding state variables for partition search.
typedef struct {
  // Parameters related to partition block size.
  PartitionBlkParams part_blk_params;

  int is_block_splittable;

  // // RD cost for the current block of given partition type.
  // RD_STATS this_rdc;

  // // RD cost summed across all blocks of partition type.
  // RD_STATS sum_rdc;

  // Array holding partition type cost.
  int tmp_partition_cost[PARTITION_TYPES];
#if CONFIG_EXT_RECUR_PARTITIONS
  int partition_cost_table[EXT_PARTITION_TYPES];
#endif

  // Pointer to partition cost buffer
  int *partition_cost;

  // RD costs for different partition types.
  int64_t none_rd;
  int64_t split_rd[SUB_PARTITIONS_SPLIT];
  // RD costs for rectangular partitions.
  // rect_part_rd[0][i] is the RD cost of ith partition index of PARTITION_HORZ.
  // rect_part_rd[1][i] is the RD cost of ith partition index of PARTITION_VERT.
  int64_t rect_part_rd[NUM_RECT_PARTS][SUB_PARTITIONS_RECT];

#if CONFIG_EXT_RECUR_PARTITIONS
  // New Simple Motion Result for PARTITION_NONE
  SMSPartitionStats none_data;
#endif  // CONFIG_EXT_RECUR_PARTITIONS

  // Flags indicating if the corresponding partition was winner or not.
  // Used to bypass similar blocks during AB partition evaluation.
  int split_ctx_is_ready[2];
  int rect_ctx_is_ready[NUM_RECT_PARTS];

  // Flags to prune/skip particular partition size evaluation.
  int terminate_partition_search;
  int partition_none_allowed;
  int partition_rect_allowed[NUM_RECT_PARTS];
#if CONFIG_EXT_RECUR_PARTITIONS
  int partition_3_allowed[NUM_RECT_PARTS];
#else
  int partition_ab_allowed[NUM_RECT_PARTS][NUM_AB_TYPES];
  int partition_4_allowed[NUM_RECT_PARTS];
#endif  // CONFIG_EXT_RECUR_PARTITIONS
  int do_rectangular_split;
#if !(CONFIG_EXT_RECUR_PARTITIONS && !KEEP_PARTITION_SPLIT)
  int do_square_split;
#endif  // !(CONFIG_EXT_RECUR_PARTITIONS && !KEEP_PARTITION_SPLIT)
  int prune_rect_part[NUM_RECT_PARTS];

  // Partition plane context index.
  int pl;

  // This flag will be set if best partition is found from the search.
  bool found_best_partition;
} PartitionSearchState;

void av1_init_partition_search_state(PartitionSearchState *search_state,
                                     MACROBLOCK *x, AV1_COMP *const cpi,
                                     const PC_TREE *pc_tree, int mi_row,
                                     int mi_col, BLOCK_SIZE bsize,
                                     BLOCK_SIZE max_sq_part,
                                     BLOCK_SIZE min_sq_part);

void av1_update_picked_ref_frames_mask(MACROBLOCK *const x, int ref_type,
                                       BLOCK_SIZE bsize, int mib_size,
                                       int mi_row, int mi_col);
#endif  // !CONFIG_REALTIME_ONLY

void av1_encode_b(const AV1_COMP *const cpi, TileDataEnc *tile_data,
                  ThreadData *td, TOKENEXTRA **tp, int mi_row, int mi_col,
                  RUN_TYPE dry_run, BLOCK_SIZE bsize, PARTITION_TYPE partition,
                  const PICK_MODE_CONTEXT *const ctx, int *rate);

#if !CONFIG_REALTIME_ONLY
void av1_encode_sb(const AV1_COMP *const cpi, ThreadData *td,
                   TileDataEnc *tile_data, TOKENEXTRA **tp, int mi_row,
                   int mi_col, RUN_TYPE dry_run, BLOCK_SIZE bsize,
                   PC_TREE *pc_tree, PARTITION_TREE *ptree, int *rate);

// Checks to see if a super block is on a horizontal image edge.
// In most cases this is the "real" edge unless there are formatting
// bars embedded in the stream.
int av1_active_h_edge(const AV1_COMP *cpi, int mi_row, int mi_step);

// Checks to see if a super block is on a vertical image edge.
// In most cases this is the "real" edge unless there are formatting
// bars embedded in the stream.
int av1_active_v_edge(const AV1_COMP *cpi, int mi_row, int mi_step);

void av1_encode_superblock(const AV1_COMP *const cpi, TileDataEnc *tile_data,
                           ThreadData *td, TOKENEXTRA **t, RUN_TYPE dry_run,
                           BLOCK_SIZE bsize, int *rate);

void av1_update_state(const AV1_COMP *const cpi, ThreadData *td,
                      const PICK_MODE_CONTEXT *const ctx, int mi_row,
                      int mi_col, BLOCK_SIZE bsize, RUN_TYPE dry_run);
#endif  // !CONFIG_REALTIME_ONLY

void av1_setup_block_rdmult(const struct AV1_COMP *cpi,
                            struct macroblock *const x, int mi_row, int mi_col,
                            BLOCK_SIZE bsize, AQ_MODE aq_mode,
                            MB_MODE_INFO *mbmi);

#if !CONFIG_REALTIME_ONLY
static INLINE int check_is_chroma_size_valid(PARTITION_TYPE partition,
                                             BLOCK_SIZE bsize, int mi_row,
                                             int mi_col, int ss_x, int ss_y,
                                             const PC_TREE *pc_tree) {
  const BLOCK_SIZE subsize = get_partition_subsize(bsize, partition);
  int is_valid = 0;
  if (subsize < BLOCK_SIZES_ALL) {
    CHROMA_REF_INFO tmp_chr_ref_info = { 1, 0, mi_row, mi_col, bsize, bsize };
    set_chroma_ref_info(mi_row, mi_col, 0, subsize, &tmp_chr_ref_info,
                        &pc_tree->chroma_ref_info, bsize, partition, ss_x,
                        ss_y);
    is_valid = get_plane_block_size(tmp_chr_ref_info.bsize_base, ss_x, ss_y) !=
               BLOCK_INVALID;
  }
  return is_valid;
}

// If any of the search_state_rd is INT_MAX, set them to 0 for the ML models
static INLINE void clip_partition_search_state_rd(
    PartitionSearchState *search_state) {
  search_state->rect_part_rd[HORZ][0] =
      (search_state->rect_part_rd[HORZ][0] < INT64_MAX
           ? search_state->rect_part_rd[HORZ][0]
           : 0);
  search_state->rect_part_rd[HORZ][1] =
      (search_state->rect_part_rd[HORZ][1] < INT64_MAX
           ? search_state->rect_part_rd[HORZ][1]
           : 0);
  search_state->rect_part_rd[VERT][0] =
      (search_state->rect_part_rd[VERT][0] < INT64_MAX
           ? search_state->rect_part_rd[VERT][0]
           : 0);
  search_state->rect_part_rd[VERT][1] =
      (search_state->rect_part_rd[VERT][1] < INT64_MAX
           ? search_state->rect_part_rd[VERT][1]
           : 0);
  search_state->split_rd[0] =
      (search_state->split_rd[0] < INT64_MAX ? search_state->split_rd[0] : 0);
  search_state->split_rd[1] =
      (search_state->split_rd[1] < INT64_MAX ? search_state->split_rd[1] : 0);
  search_state->split_rd[2] =
      (search_state->split_rd[2] < INT64_MAX ? search_state->split_rd[2] : 0);
  search_state->split_rd[3] =
      (search_state->split_rd[3] < INT64_MAX ? search_state->split_rd[3] : 0);
}

static INLINE void init_partition_allowed(PartitionSearchState *search_state,
                                          const AV1_COMP *cpi,
                                          const PC_TREE *pc_tree) {
  const PartitionBlkParams *blk_params = &search_state->part_blk_params;
  const BLOCK_SIZE bsize = blk_params->bsize;
  const int mi_row = blk_params->mi_row, mi_col = blk_params->mi_col;
  const int has_rows = blk_params->has_rows, has_cols = blk_params->has_cols;
  const int ss_x = blk_params->ss_x, ss_y = blk_params->ss_y;

  int is_chroma_size_valid_horz = check_is_chroma_size_valid(
      PARTITION_HORZ, bsize, mi_row, mi_col, ss_x, ss_y, pc_tree);

  int is_chroma_size_valid_vert = check_is_chroma_size_valid(
      PARTITION_VERT, bsize, mi_row, mi_col, ss_x, ss_y, pc_tree);

  search_state->terminate_partition_search = 0;
  search_state->partition_none_allowed = has_rows && has_cols;
  search_state->partition_rect_allowed[HORZ] =
#if CONFIG_EXT_RECUR_PARTITIONS && !KEEP_PARTITION_SPLIT
      (has_cols || (!has_rows && !has_cols)) &&
      is_partition_valid(bsize, PARTITION_HORZ) &&
#else   // CONFIG_EXT_RECUR_PARTITIONS && !KEEP_PARTITION_SPLIT
      blk_params->has_cols && is_partition_valid(bsize, PARTITION_HORZ) &&
#endif  // CONFIG_EXT_RECUR_PARTITIONS && !KEEP_PARTITION_SPLIT
      cpi->oxcf.enable_rect_partitions && is_chroma_size_valid_horz;
  search_state->partition_rect_allowed[VERT] =
#if CONFIG_EXT_RECUR_PARTITIONS && !KEEP_PARTITION_SPLIT
      (has_rows || (!has_rows && !has_cols)) &&
      is_partition_valid(bsize, PARTITION_VERT) &&
#else   // CONFIG_EXT_RECUR_PARTITIONS && !KEEP_PARTITION_SPLIT
      blk_params->has_rows && is_partition_valid(bsize, PARTITION_VERT) &&
#endif  // CONFIG_EXT_RECUR_PARTITIONS && !KEEP_PARTITION_SPLIT
      cpi->oxcf.enable_rect_partitions && is_chroma_size_valid_vert;
#if !(CONFIG_EXT_RECUR_PARTITIONS && !KEEP_PARTITION_SPLIT)
  search_state->do_square_split =
      search_state->is_block_splittable && is_square_block(bsize);
#endif  // !(CONFIG_EXT_RECUR_PARTITIONS && !KEEP_PARTITION_SPLIT)
}
#endif  // !CONFIG_REALTIME_ONLY
#ifdef __cplusplus
}  // extern "C"
#endif

#endif  // AOM_AV1_ENCODER_ENCODEFRAME_UTILS_H_