/*
 * Copyright (c) 2021, Alliance for Open Media. All rights reserved
 *
 * This source code is subject to the terms of the BSD 3-Clause Clear License
 * and the Alliance for Open Media Patent License 1.0. If the BSD 3-Clause Clear
 * License was not distributed with this source code in the LICENSE file, you
 * can obtain it at aomedia.org/license/software-license/bsd-3-c-c/.  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
 * aomedia.org/license/patent-license/.
 */

#ifndef AOM_AV1_ENCODER_TRANSFORM_SEARCH_H_
#define AOM_AV1_ENCODER_TRANSFORM_SEARCH_H_

#include "av1/common/pred_common.h"
#include "av1/encoder/encoder.h"

#ifdef __cplusplus
extern "C" {
#endif

// Set this macro as 1 to collect data about tx size selection.
#define COLLECT_TX_SIZE_DATA 0

#if COLLECT_TX_SIZE_DATA
static const char av1_tx_size_data_output_file[] = "tx_size_data.txt";
#endif

enum {
  FTXS_NONE = 0,
  FTXS_DCT_AND_1D_DCT_ONLY = 1 << 0,
  FTXS_DISABLE_TRELLIS_OPT = 1 << 1,
  FTXS_USE_TRANSFORM_DOMAIN = 1 << 2
} UENUM1BYTE(FAST_TX_SEARCH_MODE);

#if CONFIG_NEW_TX_PARTITION
#if CONFIG_TX_PARTITION_CTX
static AOM_INLINE int inter_tx_partition_cost(const MACROBLOCK *const x,
                                              TX_PARTITION_TYPE partition,
                                              BLOCK_SIZE bsize,
                                              TX_SIZE max_tx_size) {
#else
static AOM_INLINE int inter_tx_partition_cost(
    const MACROBLOCK *const x, int is_rect, TX_PARTITION_TYPE partition,
    const TXFM_CONTEXT *const above_ctx, const TXFM_CONTEXT *const left_ctx,
    BLOCK_SIZE bsize, TX_SIZE max_tx_size) {
#endif  // CONFIG_TX_PARTITION_CTX
  int cost = 0;
  const int allow_horz = allow_tx_horz_split(max_tx_size);
  const int allow_vert = allow_tx_vert_split(max_tx_size);
#if CONFIG_IMPROVEIDTX_CTXS
  const MACROBLOCKD *const xd = &x->e_mbd;
  const MB_MODE_INFO *const mbmi = xd->mi[0];
  const int is_fsc = (mbmi->fsc_mode[xd->tree_type == CHROMA_PART]);
#endif  // CONFIG_IMPROVEIDTX_CTXS
#if CONFIG_TX_PARTITION_CTX
#if CONFIG_TX_PARTITION_TYPE_EXT
  const int bsize_group = size_to_tx_part_group_lookup[bsize];
  const int txsize_group = size_to_tx_type_group_lookup[bsize];
  int do_partition = 0;
  if (allow_horz || allow_vert) {
    do_partition = (partition != TX_PARTITION_NONE);
#if CONFIG_IMPROVEIDTX_CTXS
    cost += x->mode_costs
                .txfm_do_partition_cost[is_fsc][1][bsize_group][do_partition];
#else
    cost += x->mode_costs.txfm_do_partition_cost[1][bsize_group][do_partition];
#endif  // CONFIG_IMPROVEIDTX_CTXS
  }

  if (do_partition) {
    if (allow_horz && allow_vert) {
      assert(txsize_group > 0);
      const TX_PARTITION_TYPE split4_partition =
          get_split4_partition(partition);
#if CONFIG_IMPROVEIDTX_CTXS
      cost += x->mode_costs
                  .txfm_4way_partition_type_cost[is_fsc][1][txsize_group - 1]
                                                [split4_partition - 1];
#else
      cost += x->mode_costs.txfm_4way_partition_type_cost[1][txsize_group - 1]
                                                         [split4_partition - 1];
#endif  // CONFIG_IMPROVEIDTX_CTXS
    } else if (allow_horz || allow_vert) {
      int has_first_split = 0;
      if (partition == TX_PARTITION_VERT_M || partition == TX_PARTITION_HORZ_M)
        has_first_split = 1;

      if (txsize_group) {
#if CONFIG_IMPROVEIDTX_CTXS
        cost += x->mode_costs
                    .txfm_4way_partition_type_cost[is_fsc][1][txsize_group - 1]
                                                  [has_first_split];
#else
        cost += x->mode_costs.txfm_4way_partition_type_cost[1][txsize_group - 1]
                                                           [has_first_split];
#endif  // CONFIG_IMPROVEIDTX_CTXS
      }
    }
  }
#else
  const int bsize_group = size_to_tx_part_group_lookup[bsize];
  int do_partition = 0;
  if (allow_horz || allow_vert) {
    do_partition = (partition != TX_PARTITION_NONE);
#if CONFIG_IMPROVEIDTX_CTXS
    cost += x->mode_costs
                .txfm_do_partition_cost[is_fsc][1][bsize_group][do_partition];
#else
    cost += x->mode_costs.txfm_do_partition_cost[1][bsize_group][do_partition];
#endif  // CONFIG_IMPROVEIDTX_CTXS
  }

  if (do_partition) {
    if (allow_horz && allow_vert) {
      assert(bsize_group > 0);
      const TX_PARTITION_TYPE split4_partition =
          get_split4_partition(partition);
#if CONFIG_IMPROVEIDTX_CTXS
      cost += x->mode_costs
                  .txfm_4way_partition_type_cost[is_fsc][1][bsize_group - 1]
                                                [split4_partition - 1];
#else
      cost += x->mode_costs.txfm_4way_partition_type_cost[1][bsize_group - 1]
                                                         [split4_partition - 1];
#endif  // CONFIG_IMPROVEIDTX_CTXS
    }
  }
#endif  // CONFIG_TX_PARTITION_TYPE_EXT
#else
  if (allow_horz && allow_vert) {
    const int split4_ctx_0 = txfm_partition_split4_inter_context(
        above_ctx, left_ctx, bsize, max_tx_size);
    const TX_PARTITION_TYPE split4_partition = get_split4_partition(partition);
    cost += x->mode_costs.inter_4way_txfm_partition_cost[is_rect][split4_ctx_0]
                                                        [split4_partition];
  } else if (allow_horz || allow_vert) {
    const int has_first_split = partition != TX_PARTITION_NONE;
    cost += x->mode_costs.inter_2way_txfm_partition_cost[has_first_split];
  } else {
    assert(!allow_horz && !allow_vert);
    assert(partition == PARTITION_NONE);
  }
#endif  // CONFIG_TX_PARTITION_CTX
  return cost;
}

static AOM_INLINE int intra_tx_partition_cost(const MACROBLOCK *const x,
#if !CONFIG_TX_PARTITION_CTX
                                              int is_rect,
#endif  // !CONFIG_TX_PARTITION_CTX
                                              TX_PARTITION_TYPE partition,
                                              TX_SIZE max_tx_size) {
  int cost = 0;
  const MACROBLOCKD *const xd = &x->e_mbd;
  const int allow_horz = allow_tx_horz_split(max_tx_size);
  const int allow_vert = allow_tx_vert_split(max_tx_size);
#if CONFIG_TX_PARTITION_CTX
#if CONFIG_TX_PARTITION_TYPE_EXT
  const MB_MODE_INFO *const mbmi = xd->mi[0];
  const BLOCK_SIZE bsize = mbmi->sb_type[PLANE_TYPE_Y];
  const int bsize_group = size_to_tx_part_group_lookup[bsize];
  const int txsize_group = size_to_tx_type_group_lookup[bsize];
#if CONFIG_IMPROVEIDTX_CTXS
  const int is_fsc = (mbmi->fsc_mode[xd->tree_type == CHROMA_PART]);
#endif  // CONFIG_IMPROVEIDTX_CTXS
  int do_partition = 0;
  if (allow_horz || allow_vert) {
    do_partition = (partition != TX_PARTITION_NONE);
#if CONFIG_IMPROVEIDTX_CTXS
    cost += x->mode_costs
                .txfm_do_partition_cost[is_fsc][0][bsize_group][do_partition];
#else
    cost += x->mode_costs.txfm_do_partition_cost[0][bsize_group][do_partition];
#endif  // CONFIG_IMPROVEIDTX_CTXS
  }

  if (do_partition) {
    if (allow_horz && allow_vert) {
      assert(txsize_group > 0);
      const TX_PARTITION_TYPE split4_partition =
          get_split4_partition(partition);
#if CONFIG_IMPROVEIDTX_CTXS
      cost += x->mode_costs
                  .txfm_4way_partition_type_cost[is_fsc][0][txsize_group - 1]
                                                [split4_partition - 1];
#else
      cost += x->mode_costs.txfm_4way_partition_type_cost[0][txsize_group - 1]
                                                         [split4_partition - 1];
#endif  // CONFIG_IMPROVEIDTX_CTXS
    } else if (allow_horz || allow_vert) {
      int has_first_split = 0;
      if (partition == TX_PARTITION_VERT_M || partition == TX_PARTITION_HORZ_M)
        has_first_split = 1;

      if (txsize_group) {
#if CONFIG_IMPROVEIDTX_CTXS
        cost += x->mode_costs
                    .txfm_4way_partition_type_cost[is_fsc][0][txsize_group - 1]
                                                  [has_first_split];
#else
        cost += x->mode_costs.txfm_4way_partition_type_cost[0][txsize_group - 1]
                                                           [has_first_split];
#endif  // CONFIG_IMPROVEIDTX_CTXS
      }
    }
  }
#else
  const MB_MODE_INFO *const mbmi = xd->mi[0];
  const BLOCK_SIZE bsize = mbmi->sb_type[PLANE_TYPE_Y];
  const int bsize_group = size_to_tx_part_group_lookup[bsize];
#if CONFIG_IMPROVEIDTX_CTXS
  const int is_fsc = (mbmi->fsc_mode[xd->tree_type == CHROMA_PART]);
#endif  // CONFIG_IMPROVEIDTX_CTXS
  int do_partition = 0;
  if (allow_horz || allow_vert) {
    do_partition = (partition != TX_PARTITION_NONE);
#if CONFIG_IMPROVEIDTX_CTXS
    cost += x->mode_costs
                .txfm_do_partition_cost[is_fsc][0][bsize_group][do_partition];
#else
    cost += x->mode_costs.txfm_do_partition_cost[0][bsize_group][do_partition];
#endif  // CONFIG_IMPROVEIDTX_CTXS
  }

  if (do_partition) {
    if (allow_horz && allow_vert) {
      assert(bsize_group > 0);
      const TX_PARTITION_TYPE split4_partition =
          get_split4_partition(partition);
#if CONFIG_IMPROVEIDTX_CTXS
      cost += x->mode_costs
                  .txfm_4way_partition_type_cost[is_fsc][0][bsize_group - 1]
#else
      cost += x->mode_costs.txfm_4way_partition_type_cost[0][bsize_group - 1]
#endif  // CONFIG_IMPROVEIDTX_CTXS
                                                [split4_partition - 1];
    }
  }
#endif  // CONFIG_TX_PARTITION_TYPE_EXT
#else
  if (allow_horz && allow_vert) {
    const int split4_ctx_0 = get_tx_size_context(xd);
    const TX_PARTITION_TYPE split4_partition = get_split4_partition(partition);
    cost += x->mode_costs.intra_4way_txfm_partition_cost[is_rect][split4_ctx_0]
                                                        [split4_partition];
  } else if (allow_horz || allow_vert) {
    const int has_first_split = partition != TX_PARTITION_NONE;
    cost += x->mode_costs.intra_2way_txfm_partition_cost[has_first_split];
  } else {
    assert(!allow_horz && !allow_vert);
    assert(partition == PARTITION_NONE);
  }
#endif  // CONFIG_TX_PARTITION_CTX
  return cost;
}
#endif  // CONFIG_NEW_TX_PARTITION

// Calculate the energy of 32-bit signed data in 1d buffer
// TODO(any): Implement x86 SIMD for the following function
static AOM_INLINE uint64_t sum_squares_i32(const int32_t *src, int32_t n) {
  uint64_t ss = 0;
  do {
    const int32_t v = *src++;
    ss += ((int64_t)v) * v;
  } while (--n);

  return ss;
}

static AOM_INLINE int tx_size_cost(const MACROBLOCK *const x, BLOCK_SIZE bsize,
                                   TX_SIZE tx_size) {
  assert(bsize == x->e_mbd.mi[0]->sb_type[PLANE_TYPE_Y]);
  if (x->txfm_search_params.tx_mode_search_type != TX_MODE_SELECT ||
      !block_signals_txsize(bsize))
    return 0;

  const MACROBLOCKD *const xd = &x->e_mbd;
#if CONFIG_NEW_TX_PARTITION
  if (bsize >= BLOCK_SIZES_ALL) return INT_MAX;

  (void)tx_size;
  MB_MODE_INFO *const mbmi = xd->mi[0];
  const TX_SIZE max_tx_size = max_txsize_rect_lookup[bsize];
#if CONFIG_TX_PARTITION_CTX
  return intra_tx_partition_cost(x, mbmi->tx_partition_type[0], max_tx_size);
#else
  const int is_rect = is_rect_tx(max_tx_size);
  return intra_tx_partition_cost(x, is_rect, mbmi->tx_partition_type[0],
                                 max_tx_size);
#endif  // CONFIG_TX_PARTITION_CTX
#else
  const int32_t tx_size_cat = bsize_to_tx_size_cat(bsize);
  const int depth = tx_size_to_depth(tx_size, bsize);
  const int tx_size_ctx = get_tx_size_context(xd);
  return x->mode_costs.tx_size_cost[tx_size_cat][tx_size_ctx][depth];
#endif  // CONFIG_NEW_TX_PARTITION
}

static AOM_INLINE int skip_cctx_eval_based_on_eob(int plane, int is_inter,
                                                  uint16_t eob_c1,
                                                  CctxType cctx_type) {
  if (plane != AOM_PLANE_U) return 0;
  if (eob_c1 == 0) return 1;
  if (eob_c1 == 1 && !is_inter && cctx_type != CCTX_NONE) return 1;
  return 0;
}

/*!\brief Transform type search for luma macroblock with fixed transform size.
 *
 * \ingroup transform_search
 * Search for the best transform type and return the transform coefficients RD
 * cost of current luma macroblock with the given uniform transform size.
 *
 * \param[in]    x              Pointer to structure holding the data for the
                                current encoding macroblock
 * \param[in]    cpi            Top-level encoder structure
 * \param[in]    rd_stats       Pointer to struct to keep track of the RD stats
 * \param[in]    ref_best_rd    Best RD cost seen for this block so far
 * \param[in]    bs             Size of the current macroblock
 * \param[in]    tx_size        The given transform size
 * \param[in]    ftxs_mode      Transform search mode specifying desired speed
                                and quality tradeoff
 * \param[in]    skip_trellis   Binary flag indicating if trellis optimization
                                should be skipped
 * \return       An int64_t value that is the best RD cost found.
 */
int64_t av1_uniform_txfm_yrd(const AV1_COMP *const cpi, MACROBLOCK *x,
                             RD_STATS *rd_stats, int64_t ref_best_rd,
                             BLOCK_SIZE bs, TX_SIZE tx_size,
                             FAST_TX_SEARCH_MODE ftxs_mode, int skip_trellis);

/*!\brief Recursive transform size and type search.
 *
 * \ingroup transform_search
 * Search for best transform size and type for luma inter blocks. The transform
 * block partitioning can be recursive resulting in non-uniform transform sizes.
 * The best transform size and type, if found, will be saved in the MB_MODE_INFO
 * structure, and the corresponding RD stats will be saved in rd_stats.
 *
 * \param[in]    cpi            Top-level encoder structure
 * \param[in]    x              Pointer to structure holding the data for the
                                current encoding macroblock
 * \param[in]    rd_stats       Pointer to struct to keep track of the RD stats
 * \param[in]    bsize          Current macroblock size
 * \param[in]    ref_best_rd    Best RD cost seen for this block so far
 * Nothing is returned. The selected transform size and type will be saved
 * in the MB_MODE_INFO structure.
 */
void av1_pick_recursive_tx_size_type_yrd(const AV1_COMP *cpi, MACROBLOCK *x,
                                         RD_STATS *rd_stats, BLOCK_SIZE bsize,
                                         int64_t ref_best_rd);

/*!\brief Uniform transform size and type search.
 *
 * \ingroup transform_search
 * Search for the best transform size and type for current macroblock block,
 * with the assumption that all the transform blocks have a uniform size
 * (VP9 style). The selected transform size and type will be saved in the
 * MB_MODE_INFO structure; the corresponding RD stats will be saved in rd_stats.
 * This function may be used for both intra and inter predicted blocks.
 *
 * \param[in]    cpi            Top-level encoder structure
 * \param[in]    x              Pointer to structure holding the data for the
                                current encoding macroblock
 * \param[in]    rd_stats       Pointer to struct to keep track of the RD stats
 * \param[in]    bs             Current macroblock size
 * \param[in]    ref_best_rd    Best RD cost seen for this block so far
 * Nothing is returned. The selected transform size and type will be saved
 * in the MB_MODE_INFO structure.
 */
void av1_pick_uniform_tx_size_type_yrd(const AV1_COMP *const cpi, MACROBLOCK *x,
                                       RD_STATS *rd_stats, BLOCK_SIZE bs,
                                       int64_t ref_best_rd);

/*!\brief Chroma block transform search.
 *
 * \ingroup transform_search
 * Calculate the transform coefficient RD cost for the given chroma macroblock
 * If the current mode is intra, then this function will compute the predictor.
 *
 * \param[in]    cpi            Top-level encoder structure
 * \param[in]    x              Pointer to structure holding the data for the
                                current encoding macroblock
 * \param[in]    rd_stats       Pointer to struct to keep track of the RD stats
 * \param[in]    ref_best_rd    Best RD cost seen for this block so far
 * \return       An integer value is returned. 0: early termination triggered,
                 no valid rd cost available; 1: rd cost values are valid.
 */
int av1_txfm_uvrd(const AV1_COMP *const cpi, MACROBLOCK *x, RD_STATS *rd_stats,
                  int64_t ref_best_rd);

/*!\brief Transform type search with fixed transform size.
 *
 * \ingroup transform_search
 * Search for the best transform type and calculate the transform coefficients
 * RD cost of the current transform block with the specified (uniform) transform
 * size and plane. The RD results will be saved in rd_stats.
 *
 * \param[in]    x              Pointer to structure holding the data for the
                                current encoding macroblock
 * \param[in]    cpi            Top-level encoder structure
 * \param[in]    rd_stats       Pointer to struct to keep track of the RD stats
 * \param[in]    ref_best_rd    Best RD cost seen for this block so far
 * \param[in]    current_rd     Current RD cost for this block so far
 * \param[in]    plane          Plane index
 * \param[in]    plane_bsize    Size of the current macroblock considering
                                sup-sampling
 * \param[in]    tx_size        The given transform size
 * \param[in]    ftxs_mode      Transform search mode specifying desired speed
                                and quality tradeoff
 * \param[in]    skip_trellis   Binary flag indicating if trellis optimization
                                should be skipped
 *
 * Nothing is returned. The RD results will be saved in rd_stats.
 */
void av1_txfm_rd_in_plane(MACROBLOCK *x, const AV1_COMP *cpi,
                          RD_STATS *rd_stats, int64_t ref_best_rd,
                          int64_t current_rd, int plane, BLOCK_SIZE plane_bsize,
                          TX_SIZE tx_size, FAST_TX_SEARCH_MODE ftxs_mode,
                          int skip_trellis);

/*!\brief Recursive transform size and type search.
 *
 * \ingroup transform_search
 * This function combines y and uv planes' transform search processes together
 * for inter-predicted blocks (including IntraBC), when the prediction is
 * already generated. It first does subtraction to obtain the prediction error.
 * Then it calls
 * av1_pick_recursive_tx_size_type_yrd/av1_pick_uniform_tx_size_type_yrd and
 * av1_txfm_uvrd sequentially and handles possible early terminations.
 * The RD metrics are calculated and stored in rd_stats/_y/_uv.
 *
 * \param[in]    cpi            Top-level encoder structure
 * \param[in]    x              Pointer to structure holding the data for the
                                current encoding macroblock
 * \param[in]    bsize          Current macroblock size
 * \param[in]    rd_stats       Pointer to struct to keep track of the overal RD
                                stats
 * \param[in]    rd_stats_y     Pointer to struct to keep track of the RD
                                stats for the luma plane
 * \param[in]    rd_stats_uv    Pointer to struct to keep track of the RD
                                stats for the chroma planes
 * \param[in]    mode_rate      Rate cost to encode the prediction mode info. of
                                the current macroblock
 * \param[in]    ref_best_rd    Best RD cost seen for this block so far
 *
 * \return       An integer value is returned indicating if a valid transform
                 candidate is found (1) or not (0).
 */
int av1_txfm_search(const AV1_COMP *cpi, MACROBLOCK *x, BLOCK_SIZE bsize,
                    RD_STATS *rd_stats, RD_STATS *rd_stats_y,
                    RD_STATS *rd_stats_uv, int mode_rate, int64_t ref_best_rd);

#ifdef __cplusplus
}  // extern "C"
#endif

#endif  // AOM_AV1_ENCODER_TRANSFORM_SEARCH_H_