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

#include <vector>

#include "aom_dsp/binary_codes_writer.h"
#include "av1/common/av1_common_int.h"
#include "av1/common/cnn_tflite.h"
#include "av1/tflite_models/op_registrations.h"
#include "av1/tflite_models/intra_frame_model/uv_qp0_90.h"
#include "av1/tflite_models/intra_frame_model/uv_qp91_120.h"
#include "av1/tflite_models/intra_frame_model/uv_qp121_145.h"
#include "av1/tflite_models/intra_frame_model/uv_qp146_175.h"
#include "av1/tflite_models/intra_frame_model/uv_qp176_205.h"
#include "av1/tflite_models/intra_frame_model/uv_qp206_255.h"
#include "av1/tflite_models/intra_frame_model/qp0_90.h"
#include "av1/tflite_models/intra_frame_model/qp91_120.h"
#include "av1/tflite_models/intra_frame_model/qp121_145.h"
#include "av1/tflite_models/intra_frame_model/qp146_175.h"
#include "av1/tflite_models/intra_frame_model/qp176_205.h"
#include "av1/tflite_models/intra_frame_model/qp206_255.h"
#include "av1/tflite_models/inter_frame_model/uv_qp0_90.h"
#include "av1/tflite_models/inter_frame_model/uv_qp91_120.h"
#include "av1/tflite_models/inter_frame_model/uv_qp121_145.h"
#include "av1/tflite_models/inter_frame_model/uv_qp146_175.h"
#include "av1/tflite_models/inter_frame_model/uv_qp176_205.h"
#include "av1/tflite_models/inter_frame_model/uv_qp206_255.h"
#include "av1/tflite_models/inter_frame_model/qp0_90.h"
#include "av1/tflite_models/inter_frame_model/qp91_120.h"
#include "av1/tflite_models/inter_frame_model/qp121_145.h"
#include "av1/tflite_models/inter_frame_model/qp146_175.h"
#include "av1/tflite_models/inter_frame_model/qp176_205.h"
#include "av1/tflite_models/inter_frame_model/qp206_255.h"
#if CONFIG_EXT_SUPERRES
#include "av1/tflite_models/inter_frame_model/sr5by4ra_1_tflite.h"
#include "av1/tflite_models/inter_frame_model/sr5by4ra_2_tflite.h"
#include "av1/tflite_models/inter_frame_model/sr5by4ra_3_tflite.h"
#include "av1/tflite_models/inter_frame_model/sr3by2ra_1_tflite.h"
#include "av1/tflite_models/inter_frame_model/sr3by2ra_2_tflite.h"
#include "av1/tflite_models/inter_frame_model/sr3by2ra_3_tflite.h"
#include "av1/tflite_models/inter_frame_model/sr7by4ra_1_tflite.h"
#include "av1/tflite_models/inter_frame_model/sr7by4ra_2_tflite.h"
#include "av1/tflite_models/inter_frame_model/sr7by4ra_3_tflite.h"
#include "av1/tflite_models/inter_frame_model/sr2by1ra_1_tflite.h"
#include "av1/tflite_models/inter_frame_model/sr2by1ra_2_tflite.h"
#include "av1/tflite_models/inter_frame_model/sr2by1ra_3_tflite.h"
#include "av1/tflite_models/intra_frame_model/sr5by4ai_1_tflite.h"
#include "av1/tflite_models/intra_frame_model/sr5by4ai_2_tflite.h"
#include "av1/tflite_models/intra_frame_model/sr5by4ai_3_tflite.h"
#include "av1/tflite_models/intra_frame_model/sr3by2ai_1_tflite.h"
#include "av1/tflite_models/intra_frame_model/sr3by2ai_2_tflite.h"
#include "av1/tflite_models/intra_frame_model/sr3by2ai_3_tflite.h"
#include "av1/tflite_models/intra_frame_model/sr7by4ai_1_tflite.h"
#include "av1/tflite_models/intra_frame_model/sr7by4ai_2_tflite.h"
#include "av1/tflite_models/intra_frame_model/sr7by4ai_3_tflite.h"
#include "av1/tflite_models/intra_frame_model/sr2by1ai_1_tflite.h"
#include "av1/tflite_models/intra_frame_model/sr2by1ai_2_tflite.h"
#include "av1/tflite_models/intra_frame_model/sr2by1ai_3_tflite.h"
#endif  // CONFIG_EXT_SUPERRES

#if CONFIG_CNN_GUIDED_QUADTREE
#include "av1/tflite_models/inter_frame_model/qp0_90_quadtree.h"
#include "av1/tflite_models/inter_frame_model/qp91_120_quadtree.h"
#include "av1/tflite_models/inter_frame_model/qp121_145_quadtree.h"
#include "av1/tflite_models/inter_frame_model/qp146_175_quadtree.h"
#include "av1/tflite_models/inter_frame_model/qp176_205_quadtree.h"
#include "av1/tflite_models/inter_frame_model/qp206_255_quadtree.h"
#include "av1/tflite_models/intra_frame_model/qp0_90_quadtree.h"
#include "av1/tflite_models/intra_frame_model/qp91_120_quadtree.h"
#include "av1/tflite_models/intra_frame_model/qp121_145_quadtree.h"
#include "av1/tflite_models/intra_frame_model/qp146_175_quadtree.h"
#include "av1/tflite_models/intra_frame_model/qp176_205_quadtree.h"
#include "av1/tflite_models/intra_frame_model/qp206_255_quadtree.h"
#endif

#include "common/tf_lite_includes.h"

#if CONFIG_CNN_RESTORATION

#define USE_XNNPACK 0

// Returns the TF-lite model based on the qindex.
static const unsigned char *get_intra_model_from_qindex(int qindex,
                                                        int superres_denom,
                                                        int is_luma,
                                                        int cnn_index) {
  if (qindex <= MIN_CNN_Q_INDEX) {
    assert(0);
    return nullptr;
  }
#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 CONFIG_CNN_GUIDED_QUADTREE
  if (superres_denom == SCALE_NUMERATOR) {  // quadtree
    if (is_luma) {
      if (qindex <= 90) {
        return (cnn_index == 0)   ? qp0_90_quadtree_model_tflite_data
               : (cnn_index == 1) ? qp91_120_quadtree_model_tflite_data
                                  : qp121_145_quadtree_model_tflite_data;
      } else if (qindex <= 120) {
        return (cnn_index == 0)   ? qp91_120_quadtree_model_tflite_data
               : (cnn_index == 1) ? qp0_90_quadtree_model_tflite_data
                                  : qp121_145_quadtree_model_tflite_data;
      } else if (qindex <= 145) {
        return (cnn_index == 0)   ? qp121_145_quadtree_model_tflite_data
               : (cnn_index == 1) ? qp91_120_quadtree_model_tflite_data
                                  : qp146_175_quadtree_model_tflite_data;
      } else if (qindex <= 175) {
        return (cnn_index == 0)   ? qp146_175_quadtree_model_tflite_data
               : (cnn_index == 1) ? qp121_145_quadtree_model_tflite_data
                                  : qp176_205_quadtree_model_tflite_data;
      } else if (qindex <= 205) {
        return (cnn_index == 0)   ? qp176_205_quadtree_model_tflite_data
               : (cnn_index == 1) ? qp146_175_quadtree_model_tflite_data
                                  : qp206_255_quadtree_model_tflite_data;
      } else {
        return (cnn_index == 0)   ? qp206_255_quadtree_model_tflite_data
               : (cnn_index == 1) ? qp176_205_quadtree_model_tflite_data
                                  : qp146_175_quadtree_model_tflite_data;
      }
    }
  }
#endif  // CONFIG_CNN_GUIDED_QUADTREE

  if (superres_denom == SCALE_NUMERATOR) {
    if (is_luma) {
      if (qindex < 91) {
        return (cnn_index == 0)   ? qp0_90_model_tflite_data
               : (cnn_index == 1) ? qp91_120_model_tflite_data
                                  : qp121_145_model_tflite_data;
      } else if (qindex < 121) {
        return (cnn_index == 0)   ? qp91_120_model_tflite_data
               : (cnn_index == 1) ? qp0_90_model_tflite_data
                                  : qp121_145_model_tflite_data;
      } else if (qindex < 146) {
        return (cnn_index == 0)   ? qp121_145_model_tflite_data
               : (cnn_index == 1) ? qp91_120_model_tflite_data
                                  : qp146_175_model_tflite_data;
      } else if (qindex < 176) {
        return (cnn_index == 0)   ? qp146_175_model_tflite_data
               : (cnn_index == 1) ? qp121_145_model_tflite_data
                                  : qp176_205_model_tflite_data;
      } else if (qindex < 206) {
        return (cnn_index == 0)   ? qp176_205_model_tflite_data
               : (cnn_index == 1) ? qp146_175_model_tflite_data
                                  : qp206_255_model_tflite_data;
      } else {
        return (cnn_index == 0)   ? qp206_255_model_tflite_data
               : (cnn_index == 1) ? qp176_205_model_tflite_data
                                  : qp146_175_model_tflite_data;
      }
    } else {
      assert(cnn_index == 0);
      if (qindex < 91) {
        return uv_qp0_90_model_tflite_data;
      } else if (qindex < 121) {
        return uv_qp91_120_model_tflite_data;
      } else if (qindex < 146) {
        return uv_qp121_145_model_tflite_data;
      } else if (qindex < 176) {
        return uv_qp146_175_model_tflite_data;
      } else if (qindex < 206) {
        return uv_qp176_205_model_tflite_data;
      } else {
        return uv_qp206_255_model_tflite_data;
      }
    }
  }
#if CONFIG_EXT_SUPERRES
  assert(is_luma);
#if SELECT_CNN_FOR_SUPERRES
  switch (superres_denom) {
    case 10:
      return (cnn_index == 0)   ? sr5by4ai_1_tflite
             : (cnn_index == 1) ? sr5by4ai_2_tflite
                                : sr5by4ai_3_tflite;
    case 12:
      return (cnn_index == 0)   ? sr3by2ai_1_tflite
             : (cnn_index == 1) ? sr3by2ai_2_tflite
                                : sr3by2ai_3_tflite;
    case 14:
      return (cnn_index == 0)   ? sr7by4ai_1_tflite
             : (cnn_index == 1) ? sr7by4ai_2_tflite
                                : sr7by4ai_3_tflite;
    case 16:
      return (cnn_index == 0)   ? sr2by1ai_1_tflite
             : (cnn_index == 1) ? sr2by1ai_2_tflite
                                : sr2by1ai_3_tflite;
    default: assert(0); return nullptr;
  }
#else   // SELECT_CNN_FOR_SUPERRES
  switch (superres_denom) {
    case 10:
      if (qindex < 120)
        return sr5by4ai_1_tflite;
      else if (qindex < 180)
        return sr5by4ai_2_tflite;
      else
        return sr5by4ai_3_tflite;
    case 12:
      if (qindex < 120)
        return sr3by2ai_1_tflite;
      else if (qindex < 180)
        return sr3by2ai_2_tflite;
      else
        return sr3by2ai_3_tflite;
    case 14:
      if (qindex < 120)
        return sr7by4ai_1_tflite;
      else if (qindex < 180)
        return sr7by4ai_2_tflite;
      else
        return sr7by4ai_3_tflite;
    case 16:
      if (qindex < 120)
        return sr2by1ai_1_tflite;
      else if (qindex < 180)
        return sr2by1ai_2_tflite;
      else
        return sr2by1ai_3_tflite;
    default: assert(0); return nullptr;
  }
#endif  // SELECT_CNN_FOR_SUPERRES
#endif  // CONFIG_EXT_SUPERRES
  return nullptr;
}

// Returns the TF-lite model based on the qindex.
static const unsigned char *get_inter_model_from_qindex(int qindex,
                                                        int superres_denom,
                                                        int is_luma,
                                                        int cnn_index) {
  if (qindex <= MIN_CNN_Q_INDEX) {
    assert(0);
    return nullptr;
  }
#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 CONFIG_CNN_GUIDED_QUADTREE
  if (superres_denom == SCALE_NUMERATOR) {  // quadtree
    if (is_luma) {
      if (qindex <= 90) {
        return (cnn_index == 0)   ? qp0_90_quadtree_inter_model_tflite_data
               : (cnn_index == 1) ? qp91_120_quadtree_inter_model_tflite_data
                                  : qp121_145_quadtree_inter_model_tflite_data;
      } else if (qindex <= 120) {
        return (cnn_index == 0)   ? qp91_120_quadtree_inter_model_tflite_data
               : (cnn_index == 1) ? qp0_90_quadtree_inter_model_tflite_data
                                  : qp121_145_quadtree_inter_model_tflite_data;
      } else if (qindex <= 145) {
        return (cnn_index == 0)   ? qp121_145_quadtree_inter_model_tflite_data
               : (cnn_index == 1) ? qp91_120_quadtree_inter_model_tflite_data
                                  : qp146_175_quadtree_inter_model_tflite_data;
      } else if (qindex <= 175) {
        return (cnn_index == 0)   ? qp146_175_quadtree_inter_model_tflite_data
               : (cnn_index == 1) ? qp121_145_quadtree_inter_model_tflite_data
                                  : qp176_205_quadtree_inter_model_tflite_data;
      } else if (qindex <= 205) {
        return (cnn_index == 0)   ? qp176_205_quadtree_inter_model_tflite_data
               : (cnn_index == 1) ? qp146_175_quadtree_inter_model_tflite_data
                                  : qp206_255_quadtree_inter_model_tflite_data;
      } else {
        return (cnn_index == 0)   ? qp206_255_quadtree_inter_model_tflite_data
               : (cnn_index == 1) ? qp176_205_quadtree_inter_model_tflite_data
                                  : qp146_175_quadtree_inter_model_tflite_data;
      }
    }
  }
#endif  // CONFIG_CNN_GUIDED_QUADTREE

  if (superres_denom == SCALE_NUMERATOR) {
    if (is_luma) {
      if (qindex < 91) {
        return (cnn_index == 0)   ? qp0_90_inter_model_tflite_data
               : (cnn_index == 1) ? qp91_120_inter_model_tflite_data
                                  : qp121_145_inter_model_tflite_data;
      } else if (qindex < 121) {
        return (cnn_index == 0)   ? qp91_120_inter_model_tflite_data
               : (cnn_index == 1) ? qp0_90_inter_model_tflite_data
                                  : qp121_145_inter_model_tflite_data;
      } else if (qindex < 146) {
        return (cnn_index == 0)   ? qp121_145_inter_model_tflite_data
               : (cnn_index == 1) ? qp91_120_inter_model_tflite_data
                                  : qp146_175_inter_model_tflite_data;
      } else if (qindex < 176) {
        return (cnn_index == 0)   ? qp146_175_inter_model_tflite_data
               : (cnn_index == 1) ? qp121_145_inter_model_tflite_data
                                  : qp176_205_inter_model_tflite_data;
      } else if (qindex < 206) {
        return (cnn_index == 0)   ? qp176_205_inter_model_tflite_data
               : (cnn_index == 1) ? qp146_175_inter_model_tflite_data
                                  : qp206_255_inter_model_tflite_data;
      } else {
        return (cnn_index == 0)   ? qp206_255_inter_model_tflite_data
               : (cnn_index == 1) ? qp176_205_inter_model_tflite_data
                                  : qp146_175_inter_model_tflite_data;
      }
    } else {
      assert(cnn_index == 0);
      if (qindex < 91) {
        return uv_qp0_90_inter_model_tflite_data;
      } else if (qindex < 121) {
        return uv_qp91_120_inter_model_tflite_data;
      } else if (qindex < 146) {
        return uv_qp121_145_inter_model_tflite_data;
      } else if (qindex < 176) {
        return uv_qp146_175_inter_model_tflite_data;
      } else if (qindex < 206) {
        return uv_qp176_205_inter_model_tflite_data;
      } else {
        return uv_qp206_255_inter_model_tflite_data;
      }
    }
  }

#if CONFIG_EXT_SUPERRES
  assert(is_luma);
  switch (superres_denom) {
    case 10:
      if (qindex < 120)
        return sr5by4ra_1_tflite;
      else if (qindex < 180)
        return sr5by4ra_2_tflite;
      else
        return sr5by4ra_3_tflite;
    case 12:
      if (qindex < 120)
        return sr3by2ra_1_tflite;
      else if (qindex < 180)
        return sr3by2ra_2_tflite;
      else
        return sr3by2ra_3_tflite;
    case 14:
      if (qindex < 120)
        return sr7by4ra_1_tflite;
      else if (qindex < 180)
        return sr7by4ra_2_tflite;
      else
        return sr7by4ra_3_tflite;
    case 16:
      if (qindex < 120)
        return sr2by1ra_1_tflite;
      else if (qindex < 180)
        return sr2by1ra_2_tflite;
      else
        return sr2by1ra_3_tflite;
    default: assert(0); return nullptr;
  }
#endif  // CONFIG_EXT_SUPERRES
  return nullptr;
}

#if USE_XNNPACK
static TfLiteDelegate *get_tflite_xnnpack_delegate(int num_threads) {
  TfLiteXNNPackDelegateOptions xnnpack_options =
      TfLiteXNNPackDelegateOptionsDefault();
  xnnpack_options.num_threads = AOMMAX(num_threads, 1);
  return TfLiteXNNPackDelegateCreate(&xnnpack_options);
}
#endif  // USE_XNNPACK

// Builds and returns the TFlite interpreter.
static std::unique_ptr<tflite::Interpreter> get_tflite_interpreter(
    int qindex, int superres_denom, int width, int height, int num_threads,
    int is_intra_only, int is_luma, int cnn_index
#if USE_XNNPACK
    ,
    TfLiteDelegate *xnnpack_delegate
#endif  // USE_XNNPACK
) {
  const unsigned char *const model_tflite_data =
      is_intra_only ? get_intra_model_from_qindex(qindex, superres_denom,
                                                  is_luma, cnn_index)
                    : get_inter_model_from_qindex(qindex, superres_denom,
                                                  is_luma, cnn_index);
  auto model = tflite::GetModel(model_tflite_data);
  tflite::MutableOpResolver resolver;
  RegisterSelectedOpsAllQps(&resolver);
  tflite::InterpreterBuilder builder(model, resolver);
  // TODO(urvang): Investigate if caching the interpreter object provides
  // further speed-up. May still have to re-build the interpreter if qindex
  // changes.
  std::unique_ptr<tflite::Interpreter> interpreter;
  builder(&interpreter);
  interpreter->SetNumThreads(AOMMAX(num_threads, 1));
  tflite::ErrorReporter *reporter = tflite::DefaultErrorReporter();

  // Dimension order: batch_size, height, width, num_channels.
  // Note: height comes before width here!
  const std::vector<int> in_out_dims = { 1, height, width, 1 };
  // We only need to resize the input tensor. All other tensors (including
  // output tensor) will be resized automatically.
  if (interpreter->ResizeInputTensor(interpreter->inputs()[0], in_out_dims) !=
      kTfLiteOk) {
    reporter->Report("Failed at input tensor resize");
    return nullptr;
  }

  if (interpreter->AllocateTensors() != kTfLiteOk) {
    reporter->Report("Failed at tensor allocation");
    return nullptr;
  }

#if USE_XNNPACK
  if (interpreter->ModifyGraphWithDelegate(xnnpack_delegate) != kTfLiteOk) {
    reporter->Report("Failed at modifying graph with XNNPack delegate");
    return nullptr;
  }
#endif  // USE_XNNPACK

  return interpreter;
}

extern "C" int av1_restore_cnn_img_tflite_highbd(
    int qindex, int superres_denom, const uint16_t *dgd, int width, int height,
    int dgd_stride, uint16_t *rst, int rst_stride, int num_threads,
    int bit_depth, int is_intra_only, int is_luma, int cnn_index) {
  // Ensure image can be downscaled by factor of 8 on each axis
  int padding_width = int(ceil(float(width) / 8.0) * 8);
  int padding_height = int(ceil(float(height) / 8.0) * 8);
#if USE_XNNPACK
  TfLiteDelegate *xnnpack_delegate = get_tflite_xnnpack_delegate(num_threads);
#endif  // USE_XNNPACK
  std::unique_ptr<tflite::Interpreter> interpreter = get_tflite_interpreter(
      qindex, superres_denom, padding_width, padding_height, num_threads,
      is_intra_only, is_luma, cnn_index
#if USE_XNNPACK
      ,
      xnnpack_delegate
#endif  // USE_XNNPACK
  );

  // Prepare input.
  const auto max_val = static_cast<float>((1 << bit_depth) - 1);
  const int in_stride = padding_width;
  auto input = interpreter->typed_input_tensor<float>(0);
  for (int r = 0; r < padding_height; ++r) {
    for (int c = 0; c < padding_width; ++c) {
      if (r < height && c < width) {
        input[r * in_stride + c] =
            static_cast<float>(dgd[r * dgd_stride + c]) / max_val;
        assert(input[r * in_stride + c] >= 0.0f);
        assert(input[r * in_stride + c] <= 1.0f);
      } else {
        // Padding with either zeros or by copies
        // input[r * in_stride + c] = 0;  // Pad with zeros
        int w_copy_idx = c;
        if (c >= width) {
          w_copy_idx = width + (width - c) - 1;
        }
        int h_copy_idx = r;
        if (r >= height) {
          h_copy_idx = height + (height - r) - 1;
        }
        input[r * in_stride + c] = input[h_copy_idx * in_stride + w_copy_idx];
      }
    }
  }

  // Invoke TFlite inference.
  tflite::ErrorReporter *reporter = tflite::DefaultErrorReporter();
  auto status = interpreter->Invoke();
  if (status != kTfLiteOk) {
    reporter->Report("Failed at interpreter invocation");
    return 0;
  }

  // Use the output to restore 'dgd' and store in 'rst'.
  const auto output = interpreter->typed_output_tensor<float>(0);
  const int out_stride = width;
  for (int r = 0; r < height; ++r) {
    for (int c = 0; c < width; ++c) {
      const int residue =
          static_cast<int>(output[r * out_stride + c] * max_val + 0.5);
      rst[r * rst_stride + c] =
          clip_pixel_highbd(dgd[r * dgd_stride + c] + residue, bit_depth);
    }
  }

  interpreter.reset();

#if USE_XNNPACK
  // IMPORTANT: release the interpreter before destroying the delegate.
  TfLiteXNNPackDelegateDelete(xnnpack_delegate);
#endif  // USE_XNNPACK

  return 1;
}

extern "C" void av1_restore_cnn_tflite(const AV1_COMMON *cm, int num_threads,
                                       const int apply_cnn[MAX_MB_PLANE],
                                       const int cnn_indices[MAX_MB_PLANE]) {
  YV12_BUFFER_CONFIG *buf = &cm->cur_frame->buf;
  const int is_intra_only = frame_is_intra_only(cm);
  for (int plane = 0; plane < av1_num_planes(cm); ++plane) {
    if (!apply_cnn[plane]) continue;
    const int is_luma = (plane == AOM_PLANE_Y);
    const int cnn_index = cnn_indices[plane];
    assert(cnn_index >= 0 &&
           cnn_index < av1_num_cnn_indices_for_plane(cm, plane));
    switch (plane) {
      case AOM_PLANE_Y:
        av1_restore_cnn_img_tflite_highbd(
            cm->quant_params.base_qindex, cm->superres_scale_denominator,
            CONVERT_TO_SHORTPTR(buf->y_buffer), buf->y_crop_width,
            buf->y_crop_height, buf->y_stride,
            CONVERT_TO_SHORTPTR(buf->y_buffer), buf->y_stride, num_threads,
            cm->seq_params.bit_depth, is_intra_only, is_luma, cnn_index);
        break;
      case AOM_PLANE_U:
        av1_restore_cnn_img_tflite_highbd(
            cm->quant_params.base_qindex, cm->superres_scale_denominator,
            CONVERT_TO_SHORTPTR(buf->u_buffer), buf->uv_crop_width,
            buf->uv_crop_height, buf->uv_stride,
            CONVERT_TO_SHORTPTR(buf->u_buffer), buf->uv_stride, num_threads,
            cm->seq_params.bit_depth, is_intra_only, is_luma, cnn_index);
        break;
      case AOM_PLANE_V:
        av1_restore_cnn_img_tflite_highbd(
            cm->quant_params.base_qindex, cm->superres_scale_denominator,
            CONVERT_TO_SHORTPTR(buf->v_buffer), buf->uv_crop_width,
            buf->uv_crop_height, buf->uv_stride,
            CONVERT_TO_SHORTPTR(buf->v_buffer), buf->uv_stride, num_threads,
            cm->seq_params.bit_depth, is_intra_only, is_luma, cnn_index);
        break;
      default: assert(0 && "Invalid plane index");
    }
  }
}

#if CONFIG_CNN_GUIDED_QUADTREE

// ------------------- Guided Quadtree: Common -------------------------------//

// Given single-channel input in 'dgd', generate intermediate 2-channel CNN
// output 'interm'.
static int generate_interm_guided_restoration(
    const uint16_t *dgd, int dgd_stride, int qindex, int superres_denom,
    int width, int height, int num_threads, int is_intra_only, int is_luma,
    int cnn_index, int bit_depth,
    std::vector<std::vector<std::vector<double>>> &interm) {
  // Make sure we can downscale 4 times.
  const int padding_width = (int)ceil(width * 1.0 / 16) * 16;
  const int padding_height = (int)ceil(height * 1.0 / 16) * 16;
#if USE_XNNPACK
  TfLiteDelegate *xnnpack_delegate = get_tflite_xnnpack_delegate(num_threads);
#endif  // USE_XNNPACK
  std::unique_ptr<tflite::Interpreter> interpreter = get_tflite_interpreter(
      qindex, superres_denom, padding_width, padding_height, num_threads,
      is_intra_only, is_luma, cnn_index
#if USE_XNNPACK
      ,
      xnnpack_delegate
#endif  // USE_XNNPACK
  );

  // Prepare input.
  const auto max_val = static_cast<float>((1 << bit_depth) - 1);
  const int in_stride = padding_width;
  auto input = interpreter->typed_input_tensor<float>(0);
  for (int r = 0; r < padding_height; ++r) {
    for (int c = 0; c < padding_width; ++c) {
      if (r < height && c < width) {
        input[r * in_stride + c] =
            static_cast<float>(dgd[r * dgd_stride + c]) / max_val;
        assert(input[r * in_stride + c] >= 0.0f);
        assert(input[r * in_stride + c] <= 1.0f);
      } else {
        input[r * in_stride + c] =
            static_cast<float>(dgd[AOMMIN(r, height - 1) * dgd_stride +
                                   AOMMIN(c, width - 1)]) /
            max_val;
      }
    }
  }

  // Invoke TFlite inference.
  tflite::ErrorReporter *reporter = tflite::DefaultErrorReporter();
  auto status = interpreter->Invoke();
  if (status != kTfLiteOk) {
    reporter->Report("Failed at interpreter invocation");
    return 0;
  }

  // Store the output in 'interm'.
  const auto output = interpreter->typed_output_tensor<float>(0);
  const int out_stride = padding_width;

  for (int r = 0; r < height; ++r) {
    for (int c = 0; c < width; ++c) {
      interm[r][c][0] = output[r * 2 * out_stride + c * 2] * max_val;
      interm[r][c][1] = output[r * 2 * out_stride + c * 2 + 1] * max_val;
    }
  }

  // Cleanup.
  interpreter.reset();
#if USE_XNNPACK
  // IMPORTANT: release the interpreter before destroying the delegate.
  TfLiteXNNPackDelegateDelete(xnnpack_delegate);
#endif  // USE_XNNPACK
  return 1;
}

// Get unit width and height based on max size and partition type.
static void get_unit_size(int max_unit_width, int max_unit_height,
                          GuidedQuadTreePartitionType partition_type,
                          int *unit_width, int *unit_height) {
  assert(partition_type >= 0 && partition_type < GUIDED_QT_TYPES);
  *unit_width =
      (partition_type == GUIDED_QT_NONE || partition_type == GUIDED_QT_HORZ)
          ? max_unit_width
          : max_unit_width >> 1;
  *unit_height =
      (partition_type == GUIDED_QT_NONE || partition_type == GUIDED_QT_VERT)
          ? max_unit_height
          : max_unit_height >> 1;
}

// ------------------- Guided Quadtree: Encoder ------------------------------//

// Given 2-channel intermediate output 'interm', degraded frame 'dgd' and source
// frame 'src', generates the single-channel output 'out' and corresponding
// linear combination weight pairs 'a'.
// Assumes that `width x height` area needs to be combined using unit of size
// `unit_width x unit_height`.
static void generate_linear_combination(
    const std::vector<std::vector<std::vector<double>>> &interm,
    const uint16_t *src, int src_stride, const uint16_t *dgd, int dgd_stride,
    int start_row, int end_row, int start_col, int end_col, int unit_width,
    int unit_height, const int *quadtset, int rdmult, const int *norestorecost,
    int bit_depth, std::vector<std::vector<uint16_t>> &out,
    std::vector<std::pair<int, int>> &A) {
  const int scale0 = quadtset[0];
  const int scale1 = quadtset[1];
  const int A0_min = quadtset[2];
  const int A1_min = quadtset[3];

  for (int row = start_row; row < end_row; row += unit_height) {
    const int this_start_row = row;
    const int this_end_row = AOMMIN(row + unit_height, end_row);
    for (int col = start_col; col < end_col; col += unit_width) {
      const int this_start_col = col;
      const int this_end_col = AOMMIN(col + unit_width, end_col);
      const int num_pixels =
          (this_end_row - this_start_row) * (this_end_col - this_start_col);

      // Extract some flattened arrays.
      std::vector<int> sub_r_flatten;
      sub_r_flatten.reserve(num_pixels);
      for (int i = this_start_row; i < this_end_row; i++) {
        for (int j = this_start_col; j < this_end_col; j++) {
          sub_r_flatten.push_back(src[i * src_stride + j] -
                                  dgd[i * dgd_stride + j]);
        }
      }
      assert((int)sub_r_flatten.size() == num_pixels);

      std::vector<double> sub_r0;
      sub_r0.reserve(num_pixels);
      for (int i = this_start_row; i < this_end_row; i++) {
        for (int j = this_start_col; j < this_end_col; j++) {
          sub_r0.push_back(interm[i][j][0]);
        }
      }
      assert((int)sub_r0.size() == num_pixels);

      std::vector<double> sub_r1;
      sub_r1.reserve(num_pixels);
      for (int i = this_start_row; i < this_end_row; i++) {
        for (int j = this_start_col; j < this_end_col; j++) {
          sub_r1.push_back(interm[i][j][1]);
        }
      }
      assert((int)sub_r1.size() == num_pixels);

      // Get R.
      std::vector<std::vector<double>> R(num_pixels, std::vector<double>(2));
      for (int i = 0; i < num_pixels; i++) {
        R[i][0] = sub_r0[i];
        R[i][1] = sub_r1[i];
      }

      // Get R^T.
      std::vector<std::vector<double>> R_T(2, std::vector<double>(num_pixels));
      for (int i = 0; i < num_pixels; i++) {
        R_T[0][i] = sub_r0[i];
        R_T[1][i] = sub_r1[i];
      }

      // Get R^T * R.
      double R_TDotR[2][2] = { 0 };
      for (int i = 0; i < 2; i++) {
        for (int j = 0; j < 2; j++) {
          for (int k = 0; k < num_pixels; k++) {
            R_TDotR[i][j] += R_T[i][k] * R[k][j];
          }
        }
      }

      // Get (R^T * R)^-1.
      const double value_R_TDotR =
          R_TDotR[0][0] * R_TDotR[1][1] - R_TDotR[0][1] * R_TDotR[1][0];

      double R_TDotR_inver[2][2] = {
        { R_TDotR[1][1] / value_R_TDotR, -1 * R_TDotR[0][1] / value_R_TDotR },
        { -1 * R_TDotR[1][0] / value_R_TDotR, R_TDotR[0][0] / value_R_TDotR }
      };

      // Get (R^T * R)^-1 * R^T.
      std::vector<std::vector<double>> mid(2, std::vector<double>(num_pixels));
      for (int j = 0; j < num_pixels; j++) {
        mid[0][j] =
            R_TDotR_inver[0][0] * R_T[0][j] + R_TDotR_inver[0][1] * R_T[1][j];
        mid[1][j] =
            R_TDotR_inver[1][0] * R_T[0][j] + R_TDotR_inver[1][1] * R_T[1][j];
      }

      // Compute A = (R^T * R)^-1 * R^T * residual.
      double A0 = 0;
      double A1 = 0;
      for (int i = 0; i < num_pixels; i++) {
        A0 += mid[0][i] * sub_r_flatten[i];
        A1 += mid[1][i] * sub_r_flatten[i];
      }
      A0 = A0 * scale0;
      A1 = A1 * scale1;

      // Do a finer search for best A0, A1 pair amongst four options:
      // (1) A0_floor = floor(A0), A1_floor = floor(A1)
      // (2) A0_floor, A1_floor + 1
      // (3) A0_floor + 1, A1_floor
      // (4) A0_floor + 1, A1_floor + 1
      const bool do_finer_search = true;
      if (do_finer_search) {
        double bestA0 = 0;
        double bestA1 = 0;
        double cost;
        int64_t err = 0;
        for (int i = this_start_row; i < this_end_row; i++) {
          for (int j = this_start_col; j < this_end_col; j++) {
            const int diff = src[i * src_stride + j] - dgd[i * dgd_stride + j];
            err += diff * diff;
          }
        }
        double bestcost = RDCOST_DBL_WITH_NATIVE_BD_DIST(
            rdmult, norestorecost[1] >> 4, err, bit_depth);

        // finer search
        double flrA0 = (floor(A0));
        double flrA1 = (floor(A1));
        flrA0 = AOMMIN(AOMMAX(flrA0, A0_min), A0_min + GUIDED_A_RANGE);
        flrA1 = AOMMIN(AOMMAX(flrA1, A1_min), A1_min + GUIDED_A_RANGE);
        {
          A0 = flrA0;
          A1 = flrA1;
          err = 0;
          for (int i = this_start_row; i < this_end_row; i++) {
            for (int j = this_start_col; j < this_end_col; j++) {
              int rest = int(round(dgd[i * dgd_stride + j] +
                                   A0 * interm[i][j][0] / scale0 +
                                   A1 * interm[i][j][1] / scale1));
              rest = clip_pixel_highbd(rest, bit_depth);
              const int diff = src[i * src_stride + j] - rest;
              err += diff * diff;
            }
          }
          // approx RD cost assuming GUIDED_A_PAIR_BITS bits per a0, a1 pair
          cost = RDCOST_DBL_WITH_NATIVE_BD_DIST(
              rdmult,
              (norestorecost[0] +
               (GUIDED_A_PAIR_BITS << AV1_PROB_COST_SHIFT)) >>
                  4,
              err, bit_depth);
          if (cost < bestcost) {
            bestA0 = A0;
            bestA1 = A1;
            bestcost = cost;
          }
        }
        if (flrA0 < A0_min + GUIDED_A_RANGE) {
          A0 = flrA0 + 1;
          A1 = flrA1;
          err = 0;
          for (int i = this_start_row; i < this_end_row; i++) {
            for (int j = this_start_col; j < this_end_col; j++) {
              int rest = int(round(dgd[i * dgd_stride + j] +
                                   A0 * interm[i][j][0] / scale0 +
                                   A1 * interm[i][j][1] / scale1));
              rest = clip_pixel_highbd(rest, bit_depth);
              const int diff = src[i * src_stride + j] - rest;
              err += diff * diff;
            }
          }
          // approx RD cost assuming GUIDED_A_PAIR_BITS bits per a0, a1 pair
          cost = RDCOST_DBL_WITH_NATIVE_BD_DIST(
              rdmult,
              (norestorecost[0] +
               (GUIDED_A_PAIR_BITS << AV1_PROB_COST_SHIFT)) >>
                  4,
              err, bit_depth);
          if (cost < bestcost) {
            bestA0 = A0;
            bestA1 = A1;
            bestcost = cost;
          }
        }
        if (flrA1 < A1_min + GUIDED_A_RANGE) {
          A0 = flrA0;
          A1 = flrA1 + 1;
          err = 0;
          for (int i = this_start_row; i < this_end_row; i++) {
            for (int j = this_start_col; j < this_end_col; j++) {
              int rest = int(round(dgd[i * dgd_stride + j] +
                                   A0 * interm[i][j][0] / scale0 +
                                   A1 * interm[i][j][1] / scale1));
              rest = clip_pixel_highbd(rest, bit_depth);
              const int diff = src[i * src_stride + j] - rest;
              err += diff * diff;
            }
          }
          // approx RD cost assuming GUIDED_A_PAIR_BITS bits per a0, a1 pair
          cost = RDCOST_DBL_WITH_NATIVE_BD_DIST(
              rdmult,
              (norestorecost[0] +
               (GUIDED_A_PAIR_BITS << AV1_PROB_COST_SHIFT)) >>
                  4,
              err, bit_depth);
          if (cost < bestcost) {
            bestA0 = A0;
            bestA1 = A1;
            bestcost = cost;
          }
        }
        if (flrA0 < A0_min + GUIDED_A_RANGE &&
            flrA1 < A1_min + GUIDED_A_RANGE) {
          A0 = flrA0 + 1;
          A1 = flrA1 + 1;
          err = 0;
          for (int i = this_start_row; i < this_end_row; i++) {
            for (int j = this_start_col; j < this_end_col; j++) {
              int rest = int(round(dgd[i * dgd_stride + j] +
                                   A0 * interm[i][j][0] / scale0 +
                                   A1 * interm[i][j][1] / scale1));
              rest = clip_pixel_highbd(rest, bit_depth);
              const int diff = src[i * src_stride + j] - rest;
              err += diff * diff;
            }
          }
          // approx RD cost assuming GUIDED_A_PAIR_BITS bits per a0, a1 pair
          cost = RDCOST_DBL_WITH_NATIVE_BD_DIST(
              rdmult,
              (norestorecost[0] +
               (GUIDED_A_PAIR_BITS << AV1_PROB_COST_SHIFT)) >>
                  4,
              err, bit_depth);
          if (cost < bestcost) {
            bestA0 = A0;
            bestA1 = A1;
            bestcost = cost;
          }
        }
        A0 = bestA0;
        A1 = bestA1;
      } else {
        A0 = (round(A0));
        A1 = (round(A1));
        A0 = AOMMIN(AOMMAX(A0, A0_min), A0_min + GUIDED_A_RANGE);
        A1 = AOMMIN(AOMMAX(A1, A1_min), A1_min + GUIDED_A_RANGE);
      }

      A0 = AOMMIN(AOMMAX(A0, A0_min), A0_min + GUIDED_A_RANGE);
      A1 = AOMMIN(AOMMAX(A1, A1_min), A1_min + GUIDED_A_RANGE);
      A.emplace_back((int)A0, (int)A1);
      for (int i = this_start_row; i < this_end_row; i++) {
        for (int j = this_start_col; j < this_end_col; j++) {
          const int out_unclipped = int(round(dgd[i * dgd_stride + j] +
                                              A0 * interm[i][j][0] / scale0 +
                                              A1 * interm[i][j][1] / scale1));
          out[i - start_row][j - start_col] =
              clip_pixel_highbd(out_unclipped, bit_depth);
        }
      }
    }
  }
#ifndef NDEBUG
  const auto num_units_row =
      (size_t)ceil((double)(end_row - start_row) / unit_height);
  const auto num_units_col =
      (size_t)ceil((double)(end_col - start_col) / unit_width);
  assert(A.size() == num_units_row * num_units_col);
#endif  // NDEBUG
}

// Computes SSE between 'rst' and 'src'.
static int64_t compute_sse(const std::vector<std::vector<uint16_t>> &rst,
                           const uint16_t *src, int src_stride, int start_row,
                           int end_row, int start_col, int end_col) {
  int64_t sse = 0;
  for (int r = start_row; r < end_row; ++r) {
    for (int c = start_col; c < end_col; ++c) {
      const uint16_t this_rst = rst[r - start_row][c - start_col];
      const uint16_t this_src = src[r * src_stride + c];
      const int64_t diff = (int64_t)(this_rst - this_src);
      sse += diff * diff;
    }
  }
  return sse;
}

// Computes bitrate for the given weight parameters.
static int compute_rate(const std::vector<std::pair<int, int>> &A,
                        const std::pair<int, int> &prev_A, const int *quadtset,
                        const int *norestorecosts) {
  const int A0_min = quadtset[2];
  const int A1_min = quadtset[3];
  int num_bits = 0;
  int ref0 = AOMMIN(AOMMAX(prev_A.first - A0_min, 0), GUIDED_A_RANGE);
  int ref1 = AOMMIN(AOMMAX(prev_A.second - A1_min, 0), GUIDED_A_RANGE);
  for (auto &this_A : A) {
    if (this_A.first == 0 && this_A.second == 0) {
      num_bits += norestorecosts[1];
    } else {
      num_bits += norestorecosts[0];
      num_bits += (aom_count_primitive_refsubexpfin(
                       GUIDED_A_NUM_VALUES, 1, ref0, this_A.first - A0_min) +
                   aom_count_primitive_refsubexpfin(
                       GUIDED_A_NUM_VALUES, 1, ref1, this_A.second - A1_min))
                  << AV1_PROB_COST_SHIFT;
    }
    ref0 = AOMMIN(AOMMAX(this_A.first - A0_min, 0), GUIDED_A_RANGE);
    ref1 = AOMMIN(AOMMAX(this_A.second - A1_min, 0), GUIDED_A_RANGE);
  }
  return num_bits;
}

// Given 2-channel intermediate output in 'interm' as well as 'src' and 'dgd'
// buffers, tries the given partition type on a single quadtree unit. Outputs
// the RDCost in 'this_rdcost' and restored unit in 'out'.
static void try_one_partition(
    const std::vector<std::vector<std::vector<double>>> &interm,
    GuidedQuadTreePartitionType partition_type, const uint16_t *src,
    int src_stride, const uint16_t *dgd, int dgd_stride, int start_row,
    int end_row, int start_col, int end_col, int max_unit_width,
    int max_unit_height, const int *quadtset, int rdmult,
    const std::pair<int, int> &prev_A, const int *quad_split_costs,
    const int *binary_split_costs, const int *norestorecosts, int bit_depth,
    bool is_horz_partitioning_allowed, int is_vert_partitioning_allowed,
    double *this_rdcost, std::vector<std::vector<uint16_t>> &out,
    std::vector<std::pair<int, int>> &A) {
  assert(IMPLIES(
      !is_horz_partitioning_allowed,
      partition_type == GUIDED_QT_NONE || partition_type == GUIDED_QT_VERT));
  assert(IMPLIES(
      !is_vert_partitioning_allowed,
      partition_type == GUIDED_QT_NONE || partition_type == GUIDED_QT_HORZ));
  // Get unit width and height based on partition type.
  int unit_width;
  int unit_height;
  get_unit_size(max_unit_width, max_unit_height, partition_type, &unit_width,
                &unit_height);

  // Compute restored unit, a0 and a1.
  generate_linear_combination(interm, src, src_stride, dgd, dgd_stride,
                              start_row, end_row, start_col, end_col,
                              unit_width, unit_height, quadtset, rdmult,
                              norestorecosts, bit_depth, out, A);
  assert(IMPLIES(partition_type == GUIDED_QT_NONE, A.size() == 1));
  assert(IMPLIES(partition_type == GUIDED_QT_HORZ, A.size() == 2));
  assert(IMPLIES(partition_type == GUIDED_QT_VERT, A.size() == 2));
  assert(IMPLIES(partition_type == GUIDED_QT_SPLIT, A.size() == 4));

  // Compute SSE.
  const int64_t sse =
      compute_sse(out, src, src_stride, start_row, end_row, start_col, end_col);

  // Compute Rate.
  const int a_signaling_cost =
      compute_rate(A, prev_A, quadtset, norestorecosts);
  // Partition signaling cost depending on 1, 2 or 4 possible partition types.
  const int partition_signaling_cost =
      is_horz_partitioning_allowed && is_vert_partitioning_allowed
          ? quad_split_costs[partition_type]
      : (is_horz_partitioning_allowed || is_vert_partitioning_allowed)
          ? binary_split_costs[partition_type]
          : 0;
  const int bitrate = a_signaling_cost + partition_signaling_cost;

  // Compute RDCost.
  *this_rdcost =
      RDCOST_DBL_WITH_NATIVE_BD_DIST(rdmult, bitrate >> 4, sse, bit_depth);
}

// Given intermediate restoration 'interm', source 'src' and degradade frame
// 'dgd', computes the best partitioning out of NONE, SPLIT, HORZ and VERT based
// on RD cost for the widthxheight unit starting at 'row' and 'col'.
// The split decisions are stored in 'split' and a0,a1 pairs are stored in 'A'.
static void select_quadtree_partitioning(
    const std::vector<std::vector<std::vector<double>>> &interm,
    const uint16_t *src, int src_stride, int start_row, int start_col,
    int width, int height, int quadtree_max_size, int max_unit_width,
    int max_unit_height, const int *quadtset, int rdmult,
    const std::pair<int, int> &prev_A, const int *quad_split_costs,
    const int *binary_split_costs, const int norestorecosts[2], int bit_depth,
    const uint16_t *dgd, int dgd_stride, std::vector<int> &split,
    std::vector<std::pair<int, int>> &A, double *rdcost) {
  const int end_row = AOMMIN(start_row + max_unit_height, height);
  const int end_col = AOMMIN(start_col + max_unit_width, width);
  // Check for special cases near boundary.
  const bool is_horz_partitioning_allowed =
      (max_unit_height >= quadtree_max_size);
  const bool is_vert_partitioning_allowed =
      (max_unit_width >= quadtree_max_size);
  const bool is_split_partitioning_allowed =
      is_horz_partitioning_allowed && is_vert_partitioning_allowed;

  auto best_rdcost = DBL_MAX;
  std::vector<std::pair<int, int>> best_A;
  std::vector<std::vector<uint16_t>> best_out(
      max_unit_height, std::vector<uint16_t>(max_unit_width));
  GuidedQuadTreePartitionType best_partition_type = GUIDED_QT_INVALID;

  for (int type = 0; type < GUIDED_QT_TYPES; ++type) {
    const auto this_partition_type = (GuidedQuadTreePartitionType)type;
    // Check for special cases near boundary.
    if (!is_horz_partitioning_allowed &&
        (this_partition_type == GUIDED_QT_HORZ)) {
      continue;
    }
    if (!is_vert_partitioning_allowed &&
        (this_partition_type == GUIDED_QT_VERT)) {
      continue;
    }
    if (!is_split_partitioning_allowed &&
        (this_partition_type == GUIDED_QT_SPLIT)) {
      continue;
    }
    // Try this partition type.
    double this_rdcost;
    std::vector<std::pair<int, int>> this_A;
    std::vector<std::vector<uint16_t>> this_out(
        max_unit_height, std::vector<uint16_t>(max_unit_width));
    try_one_partition(
        interm, this_partition_type, src, src_stride, dgd, dgd_stride,
        start_row, end_row, start_col, end_col, max_unit_width, max_unit_height,
        quadtset, rdmult, prev_A, quad_split_costs, binary_split_costs,
        norestorecosts, bit_depth, is_horz_partitioning_allowed,
        is_vert_partitioning_allowed, &this_rdcost, this_out, this_A);
    if (this_rdcost < best_rdcost) {
      best_rdcost = this_rdcost;
      best_A = this_A;
      best_out = this_out;
      best_partition_type = this_partition_type;
    }
  }

  // Save RDCost.
  *rdcost = best_rdcost;

  // Save a0, a1 pairs.
  for (auto &a0a1 : best_A) {
    A.push_back(a0a1);
  }

  // Save split decision.
  if (!is_horz_partitioning_allowed && !is_vert_partitioning_allowed) {
    // Nothing should be added to 'split' array.
    assert(best_partition_type == GUIDED_QT_NONE);
    return;
  }
  assert(best_partition_type >= 0 && best_partition_type < GUIDED_QT_TYPES);
  split.push_back(best_partition_type);
}

static void apply_quadtree_partitioning(
    const std::vector<std::vector<std::vector<double>>> &interm, int start_row,
    int start_col, int width, int height, int quadtree_max_size,
    int max_unit_width, int max_unit_height, const int *quadtset, int bit_depth,
    const std::vector<int> &split, size_t &split_index,
    const std::vector<std::pair<int, int>> &A, size_t &A_index, uint16_t *dgd,
    int dgd_stride);

// Top-level function to apply guided restoration on encoder side.
static int restore_cnn_quadtree_encode_img_tflite_highbd(
    YV12_BUFFER_CONFIG *source_frame, AV1_COMMON *cm, int superres_denom,
    int rdmult, const int *quad_split_costs, const int *binary_split_costs,
    int (*norestorecosts)[2], int num_threads, int bit_depth, int is_intra_only,
    int is_luma, int cnn_index, QUADInfo *quad_info, double *rdcost) {
  YV12_BUFFER_CONFIG *dgd_buf = &cm->cur_frame->buf;
  uint16_t *dgd = CONVERT_TO_SHORTPTR(dgd_buf->y_buffer);
  const int dgd_stride = dgd_buf->y_stride;
  const int qindex = cm->quant_params.base_qindex;
  const int width = cm->superres_upscaled_width;
  const int height = cm->superres_upscaled_height;

  // Get 2-channel intermediate restoration.
  std::vector<std::vector<std::vector<double>>> interm(
      height, std::vector<std::vector<double>>(width, std::vector<double>(2)));
  if (!generate_interm_guided_restoration(
          dgd, dgd_stride, qindex, superres_denom, width, height, num_threads,
          is_intra_only, is_luma, cnn_index, bit_depth, interm)) {
    return 0;
  }

  // Initialization.
  const uint16_t *src = CONVERT_TO_SHORTPTR(source_frame->y_buffer);
  const int src_stride = source_frame->y_stride;
  const int *quadtset = get_quadparm_from_qindex(
      qindex, superres_denom, is_intra_only, is_luma, cnn_index);
  const int A0_min = quadtset[2];
  const int A1_min = quadtset[3];
  const int norestore_ctx =
      get_guided_norestore_ctx(qindex, superres_denom, is_intra_only);
  const int null_norestorecosts[2] = { 0, 0 };
  const int *this_norestorecosts =
      norestore_ctx == -1 ? null_norestorecosts : norestorecosts[norestore_ctx];

  // Try all possible quadtree unit sizes.
  int best_unit_index = -1;
  std::vector<int> best_split;              // selected partitioning options.
  std::vector<std::pair<int, int>> best_A;  // selected a0, a1 weight pairs.
  double best_rdcost_total = DBL_MAX;
  for (int this_unit_index = 0; this_unit_index < GUIDED_QT_UNIT_SIZES;
       ++this_unit_index) {
    const int quadtree_max_size =
        quad_tree_get_unit_size(width, height, this_unit_index);
    // For each quadtree unit, compute the best partitioning out of
    // NONE, SPLIT, HORZ and VERT based on RD cost.
    std::vector<int> this_split;              // selected partitioning options.
    std::vector<std::pair<int, int>> this_A;  // selected a0, a1 weight pairs.
    double this_rdcost_total = 0.0;
    // Previous a0, a1 pair is mid-point of the range by default.
    std::pair<int, int> prev_A =
        std::make_pair(GUIDED_A_MID + A0_min, GUIDED_A_MID + A1_min);
    const int ext_size = quadtree_max_size * 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 : quadtree_max_size;
      for (int col = 0; col < width;) {
        const int remaining_width = width - col;
        const int this_unit_width =
            (remaining_width < ext_size) ? remaining_width : quadtree_max_size;
        double this_rdcost;
        select_quadtree_partitioning(
            interm, src, src_stride, row, col, width, height, quadtree_max_size,
            this_unit_width, this_unit_height, quadtset, rdmult, prev_A,
            quad_split_costs, binary_split_costs, this_norestorecosts,
            bit_depth, dgd, dgd_stride, this_split, this_A, &this_rdcost);
        // updates.
        this_rdcost_total += this_rdcost;
        prev_A = this_A.back();
        col += this_unit_width;
      }
      row += this_unit_height;
    }
    // Update best options.
    if (this_rdcost_total < best_rdcost_total) {
      best_unit_index = this_unit_index;
      best_split = this_split;
      best_A = this_A;
      best_rdcost_total = this_rdcost_total;
    }
  }

  // Fill in the best options.
  quad_info->unit_index = best_unit_index;
  quad_info->split_info_length = (int)best_split.size();
  quad_info->unit_info_length = (int)best_A.size();
  av1_alloc_quadtree_struct(cm, quad_info);
  for (unsigned int i = 0; i < best_split.size(); ++i) {
    quad_info->split_info[i].split = best_split[i];
  }
  for (unsigned int i = 0; i < best_A.size(); ++i) {
    quad_info->unit_info[i].xqd[0] = best_A[i].first;
    quad_info->unit_info[i].xqd[1] = best_A[i].second;
  }
  *rdcost = best_rdcost_total;

  // Apply guided restoration to 'dgd' using best options above.
  size_t split_index = 0;
  size_t A_index = 0;
  const int quadtree_max_size = quad_info->unit_size;
  const int ext_size = quadtree_max_size * 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 : quadtree_max_size;
    for (int col = 0; col < width;) {
      const int remaining_width = width - col;
      const int this_unit_width =
          (remaining_width < ext_size) ? remaining_width : quadtree_max_size;
      apply_quadtree_partitioning(
          interm, row, col, width, height, quadtree_max_size, this_unit_width,
          this_unit_height, quadtset, bit_depth, best_split, split_index,
          best_A, A_index, dgd, dgd_stride);
      col += this_unit_width;
    }
    row += this_unit_height;
  }

  return 1;
}

extern "C" int av1_restore_cnn_quadtree_encode_tflite(
    struct AV1Common *cm, YV12_BUFFER_CONFIG *source_frame, int RDMULT,
    int *quad_split_costs, int *binary_split_costs, int (*norestorecosts)[2],
    int num_threads, const int apply_cnn[MAX_MB_PLANE],
    const int cnn_indices[MAX_MB_PLANE], QUADInfo *quad_info, double *rdcost) {
  YV12_BUFFER_CONFIG *buf = &cm->cur_frame->buf;
  const int is_intra_only = frame_is_intra_only(cm);
  for (int plane = 0; plane < av1_num_planes(cm); ++plane) {
    if (!apply_cnn[plane]) continue;
    const int is_luma = (plane == AOM_PLANE_Y);
    const int cnn_index = cnn_indices[plane];
    assert(cnn_index >= 0 &&
           cnn_index < av1_num_cnn_indices_for_plane(cm, plane));
    int ret = 1;
    switch (plane) {
      case AOM_PLANE_Y:
        ret = restore_cnn_quadtree_encode_img_tflite_highbd(
            source_frame, cm, cm->superres_scale_denominator, RDMULT,
            quad_split_costs, binary_split_costs, norestorecosts, num_threads,
            cm->seq_params.bit_depth, is_intra_only, is_luma, cnn_index,
            quad_info, rdcost);
        if (ret == 0) return ret;
        break;
      case AOM_PLANE_U:
        ret = av1_restore_cnn_img_tflite_highbd(
            cm->quant_params.base_qindex, cm->superres_scale_denominator,
            CONVERT_TO_SHORTPTR(buf->u_buffer), buf->uv_crop_width,
            buf->uv_crop_height, buf->uv_stride,
            CONVERT_TO_SHORTPTR(buf->u_buffer), buf->uv_stride, num_threads,
            cm->seq_params.bit_depth, is_intra_only, is_luma, cnn_index);
        if (ret == 0) return ret;
        break;
      case AOM_PLANE_V:
        ret = av1_restore_cnn_img_tflite_highbd(
            cm->quant_params.base_qindex, cm->superres_scale_denominator,
            CONVERT_TO_SHORTPTR(buf->v_buffer), buf->uv_crop_width,
            buf->uv_crop_height, buf->uv_stride,
            CONVERT_TO_SHORTPTR(buf->v_buffer), buf->uv_stride, num_threads,
            cm->seq_params.bit_depth, is_intra_only, is_luma, cnn_index);
        if (ret == 0) return ret;
        break;
      default: assert(0 && "Invalid plane index"); return 0;
    }
  }
  return 1;
}

// ------------------- Guided Quadtree: Decoder ------------------------------//

// Given the 2-channel intermediate output in 'interm' and weight parameters,
// restores one quadtree unit in 'dgd'.
static void apply_linear_combination(
    const std::vector<std::vector<std::vector<double>>> &interm, int start_row,
    int end_row, int start_col, int end_col, int unit_width, int unit_height,
    const int *quadtset, int bit_depth,
    const std::vector<std::pair<int, int>> &A, size_t &A_index, uint16_t *dgd,
    int dgd_stride) {
  // Get scale parameters.
  const int scale0 = quadtset[0];
  const int scale1 = quadtset[1];

  for (int row = start_row; row < end_row; row += unit_height) {
    const int this_start_row = row;
    const int this_end_row = AOMMIN(row + unit_height, end_row);
    for (int col = start_col; col < end_col; col += unit_width) {
      const int this_start_col = col;
      const int this_end_col = AOMMIN(col + unit_width, end_col);

      // Get weight parameters for this unit.
      const auto this_A = A[A_index++];
      const int a0 = this_A.first;
      const int a1 = this_A.second;

      // Restore this unit.
      for (int r = this_start_row; r < this_end_row; ++r) {
        for (int c = this_start_col; c < this_end_col; ++c) {
          const int dgd_unclipped = int(round(dgd[r * dgd_stride + c] +
                                              a0 * interm[r][c][0] / scale0 +
                                              a1 * interm[r][c][1] / scale1));
          dgd[r * dgd_stride + c] = clip_pixel_highbd(dgd_unclipped, bit_depth);
        }
      }
    }
  }
}

// Given intermediate restoration 'interm', quadtree partitioning info 'split'
// and weight parameters 'A', restores the unit starting at 'row' and 'col'
// inside 'dgd'.
static void apply_quadtree_partitioning(
    const std::vector<std::vector<std::vector<double>>> &interm, int start_row,
    int start_col, int width, int height, int quadtree_max_size,
    int max_unit_width, int max_unit_height, const int *quadtset, int bit_depth,
    const std::vector<int> &split, size_t &split_index,
    const std::vector<std::pair<int, int>> &A, size_t &A_index, uint16_t *dgd,
    int dgd_stride) {
  const int end_row = AOMMIN(start_row + max_unit_height, height);
  const int end_col = AOMMIN(start_col + max_unit_width, width);
  // Check for special cases near boundary.
  const bool is_horz_partitioning_allowed =
      (max_unit_height >= quadtree_max_size);
  const bool is_vert_partitioning_allowed =
      (max_unit_width >= quadtree_max_size);

  // Get partition type.
  GuidedQuadTreePartitionType partition_type = GUIDED_QT_NONE;
  if (is_horz_partitioning_allowed || is_vert_partitioning_allowed) {
    partition_type = (GuidedQuadTreePartitionType)split[split_index++];
  }
  assert(partition_type >= 0 && partition_type < GUIDED_QT_TYPES);

  // Get unit width and height based on partition type.
  int unit_width;
  int unit_height;
  get_unit_size(max_unit_width, max_unit_height, partition_type, &unit_width,
                &unit_height);

  // Compute restored unit, a0 and a1 with given A parameters.
  apply_linear_combination(interm, start_row, end_row, start_col, end_col,
                           unit_width, unit_height, quadtset, bit_depth, A,
                           A_index, dgd, dgd_stride);
}

// Top-level function to apply guided restoration on decoder side.
static int restore_cnn_quadtree_decode_img_tflite_highbd(
    AV1_COMMON *cm, int superres_denom, int num_threads, int bit_depth,
    int is_intra_only, int is_luma, int cnn_index) {
  YV12_BUFFER_CONFIG *dgd_buf = &cm->cur_frame->buf;
  uint16_t *dgd = CONVERT_TO_SHORTPTR(dgd_buf->y_buffer);
  const int dgd_stride = dgd_buf->y_stride;
  const int qindex = cm->quant_params.base_qindex;
  const int width = cm->superres_upscaled_width;
  const int height = cm->superres_upscaled_height;

  // Get 2-channel intermediate restoration.
  std::vector<std::vector<std::vector<double>>> interm(
      height, std::vector<std::vector<double>>(width, std::vector<double>(2)));
  if (!generate_interm_guided_restoration(
          dgd, dgd_stride, qindex, superres_denom, width, height, num_threads,
          is_intra_only, is_luma, cnn_index, bit_depth, interm)) {
    return 0;
  }

  // Get quadtree params.
  const QUADInfo *const quad_info = &cm->cnn_quad_info;
  const int quadtree_max_size = quad_info->unit_size;
  const int *quadtset = get_quadparm_from_qindex(
      qindex, superres_denom, is_intra_only, is_luma, cnn_index);

  // Get partitioning types.
  std::vector<int> split;
  split.reserve(quad_info->split_info_length);
  for (int i = 0; i < quad_info->split_info_length; ++i) {
    split.push_back(quad_info->split_info[i].split);
  }

  // Get a0,a1 pairs.
  std::vector<std::pair<int, int>> A;
  A.reserve(quad_info->unit_info_length);
  for (int i = 0; i < quad_info->unit_info_length; ++i) {
    A.emplace_back(quad_info->unit_info[i].xqd[0],
                   quad_info->unit_info[i].xqd[1]);
  }

  // For each quadtree unit, apply given quadtree partitioning.
  size_t split_index = 0;
  size_t A_index = 0;
  const int ext_size = quadtree_max_size * 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 : quadtree_max_size;
    for (int col = 0; col < width;) {
      const int remaining_width = width - col;
      const int this_unit_width =
          (remaining_width < ext_size) ? remaining_width : quadtree_max_size;
      apply_quadtree_partitioning(interm, row, col, width, height,
                                  quadtree_max_size, this_unit_width,
                                  this_unit_height, quadtset, bit_depth, split,
                                  split_index, A, A_index, dgd, dgd_stride);
      col += this_unit_width;
    }
    row += this_unit_height;
  }
  assert(split_index == split.size());
  assert(A_index == A.size());
  return 1;
}

extern "C" int av1_restore_cnn_quadtree_decode_tflite(
    struct AV1Common *cm, int num_threads, int use_quadtree,
    const int apply_cnn[MAX_MB_PLANE], const int cnn_indices[MAX_MB_PLANE]) {
  YV12_BUFFER_CONFIG *buf = &cm->cur_frame->buf;
  const int is_intra_only = frame_is_intra_only(cm);
  for (int plane = 0; plane < av1_num_planes(cm); ++plane) {
    if (!apply_cnn[plane]) continue;
    const int is_luma = (plane == AOM_PLANE_Y);
    if (is_luma && !use_quadtree) continue;
    const int cnn_index = cnn_indices[plane];
    assert(cnn_index >= 0 &&
           cnn_index < av1_num_cnn_indices_for_plane(cm, plane));
    int ret = 1;
    switch (plane) {
      case AOM_PLANE_Y:
        ret = restore_cnn_quadtree_decode_img_tflite_highbd(
            cm, cm->superres_scale_denominator, num_threads,
            cm->seq_params.bit_depth, is_intra_only, is_luma, cnn_index);
        if (ret == 0) return ret;
        break;
      case AOM_PLANE_U:
        ret = av1_restore_cnn_img_tflite_highbd(
            cm->quant_params.base_qindex, cm->superres_scale_denominator,
            CONVERT_TO_SHORTPTR(buf->u_buffer), buf->uv_crop_width,
            buf->uv_crop_height, buf->uv_stride,
            CONVERT_TO_SHORTPTR(buf->u_buffer), buf->uv_stride, num_threads,
            cm->seq_params.bit_depth, is_intra_only, is_luma, cnn_index);
        if (ret == 0) return ret;
        break;
      case AOM_PLANE_V:
        ret = av1_restore_cnn_img_tflite_highbd(
            cm->quant_params.base_qindex, cm->superres_scale_denominator,
            CONVERT_TO_SHORTPTR(buf->v_buffer), buf->uv_crop_width,
            buf->uv_crop_height, buf->uv_stride,
            CONVERT_TO_SHORTPTR(buf->v_buffer), buf->uv_stride, num_threads,
            cm->seq_params.bit_depth, is_intra_only, is_luma, cnn_index);
        if (ret == 0) return ret;
        break;
      default: assert(0 && "Invalid plane index"); return 0;
    }
  }
  return 1;
}
#endif  // CONFIG_CNN_GUIDED_QUADTREE

#endif  // CONFIG_CNN_RESTORATION