/*

 * Copyright (c) 2016, 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 <assert.h>

#include <float.h>

#include <limits.h>

#include <math.h>

#include "config/aom_scale_rtcd.h"

#include "config/av1_rtcd.h"

#include "aom_dsp/aom_dsp_common.h"

#include "aom_dsp/binary_codes_writer.h"

#include "aom_dsp/mathutils.h"

#include "aom_dsp/psnr.h"

#include "aom_mem/aom_mem.h"

#include "aom_ports/mem.h"

#include "av1/common/av1_common_int.h"

#include "av1/common/quant_common.h"

#include "av1/common/restoration.h"

#include "av1/encoder/av1_quantize.h"

#include "av1/encoder/encoder.h"

#include "av1/encoder/picklpf.h"

#include "av1/encoder/pickrst.h"

// When set to RESTORE_WIENER or RESTORE_SGRPROJ only those are allowed.

// When set to RESTORE_TYPES we allow switchable.

static const RestorationType force_restore_type = RESTORE_TYPES;

// Number of Wiener iterations

#define NUM_WIENER_ITERS 5

// Penalty factor for use of dual sgr

#define DUAL_SGR_PENALTY_MULT 0.01

// Working precision for Wiener filter coefficients

#define WIENER_TAP_SCALE_FACTOR ((int64_t)1 << 16)

#define SGRPROJ_EP_GRP1_START_IDX 0

#define SGRPROJ_EP_GRP1_END_IDX 9

#define SGRPROJ_EP_GRP1_SEARCH_COUNT 4

#define SGRPROJ_EP_GRP2_3_SEARCH_COUNT 2

static const int sgproj_ep_grp1_seed[SGRPROJ_EP_GRP1_SEARCH_COUNT] = { 0, 3, 6,

                                                                       9 };

static const int sgproj_ep_grp2_3[SGRPROJ_EP_GRP2_3_SEARCH_COUNT][14] = {

  { 10, 10, 11, 11, 12, 12, 13, 13, 13, 13, -1, -1, -1, -1 },

  { 14, 14, 14, 14, 14, 14, 14, 15, 15, 15, 15, 15, 15, 15 }

};

typedef int64_t (*sse_extractor_type)(const YV12_BUFFER_CONFIG *a,

                                      const YV12_BUFFER_CONFIG *b);

typedef int64_t (*sse_part_extractor_type)(const YV12_BUFFER_CONFIG *a,

                                           const YV12_BUFFER_CONFIG *b,

                                           int hstart, int width, int vstart,

                                           int height);

typedef uint64_t (*var_part_extractor_type)(const YV12_BUFFER_CONFIG *a,

                                            int hstart, int width, int vstart,

                                            int height);

#if CONFIG_AV1_HIGHBITDEPTH

#define NUM_EXTRACTORS (3 * (1 + 1))

#else

#define NUM_EXTRACTORS 3

#endif

static const sse_part_extractor_type sse_part_extractors[NUM_EXTRACTORS] = {

  aom_get_y_sse_part,        aom_get_u_sse_part,

  aom_get_v_sse_part,

#if CONFIG_AV1_HIGHBITDEPTH

  aom_highbd_get_y_sse_part, aom_highbd_get_u_sse_part,

  aom_highbd_get_v_sse_part,

#endif

};

static const var_part_extractor_type var_part_extractors[NUM_EXTRACTORS] = {

  aom_get_y_var,        aom_get_u_var,        aom_get_v_var,

#if CONFIG_AV1_HIGHBITDEPTH

  aom_highbd_get_y_var, aom_highbd_get_u_var, aom_highbd_get_v_var,

#endif

};

static int64_t sse_restoration_unit(const RestorationTileLimits *limits,

                                    const YV12_BUFFER_CONFIG *src,

                                    const YV12_BUFFER_CONFIG *dst, int plane,

                                    int highbd) {

  return sse_part_extractors[3 * highbd + plane](

      src, dst, limits->h_start, limits->h_end - limits->h_start,

      limits->v_start, limits->v_end - limits->v_start);

}

static uint64_t var_restoration_unit(const RestorationTileLimits *limits,

                                     const YV12_BUFFER_CONFIG *src, int plane,

                                     int highbd) {

  return var_part_extractors[3 * highbd + plane](

      src, limits->h_start, limits->h_end - limits->h_start, limits->v_start,

      limits->v_end - limits->v_start);

}

typedef struct {

  // The best coefficients for Wiener or Sgrproj restoration

  WienerInfo wiener;

  SgrprojInfo sgrproj;

  // The sum of squared errors for this rtype.

  int64_t sse[RESTORE_SWITCHABLE_TYPES];

  // The rtype to use for this unit given a frame rtype as

  // index. Indices: WIENER, SGRPROJ, SWITCHABLE.

  RestorationType best_rtype[RESTORE_TYPES - 1];

  // This flag will be set based on the speed feature

  // 'prune_sgr_based_on_wiener'. 0 implies no pruning and 1 implies pruning.

  uint8_t skip_sgr_eval;

} RestUnitSearchInfo;

typedef struct {

  const YV12_BUFFER_CONFIG *src;

  YV12_BUFFER_CONFIG *dst;

  const AV1_COMMON *cm;

  const MACROBLOCK *x;

  int plane;

  int plane_width;

  int plane_height;

  RestUnitSearchInfo *rusi;

  // Speed features

  const LOOP_FILTER_SPEED_FEATURES *lpf_sf;

  uint8_t *dgd_buffer;

  int dgd_stride;

  const uint8_t *src_buffer;

  int src_stride;

  // sse and bits are initialised by reset_rsc in search_rest_type

  int64_t sse;

  int64_t bits;

  int tile_y0, tile_stripe0;

  // sgrproj and wiener are initialised by rsc_on_tile when starting the first

  // tile in the frame.

  SgrprojInfo sgrproj;

  WienerInfo wiener;

  AV1PixelRect tile_rect;

} RestSearchCtxt;

static AOM_INLINE void rsc_on_tile(void *priv) {

  RestSearchCtxt *rsc = (RestSearchCtxt *)priv;

  set_default_sgrproj(&rsc->sgrproj);

  set_default_wiener(&rsc->wiener);

  rsc->tile_stripe0 = 0;

}

static AOM_INLINE void reset_rsc(RestSearchCtxt *rsc) {

  rsc->sse = 0;

  rsc->bits = 0;

}

static AOM_INLINE void init_rsc(const YV12_BUFFER_CONFIG *src,

                                const AV1_COMMON *cm, const MACROBLOCK *x,

                                const LOOP_FILTER_SPEED_FEATURES *lpf_sf,

                                int plane, RestUnitSearchInfo *rusi,

                                YV12_BUFFER_CONFIG *dst, RestSearchCtxt *rsc) {

  rsc->src = src;

  rsc->dst = dst;

  rsc->cm = cm;

  rsc->x = x;

  rsc->plane = plane;

  rsc->rusi = rusi;

  rsc->lpf_sf = lpf_sf;

  const YV12_BUFFER_CONFIG *dgd = &cm->cur_frame->buf;

  const int is_uv = plane != AOM_PLANE_Y;

  rsc->plane_width = src->crop_widths[is_uv];

  rsc->plane_height = src->crop_heights[is_uv];

  rsc->src_buffer = src->buffers[plane];

  rsc->src_stride = src->strides[is_uv];

  rsc->dgd_buffer = dgd->buffers[plane];

  rsc->dgd_stride = dgd->strides[is_uv];

  rsc->tile_rect = av1_whole_frame_rect(cm, is_uv);

  assert(src->crop_widths[is_uv] == dgd->crop_widths[is_uv]);

  assert(src->crop_heights[is_uv] == dgd->crop_heights[is_uv]);

}

static int64_t try_restoration_unit(const RestSearchCtxt *rsc,

                                    const RestorationTileLimits *limits,

                                    const AV1PixelRect *tile_rect,

                                    const RestorationUnitInfo *rui) {

  const AV1_COMMON *const cm = rsc->cm;

  const int plane = rsc->plane;

  const int is_uv = plane > 0;

  const RestorationInfo *rsi = &cm->rst_info[plane];

  RestorationLineBuffers rlbs;

  const int bit_depth = cm->seq_params->bit_depth;

  const int highbd = cm->seq_params->use_highbitdepth;

  const YV12_BUFFER_CONFIG *fts = &cm->cur_frame->buf;

  // TODO(yunqing): For now, only use optimized LR filter in decoder. Can be

  // also used in encoder.

  const int optimized_lr = 0;

  av1_loop_restoration_filter_unit(

      limits, rui, &rsi->boundaries, &rlbs, tile_rect, rsc->tile_stripe0,

      is_uv && cm->seq_params->subsampling_x,

      is_uv && cm->seq_params->subsampling_y, highbd, bit_depth,

      fts->buffers[plane], fts->strides[is_uv], rsc->dst->buffers[plane],

      rsc->dst->strides[is_uv], cm->rst_tmpbuf, optimized_lr);

  return sse_restoration_unit(limits, rsc->src, rsc->dst, plane, highbd);

}

int64_t av1_lowbd_pixel_proj_error_c(const uint8_t *src8, int width, int height,

                                     int src_stride, const uint8_t *dat8,

                                     int dat_stride, int32_t *flt0,

                                     int flt0_stride, int32_t *flt1,

                                     int flt1_stride, int xq[2],

                                     const sgr_params_type *params) {

  int i, j;

  const uint8_t *src = src8;

  const uint8_t *dat = dat8;

  int64_t err = 0;

  if (params->r[0] > 0 && params->r[1] > 0) {

    for (i = 0; i < height; ++i) {

      for (j = 0; j < width; ++j) {

        assert(flt1[j] < (1 << 15) && flt1[j] > -(1 << 15));

        assert(flt0[j] < (1 << 15) && flt0[j] > -(1 << 15));

        const int32_t u = (int32_t)(dat[j] << SGRPROJ_RST_BITS);

        int32_t v = u << SGRPROJ_PRJ_BITS;

        v += xq[0] * (flt0[j] - u) + xq[1] * (flt1[j] - u);

        const int32_t e =

            ROUND_POWER_OF_TWO(v, SGRPROJ_RST_BITS + SGRPROJ_PRJ_BITS) - src[j];

        err += ((int64_t)e * e);

      }

      dat += dat_stride;

      src += src_stride;

      flt0 += flt0_stride;

      flt1 += flt1_stride;

    }

  } else if (params->r[0] > 0) {

    for (i = 0; i < height; ++i) {

      for (j = 0; j < width; ++j) {

        assert(flt0[j] < (1 << 15) && flt0[j] > -(1 << 15));

        const int32_t u = (int32_t)(dat[j] << SGRPROJ_RST_BITS);

        int32_t v = u << SGRPROJ_PRJ_BITS;

        v += xq[0] * (flt0[j] - u);

        const int32_t e =

            ROUND_POWER_OF_TWO(v, SGRPROJ_RST_BITS + SGRPROJ_PRJ_BITS) - src[j];

        err += ((int64_t)e * e);

      }

      dat += dat_stride;

      src += src_stride;

      flt0 += flt0_stride;

    }

  } else if (params->r[1] > 0) {

    for (i = 0; i < height; ++i) {

      for (j = 0; j < width; ++j) {

        assert(flt1[j] < (1 << 15) && flt1[j] > -(1 << 15));

        const int32_t u = (int32_t)(dat[j] << SGRPROJ_RST_BITS);

        int32_t v = u << SGRPROJ_PRJ_BITS;

        v += xq[1] * (flt1[j] - u);

        const int32_t e =

            ROUND_POWER_OF_TWO(v, SGRPROJ_RST_BITS + SGRPROJ_PRJ_BITS) - src[j];

        err += ((int64_t)e * e);

      }

      dat += dat_stride;

      src += src_stride;

      flt1 += flt1_stride;

    }

  } else {

    for (i = 0; i < height; ++i) {

      for (j = 0; j < width; ++j) {

        const int32_t e = (int32_t)(dat[j]) - src[j];

        err += ((int64_t)e * e);

      }

      dat += dat_stride;

      src += src_stride;

    }

  }

  return err;

}

#if CONFIG_AV1_HIGHBITDEPTH

int64_t av1_highbd_pixel_proj_error_c(const uint8_t *src8, int width,

                                      int height, int src_stride,

                                      const uint8_t *dat8, int dat_stride,

                                      int32_t *flt0, int flt0_stride,

                                      int32_t *flt1, int flt1_stride, int xq[2],

                                      const sgr_params_type *params) {

  const uint16_t *src = CONVERT_TO_SHORTPTR(src8);

  const uint16_t *dat = CONVERT_TO_SHORTPTR(dat8);

  int i, j;

  int64_t err = 0;

  const int32_t half = 1 << (SGRPROJ_RST_BITS + SGRPROJ_PRJ_BITS - 1);

  if (params->r[0] > 0 && params->r[1] > 0) {

    int xq0 = xq[0];

    int xq1 = xq[1];

    for (i = 0; i < height; ++i) {

      for (j = 0; j < width; ++j) {

        const int32_t d = dat[j];

        const int32_t s = src[j];

        const int32_t u = (int32_t)(d << SGRPROJ_RST_BITS);

        int32_t v0 = flt0[j] - u;

        int32_t v1 = flt1[j] - u;

        int32_t v = half;

        v += xq0 * v0;

        v += xq1 * v1;

        const int32_t e = (v >> (SGRPROJ_RST_BITS + SGRPROJ_PRJ_BITS)) + d - s;

        err += ((int64_t)e * e);

      }

      dat += dat_stride;

      flt0 += flt0_stride;

      flt1 += flt1_stride;

      src += src_stride;

    }

  } else if (params->r[0] > 0 || params->r[1] > 0) {

    int exq;

    int32_t *flt;

    int flt_stride;

    if (params->r[0] > 0) {

      exq = xq[0];

      flt = flt0;

      flt_stride = flt0_stride;

    } else {

      exq = xq[1];

      flt = flt1;

      flt_stride = flt1_stride;

    }

    for (i = 0; i < height; ++i) {

      for (j = 0; j < width; ++j) {

        const int32_t d = dat[j];

        const int32_t s = src[j];

        const int32_t u = (int32_t)(d << SGRPROJ_RST_BITS);

        int32_t v = half;

        v += exq * (flt[j] - u);

        const int32_t e = (v >> (SGRPROJ_RST_BITS + SGRPROJ_PRJ_BITS)) + d - s;

        err += ((int64_t)e * e);

      }

      dat += dat_stride;

      flt += flt_stride;

      src += src_stride;

    }

  } else {

    for (i = 0; i < height; ++i) {

      for (j = 0; j < width; ++j) {

        const int32_t d = dat[j];

        const int32_t s = src[j];

        const int32_t e = d - s;

        err += ((int64_t)e * e);

      }

      dat += dat_stride;

      src += src_stride;

    }

  }

  return err;

}

#endif  // CONFIG_AV1_HIGHBITDEPTH

static int64_t get_pixel_proj_error(const uint8_t *src8, int width, int height,

                                    int src_stride, const uint8_t *dat8,

                                    int dat_stride, int use_highbitdepth,

                                    int32_t *flt0, int flt0_stride,

                                    int32_t *flt1, int flt1_stride, int *xqd,

                                    const sgr_params_type *params) {

  int xq[2];

  av1_decode_xq(xqd, xq, params);

#if CONFIG_AV1_HIGHBITDEPTH

  if (use_highbitdepth) {

    return av1_highbd_pixel_proj_error(src8, width, height, src_stride, dat8,

                                       dat_stride, flt0, flt0_stride, flt1,

                                       flt1_stride, xq, params);

  } else {

    return av1_lowbd_pixel_proj_error(src8, width, height, src_stride, dat8,

                                      dat_stride, flt0, flt0_stride, flt1,

                                      flt1_stride, xq, params);

  }

#else

  (void)use_highbitdepth;

  return av1_lowbd_pixel_proj_error(src8, width, height, src_stride, dat8,

                                    dat_stride, flt0, flt0_stride, flt1,

                                    flt1_stride, xq, params);

#endif

}

#define USE_SGRPROJ_REFINEMENT_SEARCH 1

static int64_t finer_search_pixel_proj_error(

    const uint8_t *src8, int width, int height, int src_stride,

    const uint8_t *dat8, int dat_stride, int use_highbitdepth, int32_t *flt0,

    int flt0_stride, int32_t *flt1, int flt1_stride, int start_step, int *xqd,

    const sgr_params_type *params) {

  int64_t err = get_pixel_proj_error(

      src8, width, height, src_stride, dat8, dat_stride, use_highbitdepth, flt0,

      flt0_stride, flt1, flt1_stride, xqd, params);

  (void)start_step;

#if USE_SGRPROJ_REFINEMENT_SEARCH

  int64_t err2;

  int tap_min[] = { SGRPROJ_PRJ_MIN0, SGRPROJ_PRJ_MIN1 };

  int tap_max[] = { SGRPROJ_PRJ_MAX0, SGRPROJ_PRJ_MAX1 };

  for (int s = start_step; s >= 1; s >>= 1) {

    for (int p = 0; p < 2; ++p) {

      if ((params->r[0] == 0 && p == 0) || (params->r[1] == 0 && p == 1)) {

        continue;

      }

      int skip = 0;

      do {

        if (xqd[p] - s >= tap_min[p]) {

          xqd[p] -= s;

          err2 =

              get_pixel_proj_error(src8, width, height, src_stride, dat8,

                                   dat_stride, use_highbitdepth, flt0,

                                   flt0_stride, flt1, flt1_stride, xqd, params);

          if (err2 > err) {

            xqd[p] += s;

          } else {

            err = err2;

            skip = 1;

            // At the highest step size continue moving in the same direction

            if (s == start_step) continue;

          }

        }

        break;

      } while (1);

      if (skip) break;

      do {

        if (xqd[p] + s <= tap_max[p]) {

          xqd[p] += s;

          err2 =

              get_pixel_proj_error(src8, width, height, src_stride, dat8,

                                   dat_stride, use_highbitdepth, flt0,

                                   flt0_stride, flt1, flt1_stride, xqd, params);

          if (err2 > err) {

            xqd[p] -= s;

          } else {

            err = err2;

            // At the highest step size continue moving in the same direction

            if (s == start_step) continue;

          }

        }

        break;

      } while (1);

    }

  }

#endif  // USE_SGRPROJ_REFINEMENT_SEARCH

  return err;

}

static int64_t signed_rounded_divide(int64_t dividend, int64_t divisor) {

  if (dividend < 0)

    return (dividend - divisor / 2) / divisor;

  else

    return (dividend + divisor / 2) / divisor;

}

static AOM_INLINE void calc_proj_params_r0_r1_c(

    const uint8_t *src8, int width, int height, int src_stride,

    const uint8_t *dat8, int dat_stride, int32_t *flt0, int flt0_stride,

    int32_t *flt1, int flt1_stride, int64_t H[2][2], int64_t C[2]) {

  const int size = width * height;

  const uint8_t *src = src8;

  const uint8_t *dat = dat8;

  for (int i = 0; i < height; ++i) {

    for (int j = 0; j < width; ++j) {

      const int32_t u = (int32_t)(dat[i * dat_stride + j] << SGRPROJ_RST_BITS);

      const int32_t s =

          (int32_t)(src[i * src_stride + j] << SGRPROJ_RST_BITS) - u;

      const int32_t f1 = (int32_t)flt0[i * flt0_stride + j] - u;

      const int32_t f2 = (int32_t)flt1[i * flt1_stride + j] - u;

      H[0][0] += (int64_t)f1 * f1;

      H[1][1] += (int64_t)f2 * f2;

      H[0][1] += (int64_t)f1 * f2;

      C[0] += (int64_t)f1 * s;

      C[1] += (int64_t)f2 * s;

    }

  }

  H[0][0] /= size;

  H[0][1] /= size;

  H[1][1] /= size;

  H[1][0] = H[0][1];

  C[0] /= size;

  C[1] /= size;

}

static AOM_INLINE void calc_proj_params_r0_r1_high_bd_c(

    const uint8_t *src8, int width, int height, int src_stride,

    const uint8_t *dat8, int dat_stride, int32_t *flt0, int flt0_stride,

    int32_t *flt1, int flt1_stride, int64_t H[2][2], int64_t C[2]) {

  const int size = width * height;

  const uint16_t *src = CONVERT_TO_SHORTPTR(src8);

  const uint16_t *dat = CONVERT_TO_SHORTPTR(dat8);

  for (int i = 0; i < height; ++i) {

    for (int j = 0; j < width; ++j) {

      const int32_t u = (int32_t)(dat[i * dat_stride + j] << SGRPROJ_RST_BITS);

      const int32_t s =

          (int32_t)(src[i * src_stride + j] << SGRPROJ_RST_BITS) - u;

      const int32_t f1 = (int32_t)flt0[i * flt0_stride + j] - u;

      const int32_t f2 = (int32_t)flt1[i * flt1_stride + j] - u;

      H[0][0] += (int64_t)f1 * f1;

      H[1][1] += (int64_t)f2 * f2;

      H[0][1] += (int64_t)f1 * f2;

      C[0] += (int64_t)f1 * s;

      C[1] += (int64_t)f2 * s;

    }

  }

  H[0][0] /= size;

  H[0][1] /= size;

  H[1][1] /= size;

  H[1][0] = H[0][1];

  C[0] /= size;

  C[1] /= size;

}

static AOM_INLINE void calc_proj_params_r0_c(const uint8_t *src8, int width,

                                             int height, int src_stride,

                                             const uint8_t *dat8,

                                             int dat_stride, int32_t *flt0,

                                             int flt0_stride, int64_t H[2][2],

                                             int64_t C[2]) {

  const int size = width * height;

  const uint8_t *src = src8;

  const uint8_t *dat = dat8;

  for (int i = 0; i < height; ++i) {

    for (int j = 0; j < width; ++j) {

      const int32_t u = (int32_t)(dat[i * dat_stride + j] << SGRPROJ_RST_BITS);

      const int32_t s =

          (int32_t)(src[i * src_stride + j] << SGRPROJ_RST_BITS) - u;

      const int32_t f1 = (int32_t)flt0[i * flt0_stride + j] - u;

      H[0][0] += (int64_t)f1 * f1;

      C[0] += (int64_t)f1 * s;

    }

  }

  H[0][0] /= size;

  C[0] /= size;

}

static AOM_INLINE void calc_proj_params_r0_high_bd_c(

    const uint8_t *src8, int width, int height, int src_stride,

    const uint8_t *dat8, int dat_stride, int32_t *flt0, int flt0_stride,

    int64_t H[2][2], int64_t C[2]) {

  const int size = width * height;

  const uint16_t *src = CONVERT_TO_SHORTPTR(src8);

  const uint16_t *dat = CONVERT_TO_SHORTPTR(dat8);

  for (int i = 0; i < height; ++i) {

    for (int j = 0; j < width; ++j) {

      const int32_t u = (int32_t)(dat[i * dat_stride + j] << SGRPROJ_RST_BITS);

      const int32_t s =

          (int32_t)(src[i * src_stride + j] << SGRPROJ_RST_BITS) - u;

      const int32_t f1 = (int32_t)flt0[i * flt0_stride + j] - u;

      H[0][0] += (int64_t)f1 * f1;

      C[0] += (int64_t)f1 * s;

    }

  }

  H[0][0] /= size;

  C[0] /= size;

}

static AOM_INLINE void calc_proj_params_r1_c(const uint8_t *src8, int width,

                                             int height, int src_stride,

                                             const uint8_t *dat8,

                                             int dat_stride, int32_t *flt1,

                                             int flt1_stride, int64_t H[2][2],

                                             int64_t C[2]) {

  const int size = width * height;

  const uint8_t *src = src8;

  const uint8_t *dat = dat8;

  for (int i = 0; i < height; ++i) {

    for (int j = 0; j < width; ++j) {

      const int32_t u = (int32_t)(dat[i * dat_stride + j] << SGRPROJ_RST_BITS);

      const int32_t s =

          (int32_t)(src[i * src_stride + j] << SGRPROJ_RST_BITS) - u;

      const int32_t f2 = (int32_t)flt1[i * flt1_stride + j] - u;

      H[1][1] += (int64_t)f2 * f2;

      C[1] += (int64_t)f2 * s;

    }

  }

  H[1][1] /= size;

  C[1] /= size;

}

static AOM_INLINE void calc_proj_params_r1_high_bd_c(

    const uint8_t *src8, int width, int height, int src_stride,

    const uint8_t *dat8, int dat_stride, int32_t *flt1, int flt1_stride,

    int64_t H[2][2], int64_t C[2]) {

  const int size = width * height;

  const uint16_t *src = CONVERT_TO_SHORTPTR(src8);

  const uint16_t *dat = CONVERT_TO_SHORTPTR(dat8);

  for (int i = 0; i < height; ++i) {

    for (int j = 0; j < width; ++j) {

      const int32_t u = (int32_t)(dat[i * dat_stride + j] << SGRPROJ_RST_BITS);

      const int32_t s =

          (int32_t)(src[i * src_stride + j] << SGRPROJ_RST_BITS) - u;

      const int32_t f2 = (int32_t)flt1[i * flt1_stride + j] - u;

      H[1][1] += (int64_t)f2 * f2;

      C[1] += (int64_t)f2 * s;

    }

  }

  H[1][1] /= size;

  C[1] /= size;

}

// The function calls 3 subfunctions for the following cases :

// 1) When params->r[0] > 0 and params->r[1] > 0. In this case all elements

// of C and H need to be computed.

// 2) When only params->r[0] > 0. In this case only H[0][0] and C[0] are

// non-zero and need to be computed.

// 3) When only params->r[1] > 0. In this case only H[1][1] and C[1] are

// non-zero and need to be computed.

void av1_calc_proj_params_c(const uint8_t *src8, int width, int height,

                            int src_stride, const uint8_t *dat8, int dat_stride,

                            int32_t *flt0, int flt0_stride, int32_t *flt1,

                            int flt1_stride, int64_t H[2][2], int64_t C[2],

                            const sgr_params_type *params) {

  if ((params->r[0] > 0) && (params->r[1] > 0)) {

    calc_proj_params_r0_r1_c(src8, width, height, src_stride, dat8, dat_stride,

                             flt0, flt0_stride, flt1, flt1_stride, H, C);

  } else if (params->r[0] > 0) {

    calc_proj_params_r0_c(src8, width, height, src_stride, dat8, dat_stride,

                          flt0, flt0_stride, H, C);

  } else if (params->r[1] > 0) {

    calc_proj_params_r1_c(src8, width, height, src_stride, dat8, dat_stride,

                          flt1, flt1_stride, H, C);

  }

}

void av1_calc_proj_params_high_bd_c(const uint8_t *src8, int width, int height,

                                    int src_stride, const uint8_t *dat8,

                                    int dat_stride, int32_t *flt0,

                                    int flt0_stride, int32_t *flt1,

                                    int flt1_stride, int64_t H[2][2],

                                    int64_t C[2],

                                    const sgr_params_type *params) {

  if ((params->r[0] > 0) && (params->r[1] > 0)) {

    calc_proj_params_r0_r1_high_bd_c(src8, width, height, src_stride, dat8,

                                     dat_stride, flt0, flt0_stride, flt1,

                                     flt1_stride, H, C);

  } else if (params->r[0] > 0) {

    calc_proj_params_r0_high_bd_c(src8, width, height, src_stride, dat8,

                                  dat_stride, flt0, flt0_stride, H, C);

  } else if (params->r[1] > 0) {

    calc_proj_params_r1_high_bd_c(src8, width, height, src_stride, dat8,

                                  dat_stride, flt1, flt1_stride, H, C);

  }

}

static AOM_INLINE void get_proj_subspace(const uint8_t *src8, int width,

                                         int height, int src_stride,

                                         const uint8_t *dat8, int dat_stride,

                                         int use_highbitdepth, int32_t *flt0,

                                         int flt0_stride, int32_t *flt1,

                                         int flt1_stride, int *xq,

                                         const sgr_params_type *params) {

  int64_t H[2][2] = { { 0, 0 }, { 0, 0 } };

  int64_t C[2] = { 0, 0 };

  // Default values to be returned if the problem becomes ill-posed

  xq[0] = 0;

  xq[1] = 0;

  if (!use_highbitdepth) {

    if ((width & 0x7) == 0) {

      av1_calc_proj_params(src8, width, height, src_stride, dat8, dat_stride,

                           flt0, flt0_stride, flt1, flt1_stride, H, C, params);

    } else {

      av1_calc_proj_params_c(src8, width, height, src_stride, dat8, dat_stride,

                             flt0, flt0_stride, flt1, flt1_stride, H, C,

                             params);

    }

  }

#if CONFIG_AV1_HIGHBITDEPTH

  else {  // NOLINT

    if ((width & 0x7) == 0) {

      av1_calc_proj_params_high_bd(src8, width, height, src_stride, dat8,

                                   dat_stride, flt0, flt0_stride, flt1,

                                   flt1_stride, H, C, params);

    } else {

      av1_calc_proj_params_high_bd_c(src8, width, height, src_stride, dat8,

                                     dat_stride, flt0, flt0_stride, flt1,

                                     flt1_stride, H, C, params);

    }

  }

#endif

  if (params->r[0] == 0) {

    // H matrix is now only the scalar H[1][1]

    // C vector is now only the scalar C[1]

    const int64_t Det = H[1][1];

    if (Det == 0) return;  // ill-posed, return default values

    xq[0] = 0;

    xq[1] = (int)signed_rounded_divide(C[1] * (1 << SGRPROJ_PRJ_BITS), Det);

  } else if (params->r[1] == 0) {

    // H matrix is now only the scalar H[0][0]

    // C vector is now only the scalar C[0]

    const int64_t Det = H[0][0];

    if (Det == 0) return;  // ill-posed, return default values

    xq[0] = (int)signed_rounded_divide(C[0] * (1 << SGRPROJ_PRJ_BITS), Det);

    xq[1] = 0;

  } else {

    const int64_t Det = H[0][0] * H[1][1] - H[0][1] * H[1][0];

    if (Det == 0) return;  // ill-posed, return default values

    // If scaling up dividend would overflow, instead scale down the divisor

    const int64_t div1 = H[1][1] * C[0] - H[0][1] * C[1];

    if ((div1 > 0 && INT64_MAX / (1 << SGRPROJ_PRJ_BITS) < div1) ||

        (div1 < 0 && INT64_MIN / (1 << SGRPROJ_PRJ_BITS) > div1))

      xq[0] = (int)signed_rounded_divide(div1, Det / (1 << SGRPROJ_PRJ_BITS));

    else

      xq[0] = (int)signed_rounded_divide(div1 * (1 << SGRPROJ_PRJ_BITS), Det);

    const int64_t div2 = H[0][0] * C[1] - H[1][0] * C[0];

    if ((div2 > 0 && INT64_MAX / (1 << SGRPROJ_PRJ_BITS) < div2) ||

        (div2 < 0 && INT64_MIN / (1 << SGRPROJ_PRJ_BITS) > div2))

      xq[1] = (int)signed_rounded_divide(div2, Det / (1 << SGRPROJ_PRJ_BITS));

    else

      xq[1] = (int)signed_rounded_divide(div2 * (1 << SGRPROJ_PRJ_BITS), Det);

  }

}

static AOM_INLINE void encode_xq(int *xq, int *xqd,

                                 const sgr_params_type *params) {

  if (params->r[0] == 0) {

    xqd[0] = 0;

    xqd[1] = clamp((1 << SGRPROJ_PRJ_BITS) - xq[1], SGRPROJ_PRJ_MIN1,

                   SGRPROJ_PRJ_MAX1);

  } else if (params->r[1] == 0) {

    xqd[0] = clamp(xq[0], SGRPROJ_PRJ_MIN0, SGRPROJ_PRJ_MAX0);

    xqd[1] = clamp((1 << SGRPROJ_PRJ_BITS) - xqd[0], SGRPROJ_PRJ_MIN1,

                   SGRPROJ_PRJ_MAX1);

  } else {

    xqd[0] = clamp(xq[0], SGRPROJ_PRJ_MIN0, SGRPROJ_PRJ_MAX0);

    xqd[1] = clamp((1 << SGRPROJ_PRJ_BITS) - xqd[0] - xq[1], SGRPROJ_PRJ_MIN1,

                   SGRPROJ_PRJ_MAX1);

  }

}

// Apply the self-guided filter across an entire restoration unit.

static AOM_INLINE void apply_sgr(int sgr_params_idx, const uint8_t *dat8,

                                 int width, int height, int dat_stride,

                                 int use_highbd, int bit_depth, int pu_width,

                                 int pu_height, int32_t *flt0, int32_t *flt1,

                                 int flt_stride) {

  for (int i = 0; i < height; i += pu_height) {

    const int h = AOMMIN(pu_height, height - i);

    int32_t *flt0_row = flt0 + i * flt_stride;

    int32_t *flt1_row = flt1 + i * flt_stride;

    const uint8_t *dat8_row = dat8 + i * dat_stride;

    // Iterate over the stripe in blocks of width pu_width

    for (int j = 0; j < width; j += pu_width) {

      const int w = AOMMIN(pu_width, width - j);

      const int ret = av1_selfguided_restoration(

          dat8_row + j, w, h, dat_stride, flt0_row + j, flt1_row + j,

          flt_stride, sgr_params_idx, bit_depth, use_highbd);

      (void)ret;

      assert(!ret);

    }

  }

}

static AOM_INLINE void compute_sgrproj_err(

    const uint8_t *dat8, const int width, const int height,

    const int dat_stride, const uint8_t *src8, const int src_stride,

    const int use_highbitdepth, const int bit_depth, const int pu_width,

    const int pu_height, const int ep, int32_t *flt0, int32_t *flt1,

    const int flt_stride, int *exqd, int64_t *err) {

  int exq[2];

  apply_sgr(ep, dat8, width, height, dat_stride, use_highbitdepth, bit_depth,

            pu_width, pu_height, flt0, flt1, flt_stride);

  const sgr_params_type *const params = &av1_sgr_params[ep];

  get_proj_subspace(src8, width, height, src_stride, dat8, dat_stride,

                    use_highbitdepth, flt0, flt_stride, flt1, flt_stride, exq,

                    params);

  encode_xq(exq, exqd, params);

  *err = finer_search_pixel_proj_error(

      src8, width, height, src_stride, dat8, dat_stride, use_highbitdepth, flt0,

      flt_stride, flt1, flt_stride, 2, exqd, params);

}

static AOM_INLINE void get_best_error(int64_t *besterr, const int64_t err,

                                      const int *exqd, int *bestxqd,

                                      int *bestep, const int ep) {

  if (*besterr == -1 || err < *besterr) {

    *bestep = ep;

    *besterr = err;

    bestxqd[0] = exqd[0];

    bestxqd[1] = exqd[1];

  }

}

static SgrprojInfo search_selfguided_restoration(

    const uint8_t *dat8, int width, int height, int dat_stride,

    const uint8_t *src8, int src_stride, int use_highbitdepth, int bit_depth,

    int pu_width, int pu_height, int32_t *rstbuf, int enable_sgr_ep_pruning) {

  int32_t *flt0 = rstbuf;

  int32_t *flt1 = flt0 + RESTORATION_UNITPELS_MAX;

  int ep, idx, bestep = 0;

  int64_t besterr = -1;

  int exqd[2], bestxqd[2] = { 0, 0 };

  int flt_stride = ((width + 7) & ~7) + 8;

  assert(pu_width == (RESTORATION_PROC_UNIT_SIZE >> 1) ||

         pu_width == RESTORATION_PROC_UNIT_SIZE);

  assert(pu_height == (RESTORATION_PROC_UNIT_SIZE >> 1) ||

         pu_height == RESTORATION_PROC_UNIT_SIZE);

  if (!enable_sgr_ep_pruning) {

    for (ep = 0; ep < SGRPROJ_PARAMS; ep++) {

      int64_t err;

      compute_sgrproj_err(dat8, width, height, dat_stride, src8, src_stride,

                          use_highbitdepth, bit_depth, pu_width, pu_height, ep,

                          flt0, flt1, flt_stride, exqd, &err);

      get_best_error(&besterr, err, exqd, bestxqd, &bestep, ep);

    }

  } else {

    // evaluate first four seed ep in first group

    for (idx = 0; idx < SGRPROJ_EP_GRP1_SEARCH_COUNT; idx++) {

      ep = sgproj_ep_grp1_seed[idx];

      int64_t err;

      compute_sgrproj_err(dat8, width, height, dat_stride, src8, src_stride,

                          use_highbitdepth, bit_depth, pu_width, pu_height, ep,

                          flt0, flt1, flt_stride, exqd, &err);

      get_best_error(&besterr, err, exqd, bestxqd, &bestep, ep);

    }

    // evaluate left and right ep of winner in seed ep

    int bestep_ref = bestep;

    for (ep = bestep_ref - 1; ep < bestep_ref + 2; ep += 2) {

      if (ep < SGRPROJ_EP_GRP1_START_IDX || ep > SGRPROJ_EP_GRP1_END_IDX)

        continue;

      int64_t err;

      compute_sgrproj_err(dat8, width, height, dat_stride, src8, src_stride,

                          use_highbitdepth, bit_depth, pu_width, pu_height, ep,

                          flt0, flt1, flt_stride, exqd, &err);

      get_best_error(&besterr, err, exqd, bestxqd, &bestep, ep);

    }

    // evaluate last two group

    for (idx = 0; idx < SGRPROJ_EP_GRP2_3_SEARCH_COUNT; idx++) {

      ep = sgproj_ep_grp2_3[idx][bestep];

      int64_t err;

      compute_sgrproj_err(dat8, width, height, dat_stride, src8, src_stride,

                          use_highbitdepth, bit_depth, pu_width, pu_height, ep,

                          flt0, flt1, flt_stride, exqd, &err);

      get_best_error(&besterr, err, exqd, bestxqd, &bestep, ep);

    }

  }

  SgrprojInfo ret;

  ret.ep = bestep;

  ret.xqd[0] = bestxqd[0];

  ret.xqd[1] = bestxqd[1];

  return ret;

}

static int count_sgrproj_bits(SgrprojInfo *sgrproj_info,

                              SgrprojInfo *ref_sgrproj_info) {

  int bits = SGRPROJ_PARAMS_BITS;

  const sgr_params_type *params = &av1_sgr_params[sgrproj_info->ep];

  if (params->r[0] > 0)

    bits += aom_count_primitive_refsubexpfin(

        SGRPROJ_PRJ_MAX0 - SGRPROJ_PRJ_MIN0 + 1, SGRPROJ_PRJ_SUBEXP_K,

        ref_sgrproj_info->xqd[0] - SGRPROJ_PRJ_MIN0,

        sgrproj_info->xqd[0] - SGRPROJ_PRJ_MIN0);

  if (params->r[1] > 0)

    bits += aom_count_primitive_refsubexpfin(

        SGRPROJ_PRJ_MAX1 - SGRPROJ_PRJ_MIN1 + 1, SGRPROJ_PRJ_SUBEXP_K,

        ref_sgrproj_info->xqd[1] - SGRPROJ_PRJ_MIN1,

        sgrproj_info->xqd[1] - SGRPROJ_PRJ_MIN1);

  return bits;

}

static AOM_INLINE void search_sgrproj(const RestorationTileLimits *limits,

                                      const AV1PixelRect *tile,

                                      int rest_unit_idx, void *priv,

                                      int32_t *tmpbuf,

                                      RestorationLineBuffers *rlbs) {

  (void)rlbs;

  RestSearchCtxt *rsc = (RestSearchCtxt *)priv;

  RestUnitSearchInfo *rusi = &rsc->rusi[rest_unit_idx];

  const MACROBLOCK *const x = rsc->x;

  const AV1_COMMON *const cm = rsc->cm;