/*
 * Copyright (c) 2021, 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_AOM_AOM_EXTERNAL_PARTITION_H_
#define AOM_AOM_AOM_EXTERNAL_PARTITION_H_

/*!\defgroup aom_encoder AOMedia AOM/AV1 Encoder
 * \ingroup aom
 *
 * @{
 */
#include <stdint.h>

/*!\file
 * \brief Provides function pointer definitions for the external partition.
 */

/*!\brief Current ABI version number
 *
 * \internal
 * If this file is altered in any way that changes the ABI, this value
 * must be bumped. Examples include, but are not limited to, changing
 * types, removing or reassigning enums, adding/removing/rearranging
 * fields to structures.
 */
#define AOM_EXT_PART_ABI_VERSION 4

#ifdef __cplusplus
extern "C" {
#endif

/*!\brief Abstract external partition model handler
 */
typedef void *aom_ext_part_model_t;

/*!\brief Number of features to determine whether to skip partition none and
 * do partition split directly. The same as "FEATURE_SIZE_SMS_SPLIT".
 */
#define AOM_EXT_PART_SIZE_DIRECT_SPLIT 17

/*!\brief Number of features to use simple motion search to prune out
 * rectangular partition in some direction. The same as
 * "FEATURE_SIZE_SMS_PRUNE_PART".
 */
#define AOM_EXT_PART_SIZE_PRUNE_PART 25

/*!\brief Number of features to prune split and rectangular partition
 * after PARTITION_NONE.
 */
#define AOM_EXT_PART_SIZE_PRUNE_NONE 4

/*!\brief Number of features to terminates partition after partition none using
 * simple_motion_search features and the rate, distortion, and rdcost of
 * PARTITION_NONE. The same as "FEATURE_SIZE_SMS_TERM_NONE".
 */
#define AOM_EXT_PART_SIZE_TERM_NONE 28

/*!\brief Number of features to terminates partition after partition split.
 */
#define AOM_EXT_PART_SIZE_TERM_SPLIT 31

/*!\brief Number of features to prune rectangular partition using stats
 * collected after partition split.
 */
#define AOM_EXT_PART_SIZE_PRUNE_RECT 9

/*!\brief Number of features to prune AB partition using stats
 * collected after rectangular partition..
 */
#define AOM_EXT_PART_SIZE_PRUNE_AB 10

/*!\brief Number of features to prune 4-way partition using stats
 * collected after AB partition.
 */
#define AOM_EXT_PART_SIZE_PRUNE_4_WAY 18

/*!\brief Decision mode of the external partition model.
 * AOM_EXT_PART_WHOLE_TREE: the external partition model should provide the
 * whole partition tree for the superblock.
 *
 * AOM_EXT_PART_RECURSIVE: the external partition model provides the partition
 * decision of the current block only. The decision process starts from
 * the superblock size, down to the smallest block size (4x4) recursively.
 */
typedef enum aom_ext_part_decision_mode {
  AOM_EXT_PART_WHOLE_TREE = 0,
  AOM_EXT_PART_RECURSIVE = 1,
} aom_ext_part_decision_mode_t;

/*!\brief Config information sent to the external partition model.
 *
 * For example, the maximum superblock size determined by the sequence header.
 */
typedef struct aom_ext_part_config {
  int superblock_size;  ///< super block size (either 64x64 or 128x128)
} aom_ext_part_config_t;

/*!\brief Features pass to the external model to make partition decisions.
 * Specifically, features collected before NONE partition.
 * Features "f" are used to determine:
 * partition_none_allowed, partition_horz_allowed, partition_vert_allowed,
 * do_rectangular_split, do_square_split
 * Features "f_part2" are used to determine:
 * prune_horz, prune_vert.
 */
typedef struct aom_partition_features_before_none {
  /*! features to determine whether skip partition none and do split directly */
  float f[AOM_EXT_PART_SIZE_DIRECT_SPLIT];
  /*! features to determine whether to prune rectangular partition */
  float f_part2[AOM_EXT_PART_SIZE_PRUNE_PART];
} aom_partition_features_before_none_t;

/*!\brief Features pass to the external model to make partition decisions.
 * Specifically, features collected after NONE partition.
 */
typedef struct aom_partition_features_none {
  /*! features to prune split and rectangular partition */
  float f[AOM_EXT_PART_SIZE_PRUNE_NONE];
  /*! features to determine termination of partition */
  float f_terminate[AOM_EXT_PART_SIZE_TERM_NONE];
} aom_partition_features_none_t;

/*!\brief Features pass to the external model to make partition decisions.
 * Specifically, features collected after SPLIT partition.
 */
typedef struct aom_partition_features_split {
  /*! features to determine termination of  partition */
  float f_terminate[AOM_EXT_PART_SIZE_TERM_SPLIT];
  /*! features to determine pruning rect partition */
  float f_prune_rect[AOM_EXT_PART_SIZE_PRUNE_RECT];
} aom_partition_features_split_t;

/*!\brief Features pass to the external model to make partition decisions.
 * Specifically, features collected after RECTANGULAR partition.
 */
typedef struct aom_partition_features_rect {
  /*! features to determine pruning AB partition */
  float f[AOM_EXT_PART_SIZE_PRUNE_AB];
} aom_partition_features_rect_t;

/*!\brief Features pass to the external model to make partition decisions.
 * Specifically, features collected after AB partition: HORZ_A, HORZ_B, VERT_A,
 * VERT_B.
 */
typedef struct aom_partition_features_ab {
  /*! features to determine pruning 4-way partition */
  float f[AOM_EXT_PART_SIZE_PRUNE_4_WAY];
} aom_partition_features_ab_t;

/*!\brief Feature id to tell the external model the current stage in partition
 * pruning and what features to use to make decisions accordingly.
 */
typedef enum {
  AOM_EXT_PART_FEATURE_BEFORE_NONE,
  AOM_EXT_PART_FEATURE_BEFORE_NONE_PART2,
  AOM_EXT_PART_FEATURE_AFTER_NONE,
  AOM_EXT_PART_FEATURE_AFTER_NONE_PART2,
  AOM_EXT_PART_FEATURE_AFTER_SPLIT,
  AOM_EXT_PART_FEATURE_AFTER_SPLIT_PART2,
  AOM_EXT_PART_FEATURE_AFTER_RECT,
  AOM_EXT_PART_FEATURE_AFTER_AB
} AOM_EXT_PART_FEATURE_ID;

/*!\brief Features collected from the tpl process.
 *
 * The tpl process collects information that help measure the inter-frame
 * dependency.
 * The tpl process is computed in the unit of tpl_bsize_1d (16x16).
 * Therefore, the max number of units inside a superblock is
 * 128x128 / (16x16) = 64. Change it if the tpl process changes.
 */
typedef struct aom_sb_tpl_features {
  int available;        ///< If tpl stats are available
  int tpl_unit_length;  ///< The block length of tpl process
  int num_units;        ///< The number of units inside the current superblock
  int64_t intra_cost[64];   ///< The intra cost of each unit
  int64_t inter_cost[64];   ///< The inter cost of each unit
  int64_t mc_dep_cost[64];  ///< The motion compensated dependency cost
} aom_sb_tpl_features_t;

/*!\brief Features collected from the simple motion process.
 *
 * The simple motion process collects information by applying motion compensated
 * prediction on each block.
 * The block size is 16x16, which could be changed. If it is changed, update
 * comments and the array size here.
 */
typedef struct aom_sb_simple_motion_features {
  int unit_length;    ///< The block length of the simple motion process
  int num_units;      ///< The number of units inside the current superblock
  int block_sse[64];  ///< Sum of squared error of each unit
  int block_var[64];  ///< Variance of each unit
} aom_sb_simple_motion_features_t;

/*!\brief Features of each super block.
 *
 * Features collected for each super block before partition search.
 */
typedef struct aom_sb_features {
  /*! Features from motion search */
  aom_sb_simple_motion_features_t motion_features;
  /*! Features from tpl process */
  aom_sb_tpl_features_t tpl_features;
} aom_sb_features_t;

/*!\brief Features pass to the external model to make partition decisions.
 *
 * The encoder sends these features to the external model through
 * "func()" defined in .....
 *
 * NOTE: new member variables may be added to this structure in the future.
 * Once new features are finalized, bump the major version of libaom.
 */
typedef struct aom_partition_features {
  // Features for the current supervised multi-stage ML model.
  /*! Feature ID to indicate active features */
  AOM_EXT_PART_FEATURE_ID id;
  /*! Features collected before NONE partition */
  aom_partition_features_before_none_t before_part_none;
  /*! Features collected after NONE partition */
  aom_partition_features_none_t after_part_none;
  /*! Features collected after SPLIT partition */
  aom_partition_features_split_t after_part_split;
  /*! Features collected after RECTANGULAR partition */
  aom_partition_features_rect_t after_part_rect;
  /*! Features collected after AB partition */
  aom_partition_features_ab_t after_part_ab;

  // Features for a new ML model.
  aom_sb_features_t sb_features;  ///< Features collected for the super block
  int mi_row;                     ///< Mi_row position of the block
  int mi_col;                     ///< Mi_col position of the block
  int frame_width;                ///< Frame width
  int frame_height;               ///< Frame height
  int block_size;                 ///< As "BLOCK_SIZE" in av1/common/enums.h
  /*!
   * Valid partition types. A bitmask is used.  "1" represents the
   * corresponding type is vaild. The bitmask follows the enum order for
   * PARTITION_TYPE in "enums.h" to represent one partition type at a bit.
   * For example, 0x01 stands for only PARTITION_NONE is valid,
   * 0x09 (00...001001) stands for PARTITION_NONE and PARTITION_SPLIT are valid.
   */
  int valid_partition_types;
} aom_partition_features_t;

/*!\brief Partition decisions received from the external model.
 *
 * The encoder receives partition decisions and encodes the superblock
 * with the given partition type.
 * The encoder receives it from "func()" define in ....
 *
 * NOTE: new member variables may be added to this structure in the future.
 * Once new features are finalized, bump the major version of libaom.
 */
typedef struct aom_partition_decision {
  // Decisions for directly set partition types
  int is_final_decision;         ///< The flag whether it's the final decision
  int num_nodes;                 ///< The number of leaf nodes
  int partition_decision[2048];  ///< Partition decisions
  int current_decision;          ///< Partition decision for the current block

  // Decisions for partition type pruning
  int terminate_partition_search;  ///< Terminate further partition search
  int partition_none_allowed;      ///< Allow partition none type
  int partition_rect_allowed[2];   ///< Allow rectangular partitions
  int do_rectangular_split;        ///< Try rectangular split partition
  int do_square_split;             ///< Try square split partition
  int prune_rect_part[2];          ///< Prune rectangular partition
  int horza_partition_allowed;     ///< Allow HORZ_A partitioin
  int horzb_partition_allowed;     ///< Allow HORZ_B partitioin
  int verta_partition_allowed;     ///< Allow VERT_A partitioin
  int vertb_partition_allowed;     ///< Allow VERT_B partitioin
  int partition_horz4_allowed;     ///< Allow HORZ4 partition
  int partition_vert4_allowed;     ///< Allow VERT4 partition
} aom_partition_decision_t;

/*!\brief Encoding stats for the given partition decision.
 *
 * The encoding stats collected by encoding the superblock with the
 * given partition types.
 * The encoder sends the stats to the external model for training
 * or inference though "func()" defined in ....
 */
typedef struct aom_partition_stats {
  int rate;        ///< Rate cost of the block
  int64_t dist;    ///< Distortion of the block
  int64_t rdcost;  ///< Rate-distortion cost of the block
} aom_partition_stats_t;

/*!\brief Enum for return status.
 */
typedef enum aom_ext_part_status {
  AOM_EXT_PART_OK = 0,     ///< Status of success
  AOM_EXT_PART_ERROR = 1,  ///< Status of failure
  AOM_EXT_PART_TEST = 2,   ///< Status used for tests
} aom_ext_part_status_t;

/*!\brief Callback of creating an external partition model.
 *
 * The callback is invoked by the encoder to create an external partition
 * model.
 *
 * \param[in] priv Callback's private data
 * \param[in] part_config Config information pointer for model creation
 * \param[out] ext_part_model Pointer to the model
 */
typedef aom_ext_part_status_t (*aom_ext_part_create_model_fn_t)(
    void *priv, const aom_ext_part_config_t *part_config,
    aom_ext_part_model_t *ext_part_model);

/*!\brief Callback of sending features to the external partition model.
 *
 * The callback is invoked by the encoder to send features to the external
 * partition model.
 *
 * \param[in] ext_part_model The external model
 * \param[in] part_features Pointer to the features
 */
typedef aom_ext_part_status_t (*aom_ext_part_send_features_fn_t)(
    aom_ext_part_model_t ext_part_model,
    const aom_partition_features_t *part_features);

/*!\brief Callback of receiving partition decisions from the external
 * partition model.
 *
 * The callback is invoked by the encoder to receive partition decisions from
 * the external partition model.
 *
 * \param[in] ext_part_model The external model
 * \param[in] ext_part_decision Pointer to the partition decisions
 */
typedef aom_ext_part_status_t (*aom_ext_part_get_decision_fn_t)(
    aom_ext_part_model_t ext_part_model,
    aom_partition_decision_t *ext_part_decision);

/*!\brief Callback of sending stats to the external partition model.
 *
 * The callback is invoked by the encoder to send encoding stats to
 * the external partition model.
 *
 * \param[in] ext_part_model The external model
 * \param[in] ext_part_stats Pointer to the encoding stats
 */
typedef aom_ext_part_status_t (*aom_ext_part_send_partition_stats_fn_t)(
    aom_ext_part_model_t ext_part_model,
    const aom_partition_stats_t *ext_part_stats);

/*!\brief Callback of deleting the external partition model.
 *
 * The callback is invoked by the encoder to delete the external partition
 * model.
 *
 * \param[in] ext_part_model The external model
 */
typedef aom_ext_part_status_t (*aom_ext_part_delete_model_fn_t)(
    aom_ext_part_model_t ext_part_model);

/*!\brief Callback function set for external partition model.
 *
 * Uses can enable external partition model by registering a set of
 * callback functions with the flag: AV1E_SET_EXTERNAL_PARTITION_MODEL
 */
typedef struct aom_ext_part_funcs {
  /*!
   * Create an external partition model.
   */
  aom_ext_part_create_model_fn_t create_model;

  /*!
   * Send features to the external partition model to make partition decisions.
   */
  aom_ext_part_send_features_fn_t send_features;

  /*!
   * Get partition decisions from the external partition model.
   */
  aom_ext_part_get_decision_fn_t get_partition_decision;

  /*!
   * Send stats of the current partition to the external model.
   */
  aom_ext_part_send_partition_stats_fn_t send_partition_stats;

  /*!
   * Delete the external partition model.
   */
  aom_ext_part_delete_model_fn_t delete_model;

  /*!
   * The decision mode of the model.
   */
  aom_ext_part_decision_mode_t decision_mode;

  /*!
   * Private data for the external partition model.
   */
  void *priv;
} aom_ext_part_funcs_t;

/*!@} - end defgroup aom_encoder*/
#ifdef __cplusplus
}  // extern "C"
#endif

#endif  // AOM_AOM_AOM_EXTERNAL_PARTITION_H_