#ifndef AOM_AV1_ENCODER_INTERINTRA_ML_DATA_COLLECT_H_
#define AOM_AV1_ENCODER_INTERINTRA_ML_DATA_COLLECT_H_

#include <stdint.h>
#include "av1/encoder/block.h"
#include "av1/encoder/encoder.h"

#ifdef __cplusplus
extern "C" {
#endif

// Data collection for the interintra ML training data. The program must be
// run in single-threaded mode.

// Build the inter-predictor (with border), and record it, along with the
// intra-predictor and source image. Note that only a subset of the border
// may end up being recorded.
void av1_interintra_ml_data_collect(const AV1_COMP *const cpi,
                                    MACROBLOCK *const x, BLOCK_SIZE bsize);

// After the data is collected, it is not immediately written out, since the
// encoder may later decide to save it as a skip block. This function
// indicates that the data is safe to write out (not a skip block).
void av1_interintra_ml_data_collect_finalize();

// The block ended up being skipped. Abandon the previously collected data.
void av1_interintra_ml_data_collect_abandon();

// Is this block a valid block for collecting interintra ML data?
bool av1_interintra_ml_data_collect_valid(MACROBLOCK *const x,
                                          BLOCK_SIZE bsize);

// RD-cost of inter-intra modes should be multiplied by this number (and divided
// by 100) to bias the results into picking inter-intra modes more often.
uint8_t av1_interintra_bias();

// These functions / data structures are internal implementation details, but
// are exposed for testing.

// The data structure encoding the information we wish to write out,
// for a given plane.
typedef struct IIMLPlaneInfo {
  // Width and height of the block.
  uint8_t width;
  uint8_t height;
  // Plane, one of AOM_PLANE_Y, AOM_PLANE_U, AOM_PLANE_V.
  uint8_t plane;
  // Number of bits per pixel, either 8, 10, or 12.
  uint8_t bitdepth;
  // The border width recorded for this block.
  uint8_t border;
  // The location of this block (in pixels) in the video.
  int x;
  int y;
  int frame_order_hint;
  // Information about the rate-cost calculation.
  int32_t lambda;
  uint8_t ref_q;
  uint8_t base_q;

  // The original source block. Size is width * height. If bit-depth
  // is 10 or 12, it is actually a uint16_t pointer.
  uint8_t *source_image;

  // The top (border + width) * (border) pixels, followed by
  // the (border * height) left pixels, in raster scan order, of
  // the top-left border of the intra-predictor. If bit-depth is
  // 10 or 12, actually a uint16_t pointer.
  uint8_t *intrapred_lshape;

  // The (border + width) * (border + height) inter-predictor.
  // The inter-predictor starts at the (width * border + border) offset.
  // If bit-depth is 10 or 12, actually a uint16_t pointer.
  uint8_t *interpred;
  // 0 == intra, 1 == inter, 2 == inter-intra (smooth),
  // 3 == inter-intra (wedge); currently only inter and inter-intra are output.
  uint8_t prediction_type;
  // The (width * height) final predictor. If bit-depth is 10 or 12, actually
  // a uint16_t pointer.
  uint8_t *predictor;
  // If an inter-intra-predictor, this value represents a bias -- the percent
  // that the RD-cost for inter-intra-predictors was shifted down by, to
  // bias the data into using the inter-intra predictor more often. E.g.,
  // 90 means that the RD-cost is multiplied by 0.9, 80 means it is multiplied
  // by 0.8, etc.
  uint8_t interintra_bias;
} IIMLPlaneInfo;

// Serialize the plane info data structure into a byte array. Allocates
// memory into buf (the caller must free it afterward).
void av1_interintra_ml_data_collect_serialize(IIMLPlaneInfo *info,
                                              uint8_t **buf, size_t *buf_size);

// Copy over the top and left pixels (in an L shape) from the top-left
// corner, into the buffer. Note that source points to the start of the
// intrapred block (the border is at a negative offset).
void av1_interintra_ml_data_collect_copy_intrapred_lshape(
    uint8_t *dst, const uint8_t *src, int src_stride, int width, int height,
    int border, int bitdepth, bool is_src_hbd);

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

#endif  // AOM_AV1_ENCODER_INTERINTRA_ML_DATA_COLLECT_H_