#include <stdbool.h>
#include <stdio.h>
#include <stdlib.h>

#include "av1/common/reconinter.h"
#include "av1/common/reconintra.h"
#include "av1/encoder/interintra_ml_data_collect.h"
#include "av1/encoder/reconinter_enc.h"

#define BORDER_SIZE 4  // Only generate data with a border of 4 pixels.

// Static references to the currently captured Y/U/V planes. Held in
// memory until we can determine if this block is a skip-block (if so,
// the data is not written out).
static IIMLPlaneInfo *INFO_Y = NULL;
static IIMLPlaneInfo *INFO_U = NULL;
static IIMLPlaneInfo *INFO_V = NULL;

// Static reference to the file where the data is being stored.
static FILE *FP = NULL;

// Clean-up function called on program exit. Closes file descriptor,
// flushing writes.
static void cleanup_fp() {
  if (FP != NULL) {
    int r = fclose(FP);
    assert(r == 0);
    (void)r;
  }
}

static uint8_t read_interintra_bias() {
  const char *bias = getenv("INTERINTRA_BIAS");
  if (bias == NULL) {
    return 100;
  }
  int i = atoi(bias);
  if (i < 0 || i > 100) {
    fprintf(stderr, "INTERINTRA_BIAS must be between 0 and 100.\n");
    exit(1);
  }
  return (uint8_t)i;
}

static uint8_t INTERINTRA_BIAS = 255;

uint8_t av1_interintra_bias() {
  if (INTERINTRA_BIAS == 255) {
    INTERINTRA_BIAS = read_interintra_bias();
  }
  return INTERINTRA_BIAS;
}

// Function called on every invocation of data collection. Initializes
// the file output and registers the cleanup function (if not already done).
// The output file name is "interintra_ml_data_collect.bin" by default.
// The file name can be changed by setting INTERINTRA_ML_DATA_COLLECT
// environment variable.
static void init_first_run() {
  if (FP != NULL) {
    return;
  }
  const char *filename = getenv("INTERINTRA_ML_DATA_COLLECT");
  if (filename == NULL) {
    filename = "interintra_ml_data_collect.bin";
  }
  FP = fopen(filename, "wb");
  assert(FP != NULL);
  int r = atexit(cleanup_fp);
  assert(r == 0);
  (void)r;
}

// Copy the value from the source into the destination, at the given offsets.
// If both source and destination are high-bitdepth, will copy. Otherwise,
// if source is high-bitdepth but destination is lower, will only copy the
// first 8 bits.
static void copy_value(uint8_t *dst, int dst_offset, const uint8_t *src,
                       int src_offset, int bitdepth, bool is_hbd) {
  if (bitdepth > 8) {
    assert(is_hbd);  // If the IIMLPlaneInfo is high-bitdepth, then so must the
    // AV1 data structures.
    uint16_t *dst16 = (uint16_t *)dst;
    uint16_t *src16 = CONVERT_TO_SHORTPTR(src);
    dst16[dst_offset] = src16[src_offset];
  } else if (bitdepth == 8 && !is_hbd) {
    // A pure copy.
    dst[dst_offset] = src[src_offset];
  } else {
    // There is no case when bitdepth > 8 and !is_hbd.
    assert(bitdepth == 8 && is_hbd);
    uint16_t *src16 = CONVERT_TO_SHORTPTR(src);
    assert(src16[src_offset] <= 255);
    dst[dst_offset] = (uint8_t)src16[src_offset];
  }
}

static void copy_source_image(IIMLPlaneInfo *info, MACROBLOCK *const x) {
  // Overallocate, memory is not an issue.
  info->source_image = malloc(info->width * info->height * sizeof(uint16_t));
  assert(info->source_image != NULL);

  const struct buf_2d *ref = &x->plane[info->plane].src;
  const int stride = ref->stride;
  const uint8_t *buf = ref->buf0 + info->y * stride + info->x;
  const bool is_hbd = is_cur_buf_hbd(&x->e_mbd);

  for (int j = 0; j < info->height; ++j) {
    for (int i = 0; i < info->width; ++i) {
      copy_value(info->source_image, j * info->width + i, buf, j * stride + i,
                 info->bitdepth, is_hbd);
    }
  }
}

static void copy_intrapred_lshape(IIMLPlaneInfo *info, MACROBLOCK *const x) {
  info->intrapred_lshape =
      malloc(sizeof(uint16_t) * ((info->width + info->border) * info->border +
                                 info->border * info->height));
  assert(info->intrapred_lshape != NULL);
  MACROBLOCKD *const xd = &x->e_mbd;
  const struct macroblockd_plane *const pd = &xd->plane[info->plane];
  const struct buf_2d *ref = &pd->dst;
  const int stride = ref->stride;
  uint8_t *buf = ref->buf0 + info->y * stride + info->x;
  av1_interintra_ml_data_collect_copy_intrapred_lshape(
      info->intrapred_lshape, buf, stride, info->width, info->height,
      info->border, info->bitdepth, is_cur_buf_hbd(xd));
}

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) {
  // Point back at the start of the L-region.
  src -= (border * src_stride + border);
  int info_i = 0;
  // Copy over the top part.
  for (int j = 0; j < border; ++j) {
    for (int i = 0; i < border + width; ++i) {
      copy_value(dst, info_i++, src, j * src_stride + i, bitdepth, is_src_hbd);
    }
  }

  // Copy over the side pixels.
  for (int j = border; j < border + height; ++j) {
    for (int i = 0; i < border; ++i) {
      copy_value(dst, info_i++, src, j * src_stride + i, bitdepth, is_src_hbd);
    }
  }
}

static void copy_interpred(IIMLPlaneInfo *info, const AV1_COMP *const cpi,
                           MACROBLOCK *const x) {
  info->interpred = malloc(sizeof(uint16_t) * (info->width + info->border) *
                           (info->height + info->border));

  const AV1_COMMON *const cm = &cpi->common;
  MACROBLOCKD *const xd = &x->e_mbd;

  uint8_t *orig_buf = xd->plane[info->plane].dst.buf;
  int orig_stride = xd->plane[info->plane].dst.stride;

  uint8_t *interpred;
  int interpred_stride;
  const int border =
      av1_calc_border(xd, AOM_PLANE_Y, false /* build for obmc */);
  assert(border >= BORDER_SIZE);
  av1_alloc_buf_with_border(&interpred, &interpred_stride, border,
                            is_cur_buf_hbd(xd));

  xd->plane[info->plane].dst.buf = interpred;
  xd->plane[info->plane].dst.stride = interpred_stride;

  av1_enc_build_border_only_inter_predictor(cm, xd, xd->mi_row, xd->mi_col,
                                            info->plane);

  int info_i = 0;
  for (int j = -1 * info->border; j < info->height; ++j) {
    for (int i = -1 * info->border; i < info->width; ++i) {
      copy_value(info->interpred, info_i++, interpred, j * interpred_stride + i,
                 info->bitdepth, is_cur_buf_hbd(xd));
    }
  }
  xd->plane[info->plane].dst.buf = orig_buf;
  xd->plane[info->plane].dst.stride = orig_stride;
  av1_free_buf_with_border(interpred, interpred_stride, border,
                           is_cur_buf_hbd(xd));
}

static void copy_predictor(IIMLPlaneInfo *info, const AV1_COMP *const cpi,
                           MACROBLOCK *const x, BLOCK_SIZE bsize) {
  info->predictor = malloc(sizeof(uint16_t) * info->width * info->height);

  const AV1_COMMON *const cm = &cpi->common;
  MACROBLOCKD *const xd = &x->e_mbd;

  // Assumes that xd->plane[info->plane].dst.buf can be overwritten in
  // the inter-predictor part.
  av1_enc_build_inter_predictor(cm, xd, xd->mi_row, xd->mi_col, NULL, bsize,
                                info->plane, info->plane);

  uint8_t *predictor = xd->plane[info->plane].dst.buf;
  int stride = xd->plane[info->plane].dst.stride;
  int info_i = 0;
  for (int j = 0; j < info->height; ++j) {
    for (int i = 0; i < info->width; ++i) {
      copy_value(info->predictor, info_i++, predictor, j * stride + i,
                 info->bitdepth, is_cur_buf_hbd(xd));
    }
  }
}

// 1 == inter, 2 == inter-intra (smooth), 3 == inter-intra (wedge)
static int get_prediction_type(MB_MODE_INFO *mbmi, BLOCK_SIZE bsize) {
  if (!is_interintra_pred(mbmi)) {
    return 1;
  }
  const bool wedge_case =
      mbmi->use_wedge_interintra && is_interintra_wedge_used(bsize);
  return wedge_case ? 3 : 2;
}

// Initializes the IIMLPlaneInfo structure.
static IIMLPlaneInfo *init_plane_info(const AV1_COMP *const cpi,
                                      MACROBLOCK *const x, BLOCK_SIZE bsize,
                                      int plane) {
  MACROBLOCKD *const xd = &x->e_mbd;
  MB_MODE_INFO *mbmi = xd->mi[0];
  const AV1_COMMON *const cm = &cpi->common;
  const struct macroblockd_plane *const pd = &xd->plane[plane];
  const BLOCK_SIZE plane_bsize =
      get_plane_block_size(bsize, pd->subsampling_x, pd->subsampling_y);

  assert(is_inter_block(mbmi));
  assert(!has_second_ref(mbmi));
  assert(!is_intrabc_block(mbmi));
  // Data is packed differently for smaller block sizes. Not supported.
  assert(block_size_wide[plane_bsize] >= 4);
  assert(block_size_high[plane_bsize] >= 4);

  IIMLPlaneInfo *info = malloc(sizeof(IIMLPlaneInfo));
  assert(info != NULL);

  info->width = block_size_wide[plane_bsize];
  info->height = block_size_high[plane_bsize];
  info->plane = plane;
  info->bitdepth = (uint8_t)xd->bd;
  info->border = BORDER_SIZE;
  info->x = (xd->mi_col * MI_SIZE) >> pd->subsampling_x;
  info->y = (xd->mi_row * MI_SIZE) >> pd->subsampling_y;
  info->frame_order_hint = (int)cm->current_frame.order_hint;

  RefCntBuffer *refbuf = get_ref_frame_buf(cm, xd->mi[0]->ref_frame[0]);
  assert(refbuf != NULL);
  info->ref_q = refbuf->base_qindex;
  info->base_q = cm->base_qindex;
  info->lambda = av1_compute_rd_mult_based_on_qindex(cpi, cm->base_qindex);
  assert(info->lambda > 0);
  assert(info->lambda < 10 * 1000 * 1000);

  copy_source_image(info, x);
  copy_intrapred_lshape(info, x);
  copy_interpred(info, cpi, x);
  info->prediction_type = get_prediction_type(mbmi, bsize);
  copy_predictor(info, cpi, x, bsize);
  info->interintra_bias = av1_interintra_bias();
  return info;
}

void write_uint8_and_advance(uint8_t **buf, uint8_t b) {
  **buf = b;
  (*buf)++;
}

void write_int32_and_advance(uint8_t **buf, int32_t value) {
  assert(value >= 0);
  for (size_t i = 0; i < sizeof(value); ++i) {
    // Little-endian order.
    uint8_t byte = 0xff & value;
    write_uint8_and_advance(buf, byte);
    value >>= 8;
  }
}

void write_buffer_and_advance(uint8_t **buf, uint8_t *src, size_t size,
                              int bitdepth) {
  if (bitdepth == 8) {
    for (size_t i = 0; i < size; ++i) {
      write_uint8_and_advance(buf, src[i]);
    }
    return;
  }

  uint16_t *src16 = (uint16_t *)src;
  for (size_t i = 0; i < size; ++i) {
    // Little-endian order.
    uint8_t b1 = (uint8_t)(0xff & src16[i]);
    uint8_t b2 = (uint8_t)(0xff & (src16[i] >> 8));
    write_uint8_and_advance(buf, b1);
    write_uint8_and_advance(buf, b2);
  }
}

void av1_interintra_ml_data_collect_serialize(IIMLPlaneInfo *info,
                                              uint8_t **buf, size_t *buf_size) {
  const int bytes_per_pixel = info->bitdepth == 8 ? 1 : 2;
  const size_t source_size = info->width * info->height;
  const size_t intrapred_lshape_size =
      ((info->border + info->width) * info->border +
       info->border * info->height);
  const size_t interpred_size =
      ((info->border + info->width) * (info->border + info->height));
  *buf_size =
      1 /* width */ + 1 /* height */ + 1 /* plane */ + 1 /* bitdepth */ +
      1 /* border */ + sizeof(int32_t) /* x */ + sizeof(int32_t) /* y */ +
      sizeof(int32_t) /* frame order hint */ + sizeof(int32_t) /* lambda */ +
      1 /* base_q */ + bytes_per_pixel * source_size /* source image */ +
      bytes_per_pixel * intrapred_lshape_size /* reconstruction border */ +
      1 /* prediction_type */ + bytes_per_pixel * source_size /* predictor */ +
      1 /* ref_q */ +
      bytes_per_pixel * interpred_size /* inter-predictor + border */ +
      1 /* interintra_bias */;
  *buf = malloc(*buf_size);
  assert(*buf != NULL);
  uint8_t *start = *buf;

  write_uint8_and_advance(buf, info->width);
  write_uint8_and_advance(buf, info->height);
  write_uint8_and_advance(buf, info->plane);
  write_uint8_and_advance(buf, info->bitdepth);
  write_uint8_and_advance(buf, info->border);
  write_int32_and_advance(buf, info->x);
  write_int32_and_advance(buf, info->y);
  write_int32_and_advance(buf, info->frame_order_hint);
  write_int32_and_advance(buf, info->lambda);
  write_uint8_and_advance(buf, info->base_q);
  write_buffer_and_advance(buf, info->source_image, source_size,
                           info->bitdepth);
  write_buffer_and_advance(buf, info->intrapred_lshape, intrapred_lshape_size,
                           info->bitdepth);
  write_uint8_and_advance(buf, info->prediction_type);
  write_buffer_and_advance(buf, info->predictor, source_size, info->bitdepth);
  write_uint8_and_advance(buf, info->ref_q);
  write_buffer_and_advance(buf, info->interpred, interpred_size,
                           info->bitdepth);
  write_uint8_and_advance(buf, info->interintra_bias);
  assert(start + *buf_size == *buf);
  *buf = start;
}

void av1_interintra_ml_data_collect(const AV1_COMP *const cpi,
                                    MACROBLOCK *const x, BLOCK_SIZE bsize) {
  init_first_run();
  assert(INFO_Y == NULL);
  assert(INFO_U == NULL);
  assert(INFO_V == NULL);
  INFO_Y = init_plane_info(cpi, x, bsize, AOM_PLANE_Y);
  assert(INFO_Y != NULL);

  // Sometimes the chroma is derived from the luma. Check if chroma is
  // available.
  MACROBLOCKD *const xd = &x->e_mbd;
  if (xd->mi[0]->chroma_ref_info.is_chroma_ref) {
    INFO_U = init_plane_info(cpi, x, bsize, AOM_PLANE_U);
    INFO_V = init_plane_info(cpi, x, bsize, AOM_PLANE_V);
    assert(INFO_U != NULL);
    assert(INFO_V != NULL);
  }
}

// Write out the three planes to disk.
void av1_interintra_ml_data_collect_finalize() {
  uint8_t *buf;
  size_t buf_size;

  // If luma is not present, then neither is chroma.
  if (INFO_Y == NULL) {
    assert(INFO_U == NULL);
    assert(INFO_V == NULL);
    return;
  }
  av1_interintra_ml_data_collect_serialize(INFO_Y, &buf, &buf_size);
  size_t written = fwrite(buf, sizeof(*buf), buf_size, FP);
  assert(written == buf_size);
  free(buf);

  if (INFO_U != NULL) {
    av1_interintra_ml_data_collect_serialize(INFO_U, &buf, &buf_size);
    written = fwrite(buf, sizeof(*buf), buf_size, FP);
    assert(written == buf_size);
    free(buf);
  }

  if (INFO_V != NULL) {
    av1_interintra_ml_data_collect_serialize(INFO_V, &buf, &buf_size);
    written = fwrite(buf, sizeof(*buf), buf_size, FP);
    assert(written == buf_size);
    free(buf);
  }

  av1_interintra_ml_data_collect_abandon();
  assert(INFO_Y == NULL);
  assert(INFO_U == NULL);
  assert(INFO_V == NULL);
}

static void destroy_info(IIMLPlaneInfo *info) {
  if (info != NULL) {
    free(info->source_image);
    free(info->intrapred_lshape);
    free(info->interpred);
    free(info);
  }
}

void av1_interintra_ml_data_collect_abandon() {
  if (INFO_Y == NULL) {
    assert(INFO_U == NULL);
    assert(INFO_V == NULL);
    return;
  }
  destroy_info(INFO_Y);
  destroy_info(INFO_U);
  destroy_info(INFO_V);
  INFO_Y = NULL;
  INFO_U = NULL;
  INFO_V = NULL;
}

static int calc_intra_border(MACROBLOCK *const x, BLOCK_SIZE bsize) {
  MACROBLOCKD *const xd = &x->e_mbd;
  int border = BORDER_SIZE;
  // For each plane, check how much is available on the top and left,
  // shrinking the border if needed.
  for (int plane = AOM_PLANE_Y; plane <= AOM_PLANE_V; ++plane) {
    border = AOMMIN(border, av1_intra_top_available(xd, plane));
    border = AOMMIN(border, av1_intra_left_available(xd, plane));

    // If either the bottom or the right side of the block is partially
    // unavailable, set the entire border to 0 -- we cannot extract the
    // full border.
    const struct macroblockd_plane *const pd = &xd->plane[plane];
    const BLOCK_SIZE plane_bsize =
        get_plane_block_size(bsize, pd->subsampling_x, pd->subsampling_y);
    const TX_SIZE tx_size = max_txsize_rect_lookup[plane_bsize];
    if (av1_intra_bottom_unavailable(xd, plane, tx_size) ||
        av1_intra_right_unavailable(xd, plane, tx_size)) {
      return 0;
    }
  }
  return border;
}

bool av1_interintra_ml_data_collect_valid(MACROBLOCK *const x,
                                          BLOCK_SIZE bsize) {
  MACROBLOCKD *const xd = &x->e_mbd;
  const int inter_border =
      av1_calc_border(xd, AOM_PLANE_Y, false /* build for obmc */);
  const int intra_border = calc_intra_border(x, bsize);
  // Only process 8x8 or larger blocks -- if a dimension is 4, then the
  // chroma dimension can be 2, which has a different packing structure.
  return inter_border >= BORDER_SIZE && intra_border >= BORDER_SIZE &&
         !x->skip && block_size_wide[bsize] >= 8 && block_size_high[bsize] >= 8;
}