/*
 * 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 <assert.h>
#include <math.h>
#include <string.h>

#include "config/aom_scale_rtcd.h"

#include "aom/aom_integer.h"
#include "aom_dsp/binary_codes_writer.h"
#include "av1/common/cnn_tflite.h"
#include "av1/common/guided_quadtree.h"
#include "av1/common/reconinter.h"
#include "av1/encoder/cost.h"

// TODO(urvang@google.com): replace quantSet with struct.
// guided conv unet intra
int qp255_quadtree_model_quantSet_intra[] = { 25, 197, 0, -8 };
int qp205_quadtree_model_quantSet_intra[] = { 643189, 690747, -17, -5 };
int qp175_quadtree_model_quantSet_intra[] = { 1838, 2153, 3, -12 };
int qp145_quadtree_model_quantSet_intra[] = { 27230, 33505, -21, 3 };
int qp120_quadtree_model_quantSet_intra[] = { 133, 199, 1, -7 };
int qp90_quadtree_model_quantSet_intra[] = { 988, 1428, 0, -11 };

// guided conv unet with attention inter
int qp255_quadtree_model_quantSet_inter[] = { 56619, 15796, -9, -13 };
int qp205_quadtree_model_quantSet_inter[] = { 551342, 949167, -14, -5 };
int qp175_quadtree_model_quantSet_inter[] = { 2011, 3876, 0, -14 };
int qp145_quadtree_model_quantSet_inter[] = { 32668, 44115, -18, 2 };
int qp120_quadtree_model_quantSet_inter[] = { 20817, 19072, -12, -12 };
int qp90_quadtree_model_quantSet_inter[] = { 3455, 16494, -16, -8 };

#if CONFIG_EXT_SUPERRES
// Superres guided conv unet intra.
int sr2by1ai_1_quantset[] = { 17235, 16969, -10, -10 };
int sr2by1ai_2_quantset[] = { 76705024, 109138704, -10, -7 };
int sr2by1ai_3_quantset[] = { 335, 464, -4, -8 };

int sr3by2ai_1_quantset[] = { 385035, 450967, -8, -10 };
int sr3by2ai_2_quantset[] = { 54186, 53143, 2, -13 };
int sr3by2ai_3_quantset[] = { 608, 795, -5, -8 };

int sr5by4ai_1_quantset[] = { 57611, 154741, -2, 2 };
int sr5by4ai_2_quantset[] = { 166, 503, -12, -1 };
int sr5by4ai_3_quantset[] = { 11, 4, -11, -9 };

int sr7by4ai_1_quantset[] = { 2631403, 4529410, -8, -8 };
int sr7by4ai_2_quantset[] = { 28290, 12216, 1, -3 };
int sr7by4ai_3_quantset[] = { 11, 9, -7, -8 };

// Superres guided conv unet inter.
int sr2by1ra_1_quantset[] = { 15680, 12890, -10, -9 };
int sr2by1ra_2_quantset[] = { 85696504, 103679088, -8, -9 };
int sr2by1ra_3_quantset[] = { 2067139, 493382, -6, -6 };

int sr3by2ra_1_quantset[] = { 651704, 364241, -9, -9 };
int sr3by2ra_2_quantset[] = { 89884, 88521, 3, -12 };
int sr3by2ra_3_quantset[] = { 3133894, 9945744, -7, -11 };

int sr5by4ra_1_quantset[] = { 40478, 106889, -3, 3 };
int sr5by4ra_2_quantset[] = { 424139, 223163, 6, 0 };
int sr5by4ra_3_quantset[] = { 164313, 57994, -5, -7 };

int sr7by4ra_1_quantset[] = { 8734122, 6400229, -6, -7 };
int sr7by4ra_2_quantset[] = { 295, 36, -3, -9 };
int sr7by4ra_3_quantset[] = { 21017, 33976, -13, -6 };
#endif  // CONFIG_EXT_SUPERRES

int *get_quadparm_from_qindex(int qindex, int superres_denom, int is_intra_only,
                              int is_luma, int cnn_index) {
#if CONFIG_EXT_SUPERRES
  assert(superres_denom == SCALE_NUMERATOR || superres_denom == 10 ||
         superres_denom == 12 || superres_denom == 14 || superres_denom == 16);
#else
  assert(superres_denom == SCALE_NUMERATOR);
#endif                                      // CONFIG_EXT_SUPERRES
  if (superres_denom == SCALE_NUMERATOR) {  // quadtree
    if (is_luma) {
      if (is_intra_only) {
        if (qindex <= 90) {
          return (cnn_index == 0)   ? qp90_quadtree_model_quantSet_intra
                 : (cnn_index == 1) ? qp120_quadtree_model_quantSet_intra
                                    : qp145_quadtree_model_quantSet_intra;

        } else if (qindex <= 120) {
          return (cnn_index == 0)   ? qp120_quadtree_model_quantSet_intra
                 : (cnn_index == 1) ? qp90_quadtree_model_quantSet_intra
                                    : qp145_quadtree_model_quantSet_intra;
        } else if (qindex <= 145) {
          return (cnn_index == 0)   ? qp145_quadtree_model_quantSet_intra
                 : (cnn_index == 1) ? qp120_quadtree_model_quantSet_intra
                                    : qp175_quadtree_model_quantSet_intra;
        } else if (qindex <= 175) {
          return (cnn_index == 0)   ? qp175_quadtree_model_quantSet_intra
                 : (cnn_index == 1) ? qp145_quadtree_model_quantSet_intra
                                    : qp205_quadtree_model_quantSet_intra;
        } else if (qindex <= 205) {
          return (cnn_index == 0)   ? qp205_quadtree_model_quantSet_intra
                 : (cnn_index == 1) ? qp175_quadtree_model_quantSet_intra
                                    : qp255_quadtree_model_quantSet_intra;
        } else {
          return (cnn_index == 0)   ? qp255_quadtree_model_quantSet_intra
                 : (cnn_index == 1) ? qp205_quadtree_model_quantSet_intra
                                    : qp175_quadtree_model_quantSet_intra;
        }
      } else {
        if (qindex <= 90) {
          return (cnn_index == 0)   ? qp90_quadtree_model_quantSet_inter
                 : (cnn_index == 1) ? qp120_quadtree_model_quantSet_inter
                                    : qp145_quadtree_model_quantSet_inter;

        } else if (qindex <= 120) {
          return (cnn_index == 0)   ? qp120_quadtree_model_quantSet_inter
                 : (cnn_index == 1) ? qp90_quadtree_model_quantSet_inter
                                    : qp145_quadtree_model_quantSet_inter;
        } else if (qindex <= 145) {
          return (cnn_index == 0)   ? qp145_quadtree_model_quantSet_inter
                 : (cnn_index == 1) ? qp120_quadtree_model_quantSet_inter
                                    : qp175_quadtree_model_quantSet_inter;
        } else if (qindex <= 175) {
          return (cnn_index == 0)   ? qp175_quadtree_model_quantSet_inter
                 : (cnn_index == 1) ? qp145_quadtree_model_quantSet_inter
                                    : qp205_quadtree_model_quantSet_inter;
        } else if (qindex <= 205) {
          return (cnn_index == 0)   ? qp205_quadtree_model_quantSet_inter
                 : (cnn_index == 1) ? qp175_quadtree_model_quantSet_inter
                                    : qp255_quadtree_model_quantSet_inter;
        } else {
          return (cnn_index == 0)   ? qp255_quadtree_model_quantSet_inter
                 : (cnn_index == 1) ? qp205_quadtree_model_quantSet_inter
                                    : qp175_quadtree_model_quantSet_inter;
        }
      }
    }
  }
#if CONFIG_EXT_SUPERRES
  assert(is_luma);
  if (is_intra_only) {
#if SELECT_CNN_FOR_SUPERRES
    switch (superres_denom) {
      case 10:
        return (cnn_index == 0)   ? sr5by4ai_1_quantset
               : (cnn_index == 1) ? sr5by4ai_2_quantset
                                  : sr5by4ai_3_quantset;
      case 12:
        return (cnn_index == 0)   ? sr3by2ai_1_quantset
               : (cnn_index == 1) ? sr3by2ai_2_quantset
                                  : sr3by2ai_3_quantset;
      case 14:
        return (cnn_index == 0)   ? sr7by4ai_1_quantset
               : (cnn_index == 1) ? sr7by4ai_2_quantset
                                  : sr7by4ai_3_quantset;
      case 16:
        return (cnn_index == 0)   ? sr2by1ai_1_quantset
               : (cnn_index == 1) ? sr2by1ai_2_quantset
                                  : sr2by1ai_3_quantset;
      default: assert(0); return NULL;
    }
#else   // SELECT_CNN_FOR_SUPERRES
    switch (superres_denom) {
      case 10:
        if (qindex < 120)
          return sr5by4ai_1_quantset;
        else if (qindex < 180)
          return sr5by4ai_2_quantset;
        else
          return sr5by4ai_3_quantset;
      case 12:
        if (qindex < 120)
          return sr3by2ai_1_quantset;
        else if (qindex < 180)
          return sr3by2ai_2_quantset;
        else
          return sr3by2ai_3_quantset;
      case 14:
        if (qindex < 120)
          return sr7by4ai_1_quantset;
        else if (qindex < 180)
          return sr7by4ai_2_quantset;
        else
          return sr7by4ai_3_quantset;
      case 16:
        if (qindex < 120)
          return sr2by1ai_1_quantset;
        else if (qindex < 180)
          return sr2by1ai_2_quantset;
        else
          return sr2by1ai_3_quantset;
      default: assert(0); return NULL;
    }
#endif  // SELECT_CNN_FOR_SUPERRES
  } else {
#if SELECT_CNN_FOR_SUPERRES
    switch (superres_denom) {
      case 10:
        return (cnn_index == 0)   ? sr5by4ra_1_quantset
               : (cnn_index == 1) ? sr5by4ra_2_quantset
                                  : sr5by4ra_3_quantset;
      case 12:
        return (cnn_index == 0)   ? sr3by2ra_1_quantset
               : (cnn_index == 1) ? sr3by2ra_2_quantset
                                  : sr3by2ra_3_quantset;
      case 14:
        return (cnn_index == 0)   ? sr7by4ra_1_quantset
               : (cnn_index == 1) ? sr7by4ra_2_quantset
                                  : sr7by4ra_3_quantset;
      case 16:
        return (cnn_index == 0)   ? sr2by1ra_1_quantset
               : (cnn_index == 1) ? sr2by1ra_2_quantset
                                  : sr2by1ra_3_quantset;
      default: assert(0); return NULL;
    }
#else   // SELECT_CNN_FOR_SUPERRES
    switch (superres_denom) {
      case 10:
        if (qindex < 120)
          return sr5by4ra_1_quantset;
        else if (qindex < 180)
          return sr5by4ra_2_quantset;
        else
          return sr5by4ra_3_quantset;
      case 12:
        if (qindex < 120)
          return sr3by2ra_1_quantset;
        else if (qindex < 180)
          return sr3by2ra_2_quantset;
        else
          return sr3by2ra_3_quantset;
      case 14:
        if (qindex < 120)
          return sr7by4ra_1_quantset;
        else if (qindex < 180)
          return sr7by4ra_2_quantset;
        else
          return sr7by4ra_3_quantset;
      case 16:
        if (qindex < 120)
          return sr2by1ra_1_quantset;
        else if (qindex < 180)
          return sr2by1ra_2_quantset;
        else
          return sr2by1ra_3_quantset;
      default: assert(0); return NULL;
    }
#endif  // SELECT_CNN_FOR_SUPERRES
  }
#endif  // CONFIG_EXT_SUPERRES
  return NULL;
}

#if CONFIG_CNN_GUIDED_QUADTREE
void quad_copy(const QUADInfo *src, QUADInfo *dst, struct AV1Common *cm) {
  dst->unit_index = src->unit_index;
  dst->unit_size = src->unit_size;
  dst->split_info_length = src->split_info_length;
  dst->unit_info_length = src->unit_info_length;
  av1_alloc_quadtree_struct(cm, dst);
  for (int i = 0; i < dst->split_info_length; ++i) {
    dst->split_info[i].split = src->split_info[i].split;
  }
  for (int i = 0; i < dst->unit_info_length; ++i) {
    dst->unit_info[i].xqd[0] = src->unit_info[i].xqd[0];
    dst->unit_info[i].xqd[1] = src->unit_info[i].xqd[1];
  }
  dst->signaled = src->signaled;
}

int quad_tree_get_max_unit_info_length(int width, int height, int unit_length) {
  int unit_info_length = 0;
  const int ext_size = unit_length * 3 / 2;
  for (int row = 0; row < height;) {
    const int remaining_height = height - row;
    const int this_unit_height =
        (remaining_height < ext_size) ? remaining_height : unit_length;
    for (int col = 0; col < width;) {
      const int remaining_width = width - col;
      const int this_unit_width =
          (remaining_width < ext_size) ? remaining_width : unit_length;
      // Check for special cases near boundary.
      const bool is_horz_partitioning_allowed =
          (this_unit_height >= unit_length);
      const bool is_vert_partitioning_allowed =
          (this_unit_width >= unit_length);
      if (!is_horz_partitioning_allowed && !is_vert_partitioning_allowed) {
        // Implicitly no split, so single unit info will be signaled.
        ++unit_info_length;
      } else {
        // Assume maximum possible sub-units.
        const int max_sub_units =
            is_horz_partitioning_allowed && is_vert_partitioning_allowed ? 4
                                                                         : 2;
        unit_info_length += max_sub_units;
      }
      col += this_unit_width;
    }
    row += this_unit_height;
  }
  return unit_info_length;
}

int quad_tree_get_split_info_length(int width, int height, int unit_length) {
  int split_info_len = 0;
  const int ext_size = unit_length * 3 / 2;
  for (int row = 0; row < height;) {
    const int remaining_height = height - row;
    const int this_unit_height =
        (remaining_height < ext_size) ? remaining_height : unit_length;
    for (int col = 0; col < width;) {
      const int remaining_width = width - col;
      const int this_unit_width =
          (remaining_width < ext_size) ? remaining_width : unit_length;
      // Check for special cases near boundary.
      const bool is_horz_partitioning_allowed =
          (this_unit_height >= unit_length);
      const bool is_vert_partitioning_allowed =
          (this_unit_width >= unit_length);
      if (is_horz_partitioning_allowed || is_vert_partitioning_allowed) {
        ++split_info_len;
      }
      col += this_unit_width;
    }
    row += this_unit_height;
  }
  return split_info_len;
}

void av1_alloc_quadtree_struct(struct AV1Common *cm, QUADInfo *quad_info) {
  if (quad_info->unit_info != NULL) {
    aom_free(quad_info->unit_info);
    quad_info->unit_info = NULL;
  }
  if (quad_info->split_info != NULL) {
    aom_free(quad_info->split_info);
    quad_info->split_info = NULL;
  }
  quad_info->unit_size = quad_tree_get_unit_size(cm->superres_upscaled_width,
                                                 cm->superres_upscaled_height,
                                                 quad_info->unit_index);

  if (quad_info->split_info_length > 0) {
    CHECK_MEM_ERROR(
        cm, quad_info->split_info,
        (QUADSplitInfo *)aom_memalign(
            16, sizeof(*quad_info->split_info) * quad_info->split_info_length));
  }

  assert(quad_info->unit_info_length > 0);

  CHECK_MEM_ERROR(
      cm, quad_info->unit_info,
      (QUADUnitInfo *)aom_memalign(
          16, sizeof(*quad_info->unit_info) * quad_info->unit_info_length));

  quad_info->signaled = 0;
}

void av1_free_quadtree_struct(QUADInfo *quad_info) {
  if (quad_info->unit_info != NULL) {
    aom_free(quad_info->unit_info);
    quad_info->unit_info = NULL;
  }
  if (quad_info->split_info != NULL) {
    aom_free(quad_info->split_info);
    quad_info->split_info = NULL;
  }
  memset(quad_info, 0, sizeof(*quad_info));
}
#endif  // CONFIG_CNN_GUIDED_QUADTREE