/*
 * 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/.
 */

#include "avm_ports/avm_timer.h"
#include "av2/encoder/encodeframe_utils.h"
#include "av2/encoder/partition_ml.h"

// Computes residual stats on a transformed and quantized residual of the
// block. This is used as ML features for prediction. The information computed
// is NNZ (Number of Non-Zero coefficients of the transformed and quantized
// residual), MAX_COEFF, PSNR.
void compute_residual_stats(AV2_COMP *const cpi, ThreadData *td, MACROBLOCK *x,
                            BLOCK_SIZE bsize, ResidualStats *out) {
  AV2_COMMON *cm = &cpi->common;
  MACROBLOCKD *const xd = &x->e_mbd;
  TX_SIZE tx_size = max_txsize_rect_lookup[bsize];
  const int plane = AVM_PLANE_Y;
  const int block = 0;
  struct macroblock_plane *const p = &x->plane[plane];
  struct macroblockd_plane *const pd = &xd->plane[plane];

  int old_tree_type = xd->tree_type;
  xd->tree_type = LUMA_PART;

  memset(out, 0, sizeof(ResidualStats));

  av2_subtract_plane(x, bsize, plane, cm->width, cm->height);

  const uint16_t *src = x->plane[0].src.buf;
  const uint16_t *dst = xd->plane[0].dst.buf;
  const int src_stride = x->plane[0].src.stride;
  const int dst_stride = xd->plane[0].dst.stride;
  out->var = cpi->fn_ptr[bsize].vf(src, src_stride, dst, dst_stride, &out->sse);

  const int num_blk = mi_size_wide[bsize] * mi_size_high[bsize];
  struct avm_internal_error_info error;
  AVM_CHECK_MEM_ERROR(&error, p->eobs,
                      avm_memalign(32, num_blk * sizeof(p->eobs[0])));
  p->coeff = td->shared_coeff_buf.coeff_buf[plane];
  p->qcoeff = td->shared_coeff_buf.qcoeff_buf[plane];
  p->dqcoeff = td->shared_coeff_buf.dqcoeff_buf[plane];
  tran_low_t *const dqcoeff = p->dqcoeff + BLOCK_OFFSET(block);
  tran_low_t *const qcoeff = p->qcoeff + BLOCK_OFFSET(block);
  tran_low_t *const coeff = p->coeff + BLOCK_OFFSET(block);
  AVM_CHECK_MEM_ERROR(&error, p->bobs,
                      avm_memalign(32, num_blk * sizeof(p->bobs[0])));
  AVM_CHECK_MEM_ERROR(
      &error, p->txb_entropy_ctx,
      avm_memalign(32, num_blk * sizeof(p->txb_entropy_ctx[0])));

  TxfmParam txfm_param;
  QUANT_PARAM quant_param;

  av2_setup_xform(cm, x, plane, tx_size, DCT_DCT, CCTX_NONE, &txfm_param);
  av2_setup_quant(tx_size, 0, AV2_XFORM_QUANT_B, cpi->oxcf.q_cfg.quant_b_adapt,
                  &quant_param);
  av2_setup_qmatrix(&cm->quant_params, xd, plane, tx_size, DCT_DCT,
                    &quant_param);
  av2_xform_quant(cm, x, plane, block, 0, 0, bsize, &txfm_param, &quant_param);
  const int n_coeffs = av2_get_max_eob(txfm_param.tx_size);
  for (int i = 0; i < n_coeffs; i++) {
    int abs_qcoeff = abs(qcoeff[i]);
    out->satd += abs(coeff[i]);
    out->satdq += abs_qcoeff;
    out->q_coeff_max = AVMMAX(out->q_coeff_max, abs_qcoeff);
    out->q_coeff_nonz += qcoeff[i] != 0;
  }

  if (p->eobs[block]) {
    txfm_param.eob = p->eobs[block];

    av2_highbd_inv_txfm_add(dqcoeff, pd->dst.buf, pd->dst.stride, &txfm_param);
  }
  int sse = 0;
  for (int i = 0; i < block_size_high[bsize]; i++) {
    for (int j = 0; j < block_size_wide[bsize]; j++) {
      int d = pd->dst.buf[i * pd->dst.stride + j] -
              x->plane[plane].src.buf[i * x->plane[plane].src.stride + j];
      sse += d * d;
    }
  }
  double mse =
      ((double)sse) / (block_size_high[bsize] * block_size_wide[bsize]);
  out->psnr = (float)(sse == 0 ? 70 : AVMMIN(70, 20 * log10(255 / sqrt(mse))));

  // TODO: figure out the way to do it w/o allocations
  p->coeff = NULL;
  p->qcoeff = NULL;
  p->dqcoeff = NULL;
  avm_free(p->eobs);
  p->eobs = NULL;
  avm_free(p->bobs);
  p->bobs = NULL;
  avm_free(p->txb_entropy_ctx);
  p->txb_entropy_ctx = NULL;
  xd->tree_type = old_tree_type;
}

#define MAX_BLK_SIZE (MAX_TX_SIZE << 1)
#define MAX_BLK_SQUARE (MAX_BLK_SIZE * MAX_BLK_SIZE)

static AVM_INLINE void av2_ml_part_split_features_square(AV2_COMP *const cpi,
                                                         MACROBLOCK *x,
                                                         int mi_row, int mi_col,
                                                         BLOCK_SIZE bsize,
                                                         float *out_features) {
  const AV2_COMMON *const cm = &cpi->common;
  MACROBLOCKD *xd = &x->e_mbd;
  MB_MODE_INFO *const mbmi = xd->mi[0];
  const int w_mi = mi_size_wide[bsize];
  const int h_mi = mi_size_high[bsize];
  DECLARE_ALIGNED(16, uint16_t, intrapred[MAX_TX_SQUARE]);

  // plus top line and left column
  BLOCK_SIZE subsize_sq = get_partition_subsize(
      get_partition_subsize(bsize, PARTITION_HORZ), PARTITION_VERT);
  if (subsize_sq == BLOCK_INVALID) {
    subsize_sq = get_partition_subsize(
        get_partition_subsize(bsize, PARTITION_VERT), PARTITION_HORZ);
  }

  if (subsize_sq != BLOCK_INVALID) {
    const int w_sub_mi = mi_size_wide[subsize_sq];
    const int h_sub_mi = mi_size_high[subsize_sq];
    TX_SIZE tx_sub_size = max_txsize_rect_lookup[subsize_sq];
    unsigned int best_sub_sse[2][2][3] = {
      { { INT_MAX, INT_MAX, INT_MAX }, { INT_MAX, INT_MAX, INT_MAX } },
      { { INT_MAX, INT_MAX, INT_MAX }, { INT_MAX, INT_MAX, INT_MAX } }
    };
    unsigned int best_sub_var[2][2][3] = {
      { { INT_MAX, INT_MAX, INT_MAX }, { INT_MAX, INT_MAX, INT_MAX } },
      { { INT_MAX, INT_MAX, INT_MAX }, { INT_MAX, INT_MAX, INT_MAX } }
    };
    PREDICTION_MODE best_sub_mode[2][2][3] = {
      { { MODE_INVALID, MODE_INVALID, MODE_INVALID },
        { MODE_INVALID, MODE_INVALID, MODE_INVALID } },
      { { MODE_INVALID, MODE_INVALID, MODE_INVALID },
        { MODE_INVALID, MODE_INVALID, MODE_INVALID } }
    };

    for (int row_off = 0, r_idx = 0; row_off < h_mi;
         row_off += h_sub_mi, ++r_idx) {
      int mi_row_left = xd->tile.mi_row_end - mi_row - row_off;
      // Don't process beyond the tile boundary
      if (mi_row_left < 0) break;
      for (int col_off = 0, c_idx = 0; col_off < w_mi;
           col_off += w_sub_mi, ++c_idx) {
        int mi_col_left = xd->tile.mi_col_end - mi_col - col_off;
        // Don't process beyond the tile boundary
        if (mi_col_left < 0) break;
        int src_off = (row_off << 2) * x->plane[0].src.stride + (col_off << 2);
        xd->mb_to_top_edge = -GET_MV_SUBPEL((mi_row + row_off) * MI_SIZE);
        xd->mb_to_left_edge = -GET_MV_SUBPEL((mi_col + col_off) * MI_SIZE);
        mbmi->sb_type[0] = subsize_sq;
        xd->up_available = (mi_row + row_off) > 0;
        xd->left_available = (mi_col + col_off) > 0;

        for (PREDICTION_MODE intra_sub_mode = INTRA_MODE_START;
             intra_sub_mode < INTRA_MODE_END; ++intra_sub_mode) {
          memset(intrapred, 0, sizeof(intrapred));
          xd->up_available = (mi_row + row_off) > 0;
          xd->left_available = (mi_col + col_off) > 0;
          av2_predict_intra_block(
              cm, xd, w_sub_mi << MI_SIZE_LOG2, h_sub_mi << MI_SIZE_LOG2,
              tx_sub_size, intra_sub_mode, 0, 0, x->plane[0].src.buf + src_off,
              x->plane[0].src.stride, intrapred, MAX_TX_SIZE, 0, 0, 0);

          unsigned int curr_sse = 0, curr_var = 0;
          curr_var = cpi->fn_ptr[enc_txsize_to_bsize[tx_sub_size]].vf(
              x->plane[0].src.buf + src_off, x->plane[0].src.stride, intrapred,
              MAX_TX_SIZE, &curr_sse);
          for (int cand = 0; cand < 3; cand++) {
            if (curr_sse < best_sub_sse[r_idx][c_idx][cand]) {
              for (int s = 2; s > cand; s--) {
                best_sub_sse[r_idx][c_idx][s] =
                    best_sub_sse[r_idx][c_idx][s - 1];
                best_sub_var[r_idx][c_idx][s] =
                    best_sub_var[r_idx][c_idx][s - 1];
                best_sub_mode[r_idx][c_idx][s] =
                    best_sub_mode[r_idx][c_idx][s - 1];
              }
              best_sub_sse[r_idx][c_idx][cand] = curr_sse;
              best_sub_var[r_idx][c_idx][cand] = curr_var;
              best_sub_mode[r_idx][c_idx][cand] = intra_sub_mode;
              break;
            }
          }
        }
        if (out_features) {
          const int sub_area_log2 =
              mi_size_wide_log2[subsize_sq] + mi_size_high_log2[subsize_sq] + 4;
          for (int cand = 0; cand < 3; ++cand) {
            int foff = r_idx * 4 + c_idx * 2 + cand * 8;
            out_features[FEATURE_INTRA_NORM_BEST_SSE_0_00 + foff] = logf(
                1.0f + (best_sub_sse[r_idx][c_idx][cand] >> sub_area_log2));
            out_features[FEATURE_INTRA_NORM_BEST_VAR_0_00 + foff] = logf(
                1.0f + (best_sub_var[r_idx][c_idx][cand] >> sub_area_log2));
          }
        }
      }
    }
  }
}

static AVM_INLINE void av2_ml_part_split_features_none(AV2_COMP *const cpi,
                                                       MACROBLOCK *x,
                                                       int mi_row, int mi_col,
                                                       BLOCK_SIZE bsize,
                                                       float *out_features) {
  const AV2_COMMON *const cm = &cpi->common;
  MACROBLOCKD *xd = &x->e_mbd;
  MB_MODE_INFO *const mbmi = xd->mi[0];
  const int w_mi = mi_size_wide[bsize];
  const int h_mi = mi_size_high[bsize];

  TX_SIZE tx_size = max_txsize_rect_lookup[bsize];

  unsigned int tx_w = tx_size_wide_unit[tx_size];
  unsigned int tx_h = tx_size_high_unit[tx_size];
  DECLARE_ALIGNED(16, uint16_t, intrapred[MAX_BLK_SQUARE]);

  xd->mb_to_top_edge = -GET_MV_SUBPEL(mi_row * MI_SIZE);
  xd->mb_to_left_edge = -GET_MV_SUBPEL(mi_col * MI_SIZE);
  mbmi->sb_type[0] = bsize;
  unsigned int best_sse[3] = { INT_MAX, INT_MAX, INT_MAX };
  unsigned int best_var[3] = { 0, 0, 0 };
  PREDICTION_MODE best_mode[3] = { MODE_INVALID, MODE_INVALID, MODE_INVALID };
  for (PREDICTION_MODE intra_mode = INTRA_MODE_START;
       intra_mode < INTRA_MODE_END; ++intra_mode) {
    unsigned int curr_sse = 0, curr_var = 0;
    memset(intrapred, 0, sizeof(intrapred));
    for (int row_off = 0; row_off < h_mi; row_off += tx_h) {
      for (int col_off = 0; col_off < w_mi; col_off += tx_w) {
        int src_off = (row_off << 2) * x->plane[0].src.stride + (col_off << 2);
        int intr_off = (row_off << 2) * MAX_BLK_SIZE + (col_off << 2);
        xd->up_available = (mi_row + row_off) > 0;
        xd->left_available = (mi_col + col_off) > 0;
        av2_predict_intra_block(cm, xd, w_mi << MI_SIZE_LOG2,
                                h_mi << MI_SIZE_LOG2, tx_size, intra_mode, 0, 0,
                                x->plane[0].src.buf + src_off,
                                x->plane[0].src.stride, intrapred + intr_off,
                                MAX_BLK_SIZE, 0, 0, 0);
        unsigned int tmp = 0;
        curr_var += cpi->fn_ptr[enc_txsize_to_bsize[tx_size]].vf(
            x->plane[0].src.buf + src_off, x->plane[0].src.stride,
            intrapred + intr_off, MAX_BLK_SIZE, &tmp);
        curr_sse += tmp;
      }
    }
    for (int cand = 0; cand < 3; cand++) {
      if (curr_sse < best_sse[cand]) {
        for (int s = 2; s > cand; s--) {
          best_sse[s] = best_sse[s - 1];
          best_var[s] = best_var[s - 1];
          best_mode[s] = best_mode[s - 1];
        }
        best_sse[cand] = curr_sse;
        best_var[cand] = curr_var;
        best_mode[cand] = intra_mode;
        break;
      }
    }
  }
  if (out_features) {
    const int blk_area_log2 =
        mi_size_wide_log2[bsize] + mi_size_high_log2[bsize] + 4;
    out_features[FEATURE_INTRA_NORM_BEST_0_SSE] =
        logf(1.0f + (best_sse[0] >> blk_area_log2));
    out_features[FEATURE_INTRA_NORM_BEST_0_VAR] =
        logf(1.0f + (best_var[0] >> blk_area_log2));
    out_features[FEATURE_INTRA_NORM_BEST_1_SSE] =
        logf(1.0f + (best_sse[1] >> blk_area_log2));
    out_features[FEATURE_INTRA_NORM_BEST_1_VAR] =
        logf(1.0f + (best_var[1] >> blk_area_log2));
    out_features[FEATURE_INTRA_NORM_BEST_2_SSE] =
        logf(1.0f + (best_sse[2] >> blk_area_log2));
    out_features[FEATURE_INTRA_NORM_BEST_2_VAR] =
        logf(1.0f + (best_var[2] >> blk_area_log2));
  }
}

static AVM_INLINE void av2_ml_part_split_features(AV2_COMP *const cpi,
                                                  MACROBLOCK *x, int mi_row,
                                                  int mi_col, BLOCK_SIZE bsize,
                                                  float *out_features) {
  MACROBLOCKD *xd = &x->e_mbd;
  MB_MODE_INFO *const mbmi = xd->mi[0];

  av2_setup_src_planes(x, cpi->source, mi_row, mi_col, 1, NULL);

  if (out_features) {
    // Q_INDEX
    const int dc_q =
        av2_dc_quant_QTX(x->qindex, 0, cpi->common.seq_params.base_y_dc_delta_q,
                         xd->bd) >>
        (xd->bd - 8);
    out_features[FEATURE_INTRA_LOG_QP_SQUARED] =
        logf(1.0f + (float)((int64_t)dc_q * (int64_t)dc_q) /
                        (256 << (2 * QUANT_TABLE_BITS)));

    // Neighbor stuff
    const int has_above = !!xd->above_mbmi;
    const int has_left = !!xd->left_mbmi;
    const BLOCK_SIZE above_bsize =
        has_above ? xd->above_mbmi->sb_type[xd->tree_type == CHROMA_PART]
                  : bsize;
    const BLOCK_SIZE left_bsize =
        has_left ? xd->left_mbmi->sb_type[xd->tree_type == CHROMA_PART] : bsize;

    out_features[FEATURE_INTRA_HAS_ABOVE] = (float)has_above;
    out_features[FEATURE_INTRA_LOG_ABOVE_WIDTH] =
        (float)mi_size_wide_log2[above_bsize];
    out_features[FEATURE_INTRA_LOG_ABOVE_HEIGHT] =
        (float)mi_size_high_log2[above_bsize];
    out_features[FEATURE_INTRA_HAS_LEFT] = (float)has_left;
    out_features[FEATURE_INTRA_LOG_LEFT_WIDTH] =
        (float)mi_size_wide_log2[left_bsize];
    out_features[FEATURE_INTRA_LOG_LEFT_HEIGHT] =
        (float)mi_size_high_log2[left_bsize];
  }

  int old1 = xd->mb_to_top_edge;
  int old2 = xd->mb_to_left_edge;
  int old3 = mbmi->sb_type[0];
  int old4 = mbmi->mrl_index;
  int old5 = mbmi->multi_line_mrl;
  mbmi->multi_line_mrl = 0;
  mbmi->mrl_index = 0;

  av2_ml_part_split_features_square(cpi, x, mi_row, mi_col, bsize,
                                    out_features);
  av2_ml_part_split_features_none(cpi, x, mi_row, mi_col, bsize, out_features);

  xd->mb_to_top_edge = old1;
  xd->mb_to_left_edge = old2;
  mbmi->sb_type[0] = old3;
  mbmi->mrl_index = old4;
  mbmi->multi_line_mrl = old5;
  avm_clear_system_state();
}

bool model_in_list(MODEL_TYPE model_type, MODEL_TYPE *out, int num_models) {
  for (int i = 0; i < num_models; i++)
    if (out[i] == model_type) return true;
  return false;
}

struct ModelParams {
  float thresh_low;
  float thresh_high;
  int qp_low;
  int qp_high;
};

#define TRY_MODEL(tgt_intra, tgt_level, tgt_bsize, model_type, low, high, \
                  qp_low, qp_high)                                        \
  {                                                                       \
    if (tgt_intra == intra && (tgt_level == harsh_level) &&               \
        tgt_bsize == bsize && *num_models < 4) {                          \
      if (!model_in_list(model_type, out, *num_models)) {                 \
        out[*num_models] = model_type;                                    \
        struct ModelParams tmp = { low, high, qp_low, qp_high };          \
        params[*num_models] = tmp;                                        \
        *num_models += 1;                                                 \
      }                                                                   \
    }                                                                     \
  }

static void get_model_type(bool intra, BLOCK_SIZE bsize, int harsh_level,
                           MODEL_TYPE *out, struct ModelParams *params,
                           int *num_models) {
  *num_models = 0;

  // CWG-E158
  //        intra? lvl sz  model                       low    high  qp_l qp_h
  TRY_MODEL(false, 1, 12, MODEL_INTER_NONE_64X64_110, 0.15f, 0.55f, 75, 110)
  TRY_MODEL(false, 0, 12, MODEL_INTER_NONE_64X64_110, 0.10f, 0.68f, 75, 110)
  TRY_MODEL(false, -1, 12, MODEL_INTER_NONE_64X64_110, 0.07f, 0.76f, 75, 110)
  TRY_MODEL(false, -2, 12, MODEL_INTER_NONE_64X64_110, 0.05f, 0.85f, 75, 110)
  TRY_MODEL(false, 1, 12, MODEL_INTER_NONE_64X64_135, 0.17f, 0.57f, 111, 135)
  TRY_MODEL(false, 0, 12, MODEL_INTER_NONE_64X64_135, 0.10f, 0.70f, 111, 135)
  TRY_MODEL(false, -1, 12, MODEL_INTER_NONE_64X64_135, 0.07f, 0.78f, 111, 135)
  TRY_MODEL(false, -2, 12, MODEL_INTER_NONE_64X64_135, 0.05f, 0.87f, 111, 135)
  TRY_MODEL(false, 1, 11, MODEL_INTER_NONE_BS11_110, 0.17f, 0.62f, 75, 110)
  TRY_MODEL(false, 0, 11, MODEL_INTER_NONE_BS11_110, 0.12f, 0.73f, 75, 110)
  TRY_MODEL(false, -1, 11, MODEL_INTER_NONE_BS11_110, 0.09f, 0.80f, 75, 110)
  TRY_MODEL(false, -2, 11, MODEL_INTER_NONE_BS11_110, 0.06f, 0.88f, 75, 110)
  TRY_MODEL(false, 1, 11, MODEL_INTER_NONE_BS11_135, 0.09f, 0.49f, 111, 135)
  TRY_MODEL(false, 0, 11, MODEL_INTER_NONE_BS11_135, 0.05f, 0.62f, 111, 135)
  TRY_MODEL(false, -1, 11, MODEL_INTER_NONE_BS11_135, 0.05f, 0.72f, 111, 135)
  TRY_MODEL(false, -2, 11, MODEL_INTER_NONE_BS11_135, 0.05f, 0.84f, 111, 135)
  TRY_MODEL(false, 1, 10, MODEL_INTER_NONE_BS10_110, 0.18f, 0.60f, 75, 110)
  TRY_MODEL(false, 0, 10, MODEL_INTER_NONE_BS10_110, 0.12f, 0.72f, 75, 110)
  TRY_MODEL(false, -1, 10, MODEL_INTER_NONE_BS10_110, 0.09f, 0.79f, 75, 110)
  TRY_MODEL(false, -2, 10, MODEL_INTER_NONE_BS10_110, 0.06f, 0.88f, 75, 110)
  TRY_MODEL(false, 1, 10, MODEL_INTER_NONE_BS10_135, 0.19f, 0.63f, 111, 135)
  TRY_MODEL(false, 0, 10, MODEL_INTER_NONE_BS10_135, 0.12f, 0.74f, 111, 135)
  TRY_MODEL(false, -1, 10, MODEL_INTER_NONE_BS10_135, 0.09f, 0.82f, 111, 135)
  TRY_MODEL(false, -2, 10, MODEL_INTER_NONE_BS10_135, 0.05f, 0.90f, 111, 135)
  TRY_MODEL(false, 1, 9, MODEL_INTER_NONE_32X32_110, 0.18f, 0.71f, 75, 110)
  TRY_MODEL(false, 0, 9, MODEL_INTER_NONE_32X32_110, 0.12f, 0.80f, 75, 110)
  TRY_MODEL(false, -1, 9, MODEL_INTER_NONE_32X32_110, 0.08f, 0.85f, 75, 110)
  TRY_MODEL(false, -2, 9, MODEL_INTER_NONE_32X32_110, 0.05f, 0.91f, 75, 110)
  TRY_MODEL(false, 1, 9, MODEL_INTER_NONE_32X32_135, 0.16f, 0.61f, 111, 135)
  TRY_MODEL(false, 0, 9, MODEL_INTER_NONE_32X32_135, 0.11f, 0.71f, 111, 135)
  TRY_MODEL(false, -1, 9, MODEL_INTER_NONE_32X32_135, 0.08f, 0.78f, 111, 135)
  TRY_MODEL(false, -2, 9, MODEL_INTER_NONE_32X32_135, 0.05f, 0.87f, 111, 135)
  TRY_MODEL(false, 1, 8, MODEL_INTER_NONE_BS8_110, 0.14f, 0.69f, 75, 110)
  TRY_MODEL(false, 0, 8, MODEL_INTER_NONE_BS8_110, 0.11f, 0.77f, 75, 110)
  TRY_MODEL(false, -1, 8, MODEL_INTER_NONE_BS8_110, 0.10f, 0.82f, 111, 135)
  TRY_MODEL(false, -2, 8, MODEL_INTER_NONE_BS8_110, 0.07f, 0.88f, 111, 135)
  TRY_MODEL(false, 1, 8, MODEL_INTER_NONE_BS8_135, 0.21f, 0.74f, 111, 135)
  TRY_MODEL(false, 0, 8, MODEL_INTER_NONE_BS8_135, 0.14f, 0.82f, 111, 135)
  TRY_MODEL(false, -1, 8, MODEL_INTER_NONE_BS8_135, 0.10f, 0.88f, 111, 135)
  TRY_MODEL(false, -2, 8, MODEL_INTER_NONE_BS8_135, 0.05f, 0.94f, 111, 135)
  TRY_MODEL(false, 1, 7, MODEL_INTER_NONE_BS7_110, 0.15f, 0.70f, 75, 110)
  TRY_MODEL(false, 0, 7, MODEL_INTER_NONE_BS7_110, 0.11f, 0.78f, 75, 110)
  TRY_MODEL(false, -1, 7, MODEL_INTER_NONE_BS7_110, 0.09f, 0.83f, 75, 110)
  TRY_MODEL(false, -2, 7, MODEL_INTER_NONE_BS7_110, 0.06f, 0.89f, 75, 110)
  TRY_MODEL(false, 1, 7, MODEL_INTER_NONE_BS7_135, 0.21f, 0.73f, 111, 135)
  TRY_MODEL(false, 0, 7, MODEL_INTER_NONE_BS7_135, 0.14f, 0.80f, 111, 135)
  TRY_MODEL(false, -1, 7, MODEL_INTER_NONE_BS7_135, 0.10f, 0.84f, 111, 135)
  TRY_MODEL(false, -2, 7, MODEL_INTER_NONE_BS7_135, 0.05f, 0.91f, 111, 135)
  TRY_MODEL(false, 1, 6, MODEL_INTER_NONE_16X16_110, 0.22f, 0.85f, 75, 110)
  TRY_MODEL(false, 0, 6, MODEL_INTER_NONE_16X16_110, 0.15f, 0.90f, 75, 110)
  TRY_MODEL(false, -1, 6, MODEL_INTER_NONE_16X16_110, 0.12f, 0.93f, 75, 110)
  TRY_MODEL(false, -2, 6, MODEL_INTER_NONE_16X16_110, 0.08f, 0.96f, 75, 110)
  TRY_MODEL(false, 1, 6, MODEL_INTER_NONE_16X16_135, 0.25f, 0.77f, 111, 135)
  TRY_MODEL(false, 0, 6, MODEL_INTER_NONE_16X16_135, 0.15f, 0.84f, 111, 135)
  TRY_MODEL(false, -1, 6, MODEL_INTER_NONE_16X16_135, 0.09f, 0.87f, 111, 135)
  TRY_MODEL(false, -2, 6, MODEL_INTER_NONE_16X16_135, 0.05f, 0.92f, 111, 135)
  TRY_MODEL(false, 1, 5, MODEL_INTER_NONE_BS5_110, 0.17f, 0.82f, 75, 110)
  TRY_MODEL(false, 0, 5, MODEL_INTER_NONE_BS5_110, 0.11f, 0.87f, 75, 110)
  TRY_MODEL(false, -1, 5, MODEL_INTER_NONE_BS5_110, 0.08f, 0.90f, 75, 110)
  TRY_MODEL(false, -2, 5, MODEL_INTER_NONE_BS5_110, 0.06f, 0.93f, 75, 110)
  TRY_MODEL(false, 1, 5, MODEL_INTER_NONE_BS5_135, 0.18f, 0.84f, 111, 135)
  TRY_MODEL(false, 0, 5, MODEL_INTER_NONE_BS5_135, 0.12f, 0.89f, 111, 135)
  TRY_MODEL(false, -1, 5, MODEL_INTER_NONE_BS5_135, 0.10f, 0.92f, 111, 135)
  TRY_MODEL(false, -2, 5, MODEL_INTER_NONE_BS5_135, 0.07f, 0.96f, 111, 135)
  TRY_MODEL(false, 1, 4, MODEL_INTER_NONE_BS4_110, 0.18f, 0.80f, 75, 110)
  TRY_MODEL(false, 0, 4, MODEL_INTER_NONE_BS4_110, 0.12f, 0.86f, 75, 110)
  TRY_MODEL(false, -1, 4, MODEL_INTER_NONE_BS4_110, 0.09f, 0.89f, 75, 110)
  TRY_MODEL(false, -2, 4, MODEL_INTER_NONE_BS4_110, 0.06f, 0.93f, 75, 110)
  TRY_MODEL(false, 1, 4, MODEL_INTER_NONE_BS4_135, 0.19f, 0.84f, 111, 135)
  TRY_MODEL(false, 0, 4, MODEL_INTER_NONE_BS4_135, 0.13f, 0.88f, 111, 135)
  TRY_MODEL(false, -1, 4, MODEL_INTER_NONE_BS4_135, 0.11f, 0.91f, 111, 135)
  TRY_MODEL(false, -2, 4, MODEL_INTER_NONE_BS4_135, 0.09f, 0.94f, 111, 135)
  // CWG-E070
  TRY_MODEL(true, -2, 15, MODEL_128X128, 0.3f, 0.98f, 60, 116)
  TRY_MODEL(true, -2, 12, MODEL_64X64, 0.03f, 0.7f, 60, 116)
  TRY_MODEL(true, -2, 9, MODEL_32X32, 0.052f, 0.45f, 60, 116)
  TRY_MODEL(true, -2, 6, MODEL_16X16, 0.002f, 0.35f, 60, 116)
  TRY_MODEL(true, -1, 15, MODEL_128X128, 0.4f, 0.97f, 60, 116)
  TRY_MODEL(true, -1, 12, MODEL_64X64, 0.1f, 0.7f, 60, 116)
  TRY_MODEL(true, -1, 9, MODEL_32X32, 0.052f, 0.45f, 60, 116)
  TRY_MODEL(true, -1, 6, MODEL_16X16, 0.005f, 0.32f, 60, 116)
  TRY_MODEL(true, 0, 15, MODEL_128X128, 0.5f, 0.95f, 60, 116)
  TRY_MODEL(true, 0, 12, MODEL_64X64, 0.2f, 0.66f, 60, 116)
  TRY_MODEL(true, 0, 9, MODEL_32X32, 0.07f, 0.44f, 60, 116)
  TRY_MODEL(true, 0, 6, MODEL_16X16, 0.01f, 0.3f, 60, 116)
  TRY_MODEL(true, 1, 15, MODEL_128X128, 0.7f, 0.95f, 60, 116)
  TRY_MODEL(true, 1, 12, MODEL_64X64, 0.25f, 0.7f, 60, 116)
  TRY_MODEL(true, 1, 9, MODEL_32X32, 0.07f, 0.25f, 60, 116)
  TRY_MODEL(true, 1, 6, MODEL_16X16, 0.01f, 0.15f, 60, 116)

  // CWG-E135
  TRY_MODEL(false, -2, 12, MODEL_INTER_SPLIT_64X64, 0.005f, 1.f, 75, 150)
  TRY_MODEL(false, -2, 9, MODEL_INTER_SPLIT_32X32, 0.005f, 1.f, 75, 150)
  TRY_MODEL(false, -2, 6, MODEL_INTER_SPLIT_16X16, 0.02f, 1.f, 75, 150)
  TRY_MODEL(false, -2, 3, MODEL_INTER_SPLIT_8X8, 0.1f, 1.f, 75, 150)
  TRY_MODEL(false, -1, 12, MODEL_INTER_SPLIT_64X64, 0.01f, 1.f, 75, 150)
  TRY_MODEL(false, -1, 9, MODEL_INTER_SPLIT_32X32, 0.02f, 1.f, 75, 150)
  TRY_MODEL(false, -1, 6, MODEL_INTER_SPLIT_16X16, 0.05f, 1.f, 75, 150)
  TRY_MODEL(false, -1, 3, MODEL_INTER_SPLIT_8X8, 0.25f, 1.f, 75, 150)
  TRY_MODEL(false, 0, 12, MODEL_INTER_SPLIT_64X64, 0.01f, 0.85f, 75, 150)
  TRY_MODEL(false, 0, 9, MODEL_INTER_SPLIT_32X32, 0.02f, 0.9f, 75, 150)
  TRY_MODEL(false, 0, 6, MODEL_INTER_SPLIT_16X16, 0.05f, 1.f, 75, 150)
  TRY_MODEL(false, 0, 3, MODEL_INTER_SPLIT_8X8, 0.25f, 1.f, 75, 150)
  TRY_MODEL(false, 1, 12, MODEL_INTER_SPLIT_64X64, 0.02f, 0.8f, 75, 150)
  TRY_MODEL(false, 1, 9, MODEL_INTER_SPLIT_32X32, 0.04f, 0.85f, 75, 150)
  TRY_MODEL(false, 1, 6, MODEL_INTER_SPLIT_16X16, 0.1f, 1.f, 75, 150)
  TRY_MODEL(false, 1, 3, MODEL_INTER_SPLIT_8X8, 0.35f, 1.f, 75, 150)
}

static float log_mag(MV mv) {
  double mag = sqrt(mv.col * mv.col + mv.row * mv.row);
  return (float)logl(1.0f + mag);
}

static float angle_rad(MV mv) {
  double mag = sqrt(mv.col * mv.col + mv.row * mv.row);
  return (float)(mag == 0 ? 0 : asin(mv.row / mag));
}

static void blk_features(float *out_features, int o_psnr, int o_log_mag,
                         int o_satdq, int o_satd, SimpleMotionData *sms,
                         int blk_area) {
  out_features[o_psnr + 0] = sms->residual_stats.psnr - 35;
  out_features[o_psnr + 1] = ((float)sms->residual_stats.q_coeff_max) / 1024;
  out_features[o_psnr + 2] =
      ((float)sms->residual_stats.q_coeff_nonz) / blk_area;
  out_features[o_log_mag + 0] = log_mag(sms->submv);
  out_features[o_log_mag + 1] = angle_rad(sms->submv);
  out_features[o_satdq] = logf(1.0f + sms->residual_stats.satdq);
  out_features[o_satd] = logf(1.0f + sms->residual_stats.satd);
}

static void av2_ml_part_split_features_inter(
    AV2_COMP *const cpi, MACROBLOCK *x, int mi_row, int mi_col,
    BLOCK_SIZE bsize, const TileInfo *tile_info, ThreadData *td,
    bool search_none_after_rect, float *out_features) {
  if (cpi->common.current_frame.frame_type != INTER_FRAME) return;
  const MACROBLOCKD *xd = &x->e_mbd;

  SimpleMotionData *blk_none =
      av2_get_sms_data(cpi, tile_info, x, mi_row, mi_col, bsize, td, true, 1);

  BLOCK_SIZE subsize_sq = get_partition_subsize(
      get_partition_subsize(bsize, PARTITION_HORZ), PARTITION_VERT);
  BLOCK_SIZE subsize_hor = get_partition_subsize(bsize, PARTITION_HORZ);
  BLOCK_SIZE subsize_ver = get_partition_subsize(bsize, PARTITION_VERT);
  if (subsize_sq == BLOCK_INVALID) {
    subsize_sq = get_partition_subsize(
        get_partition_subsize(bsize, PARTITION_VERT), PARTITION_HORZ);
  }

  if (subsize_sq != BLOCK_INVALID) {
    int w_sub_mi = mi_size_wide[subsize_sq];
    int h_sub_mi = mi_size_high[subsize_sq];

    SimpleMotionData *blk_sq_0 = av2_get_sms_data(
        cpi, tile_info, x, mi_row, mi_col, subsize_sq, td, true, 1);
    SimpleMotionData *blk_sq_1 = av2_get_sms_data(
        cpi, tile_info, x, mi_row, mi_col + w_sub_mi, subsize_sq, td, true, 1);
    SimpleMotionData *blk_sq_2 = av2_get_sms_data(
        cpi, tile_info, x, mi_row + h_sub_mi, mi_col, subsize_sq, td, true, 1);
    SimpleMotionData *blk_sq_3 =
        av2_get_sms_data(cpi, tile_info, x, mi_row + h_sub_mi,
                         mi_col + w_sub_mi, subsize_sq, td, true, 1);

    if (out_features) {
      int blk_area = block_size_wide[bsize] * block_size_high[bsize];
      out_features[FEATURE_INTER_RD_MULT] = logf(1.0f + blk_none->rdmult);
      out_features[FEATURE_INTER_SWITCH] = search_none_after_rect;
      out_features[FEATURE_INTER_PART_T] = xd->tree_type;

      blk_features(out_features, FEATURE_INTER_FULL_PSNR,
                   FEATURE_INTER_FULL_LOG_MAG, FEATURE_INTER_FULL_LOG_SATDQ,
                   FEATURE_INTER_FULL_LOG_SATD, blk_none, blk_area);
      blk_features(out_features, FEATURE_INTER_SQ_0_PSNR,
                   FEATURE_INTER_SQ_0_LOG_MAG, FEATURE_INTER_SQ_0_LOG_SATDQ,
                   FEATURE_INTER_SQ_0_LOG_SATD, blk_sq_0, blk_area);
      blk_features(out_features, FEATURE_INTER_SQ_1_PSNR,
                   FEATURE_INTER_SQ_1_LOG_MAG, FEATURE_INTER_SQ_1_LOG_SATDQ,
                   FEATURE_INTER_SQ_1_LOG_SATD, blk_sq_1, blk_area);
      blk_features(out_features, FEATURE_INTER_SQ_2_PSNR,
                   FEATURE_INTER_SQ_2_LOG_MAG, FEATURE_INTER_SQ_2_LOG_SATDQ,
                   FEATURE_INTER_SQ_2_LOG_SATD, blk_sq_2, blk_area);
      blk_features(out_features, FEATURE_INTER_SQ_3_PSNR,
                   FEATURE_INTER_SQ_3_LOG_MAG, FEATURE_INTER_SQ_3_LOG_SATDQ,
                   FEATURE_INTER_SQ_3_LOG_SATD, blk_sq_3, blk_area);
    }
  }

  // The horizontal and vertical features are only used by the partition NONE
  // pruning models. When partition NONE is disabled, do not calculate them
  // this is because the SMS will cache the block statistics the first time
  // it encounters it, and the partitions below introduce a different order
  // of iterating those blocks, which creates the occasional stats mismatch
  // when the NONE partitions are turned off.
  if (subsize_hor != BLOCK_INVALID && cpi->sf.part_sf.prune_none_with_ml) {
    int h_sub_mi = mi_size_high[subsize_hor];
    SimpleMotionData *blk_hor_0 = av2_get_sms_data(
        cpi, tile_info, x, mi_row, mi_col, subsize_hor, td, true, 1);
    SimpleMotionData *blk_hor_1 = av2_get_sms_data(
        cpi, tile_info, x, mi_row + h_sub_mi, mi_col, subsize_hor, td, true, 1);

    if (out_features) {
      int blk_area = block_size_wide[bsize] * block_size_high[bsize];
      blk_features(out_features, FEATURE_INTER_HOR_0_PSNR,
                   FEATURE_INTER_HOR_0_LOG_MAG, FEATURE_INTER_HOR_0_LOG_SATDQ,
                   FEATURE_INTER_HOR_0_LOG_SATD, blk_hor_0, blk_area);
      blk_features(out_features, FEATURE_INTER_HOR_1_PSNR,
                   FEATURE_INTER_HOR_1_LOG_MAG, FEATURE_INTER_HOR_1_LOG_SATDQ,
                   FEATURE_INTER_HOR_1_LOG_SATD, blk_hor_1, blk_area);
    }
  }

  if (subsize_ver != BLOCK_INVALID && cpi->sf.part_sf.prune_none_with_ml) {
    int w_sub_mi = mi_size_wide[subsize_ver];
    SimpleMotionData *blk_ver_0 = av2_get_sms_data(
        cpi, tile_info, x, mi_row, mi_col, subsize_ver, td, true, 1);
    SimpleMotionData *blk_ver_1 = av2_get_sms_data(
        cpi, tile_info, x, mi_row, mi_col + w_sub_mi, subsize_ver, td, true, 1);

    if (out_features) {
      int blk_area = block_size_wide[bsize] * block_size_high[bsize];

      blk_features(out_features, FEATURE_INTER_VER_0_PSNR,
                   FEATURE_INTER_VER_0_LOG_MAG, FEATURE_INTER_VER_0_LOG_SATDQ,
                   FEATURE_INTER_VER_0_LOG_SATD, blk_ver_0, blk_area);
      blk_features(out_features, FEATURE_INTER_VER_1_PSNR,
                   FEATURE_INTER_VER_1_LOG_MAG, FEATURE_INTER_VER_1_LOG_SATDQ,
                   FEATURE_INTER_VER_1_LOG_SATD, blk_ver_1, blk_area);
    }
  }

  avm_clear_system_state();
}

int av2_ml_part_split_infer(AV2_COMP *const cpi, MACROBLOCK *x, int mi_row,
                            int mi_col, BLOCK_SIZE bsize,
                            const TileInfo *tile_info, ThreadData *td,
                            bool search_none_after_rect, bool *prune_list) {
  for (int i = 0; i < 2; i++) {
    prune_list[i] = false;
  }

  const int mi_high = mi_size_high[bsize];
  const int mi_wide = mi_size_wide[bsize];
  const AV2_COMMON *const cm = &cpi->common;
  if (mi_col + mi_wide > cm->mi_params.mi_cols ||
      mi_row + mi_high > cm->mi_params.mi_rows)
    return ML_PART_DONT_FORCE;

  const MACROBLOCKD *xd = &x->e_mbd;
  int qp = cpi->common.quant_params.base_qindex;
  bool key_frame = cpi->common.current_frame.frame_type == KEY_FRAME;

  int qp_offset;
  switch (cm->seq_params.bit_depth) {
    case AVM_BITS_10: qp_offset = qindex_10b_offset[1]; break;
    case AVM_BITS_12: qp_offset = qindex_12b_offset[1]; break;
    default: qp_offset = 0; break;
  }

  int harsh_level = key_frame ? cpi->sf.part_sf.prune_split_ml_level
                              : cpi->sf.part_sf.prune_split_ml_level_inter;

  // use intra model only for key frames for now
  if (xd->tree_type == CHROMA_PART) return ML_PART_DONT_FORCE;
  MODEL_TYPE model_types[4] = { 0, 0, 0, 0 };
  struct ModelParams model_params[4];
  int num_models = 0;
  get_model_type(key_frame, bsize, harsh_level, model_types, model_params,
                 &num_models);
  float ml_input[FEATURE_INTER_MAX] = { 0.0f };
  float ml_output[1] = { 0.0f };
  bool has_features = false;
  for (int mi = 0; mi < num_models; mi++) {
    MODEL_TYPE model_type = model_types[mi];
    struct ModelParams params = model_params[mi];
    bool model_disabled =
        qp > (params.qp_high + qp_offset) || qp < (params.qp_low + qp_offset);
    model_disabled |= get_model_part_type(model_type) == PT_NONE &&
                      !cpi->sf.part_sf.prune_none_with_ml;
    if (model_disabled) continue;

    if (!has_features) {
      if (key_frame) {
        av2_ml_part_split_features(cpi, x, mi_row, mi_col, bsize, ml_input);
      } else {
        av2_ml_part_split_features_inter(cpi, x, mi_row, mi_col, bsize,
                                         tile_info, td, search_none_after_rect,
                                         ml_input);
      }
      has_features = true;
    }

    // printf("%s l:%f h:%f qpl:%d qph:%d\n", get_model_name(model_type),
    //        params.thresh_low, params.thresh_high, params.qp_low,
    //        params.qp_high);
    bool had_error = av2_part_prune_tflite_exec(&td->partition_model, ml_input,
                                                ml_output, model_type);

    assert(!had_error);

    if (had_error)
      continue;
    else {
      int part_type = get_model_part_type(model_type);
      if (part_type >= 0 && part_type < 2) {
        bool low_test = ml_output[0] < params.thresh_low;
        bool high_test = ml_output[0] > params.thresh_high;
        if (high_test) return ML_PART_FORCE_NONE + part_type;
        if (low_test) prune_list[part_type] = true;
      }
    }
  }

  return ML_PART_DONT_FORCE;
}