/*

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

#include <stdio.h>

#include <stdlib.h>

#include <string.h>

#include "aom_dsp/aom_dsp_common.h"

#include "aom_dsp/mathutils.h"

#include "aom_dsp/noise_model.h"

#include "aom_dsp/noise_util.h"

#include "aom_mem/aom_mem.h"

#include "aom_ports/mem.h"

#include "aom_scale/yv12config.h"

#define kLowPolyNumParams 3

static const int kMaxLag = 4;

// Defines a function that can be used to obtain the mean of a block for the

// provided data type (uint8_t, or uint16_t)

#define GET_BLOCK_MEAN(INT_TYPE, suffix)                                    \

  static double get_block_mean_##suffix(const INT_TYPE *data, int w, int h, \

                                        int stride, int x_o, int y_o,       \

                                        int block_size) {                   \

    const int max_h = AOMMIN(h - y_o, block_size);                          \

    const int max_w = AOMMIN(w - x_o, block_size);                          \

    double block_mean = 0;                                                  \

    for (int y = 0; y < max_h; ++y) {                                       \

      for (int x = 0; x < max_w; ++x) {                                     \

        block_mean += data[(y_o + y) * stride + x_o + x];                   \

      }                                                                     \

    }                                                                       \

    return block_mean / (max_w * max_h);                                    \

  }

Unexecuted instantiation: noise_model.c:get_block_mean_highbd

Unexecuted instantiation: noise_model.c:get_block_mean_lowbd

GET_BLOCK_MEAN(uint8_t, lowbd)

GET_BLOCK_MEAN(uint16_t, highbd)

static inline double get_block_mean(const uint8_t *data, int w, int h,

                                    int stride, int x_o, int y_o,

                                    int block_size, int use_highbd) {

  if (use_highbd)

    return get_block_mean_highbd((const uint16_t *)data, w, h, stride, x_o, y_o,

                                 block_size);

  return get_block_mean_lowbd(data, w, h, stride, x_o, y_o, block_size);

}

// Defines a function that can be used to obtain the variance of a block

// for the provided data type (uint8_t, or uint16_t)

#define GET_NOISE_VAR(INT_TYPE, suffix)                                  \

  static double get_noise_var_##suffix(                                  \

      const INT_TYPE *data, const INT_TYPE *denoised, int stride, int w, \

      int h, int x_o, int y_o, int block_size_x, int block_size_y) {     \

    const int max_h = AOMMIN(h - y_o, block_size_y);                     \

    const int max_w = AOMMIN(w - x_o, block_size_x);                     \

    double noise_var = 0;                                                \

    double noise_mean = 0;                                               \

    for (int y = 0; y < max_h; ++y) {                                    \

      for (int x = 0; x < max_w; ++x) {                                  \

        double noise = (double)data[(y_o + y) * stride + x_o + x] -      \

                       denoised[(y_o + y) * stride + x_o + x];           \

        noise_mean += noise;                                             \

        noise_var += noise * noise;                                      \

      }                                                                  \

    }                                                                    \

    noise_mean /= (max_w * max_h);                                       \

    return noise_var / (max_w * max_h) - noise_mean * noise_mean;        \

  }

Unexecuted instantiation: noise_model.c:get_noise_var_highbd

Unexecuted instantiation: noise_model.c:get_noise_var_lowbd

GET_NOISE_VAR(uint8_t, lowbd)

GET_NOISE_VAR(uint16_t, highbd)

static inline double get_noise_var(const uint8_t *data, const uint8_t *denoised,

                                   int w, int h, int stride, int x_o, int y_o,

                                   int block_size_x, int block_size_y,

                                   int use_highbd) {

  if (use_highbd)

    return get_noise_var_highbd((const uint16_t *)data,

                                (const uint16_t *)denoised, w, h, stride, x_o,

                                y_o, block_size_x, block_size_y);

  return get_noise_var_lowbd(data, denoised, w, h, stride, x_o, y_o,

                             block_size_x, block_size_y);

}

static void equation_system_clear(aom_equation_system_t *eqns) {

  const int n = eqns->n;

  memset(eqns->A, 0, sizeof(*eqns->A) * n * n);

  memset(eqns->x, 0, sizeof(*eqns->x) * n);

  memset(eqns->b, 0, sizeof(*eqns->b) * n);

}

static void equation_system_copy(aom_equation_system_t *dst,

                                 const aom_equation_system_t *src) {

  const int n = dst->n;

  memcpy(dst->A, src->A, sizeof(*dst->A) * n * n);

  memcpy(dst->x, src->x, sizeof(*dst->x) * n);

  memcpy(dst->b, src->b, sizeof(*dst->b) * n);

}

static int equation_system_init(aom_equation_system_t *eqns, int n) {

  eqns->A = (double *)aom_malloc(sizeof(*eqns->A) * n * n);

  eqns->b = (double *)aom_malloc(sizeof(*eqns->b) * n);

  eqns->x = (double *)aom_malloc(sizeof(*eqns->x) * n);

  eqns->n = n;

  if (!eqns->A || !eqns->b || !eqns->x) {

    fprintf(stderr, "Failed to allocate system of equations of size %d\n", n);

    aom_free(eqns->A);

    aom_free(eqns->b);

    aom_free(eqns->x);

    memset(eqns, 0, sizeof(*eqns));

    return 0;

  }

  equation_system_clear(eqns);

  return 1;

}

static int equation_system_solve(aom_equation_system_t *eqns) {

  const int n = eqns->n;

  double *b = (double *)aom_malloc(sizeof(*b) * n);

  double *A = (double *)aom_malloc(sizeof(*A) * n * n);

  int ret = 0;

  if (A == NULL || b == NULL) {

    fprintf(stderr, "Unable to allocate temp values of size %dx%d\n", n, n);

    aom_free(b);

    aom_free(A);

    return 0;

  }

  memcpy(A, eqns->A, sizeof(*eqns->A) * n * n);

  memcpy(b, eqns->b, sizeof(*eqns->b) * n);

  ret = linsolve(n, A, eqns->n, b, eqns->x);

  aom_free(b);

  aom_free(A);

  if (ret == 0) {

    return 0;

  }

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

    if (isnan((float)eqns->x[i])) {

      return 0;

    }

  }

  return 1;

}

static void equation_system_add(aom_equation_system_t *dest,

                                aom_equation_system_t *src) {

  const int n = dest->n;

  int i, j;

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

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

      dest->A[i * n + j] += src->A[i * n + j];

    }

    dest->b[i] += src->b[i];

  }

}

static void equation_system_free(aom_equation_system_t *eqns) {

  if (!eqns) return;

  aom_free(eqns->A);

  aom_free(eqns->b);

  aom_free(eqns->x);

  memset(eqns, 0, sizeof(*eqns));

}

static void noise_strength_solver_clear(aom_noise_strength_solver_t *solver) {

  equation_system_clear(&solver->eqns);

  solver->num_equations = 0;

  solver->total = 0;

}

static void noise_strength_solver_add(aom_noise_strength_solver_t *dest,

                                      aom_noise_strength_solver_t *src) {

  equation_system_add(&dest->eqns, &src->eqns);

  dest->num_equations += src->num_equations;

  dest->total += src->total;

}

// Return the number of coefficients required for the given parameters

static int num_coeffs(const aom_noise_model_params_t params) {

  const int n = 2 * params.lag + 1;

  switch (params.shape) {

    case AOM_NOISE_SHAPE_DIAMOND: return params.lag * (params.lag + 1);

    case AOM_NOISE_SHAPE_SQUARE: return (n * n) / 2;

  }

  return 0;

}

static int noise_state_init(aom_noise_state_t *state, int n, int bit_depth) {

  const int kNumBins = 20;

  if (!equation_system_init(&state->eqns, n)) {

    fprintf(stderr, "Failed initialization noise state with size %d\n", n);

    return 0;

  }

  state->ar_gain = 1.0;

  state->num_observations = 0;

  return aom_noise_strength_solver_init(&state->strength_solver, kNumBins,

                                        bit_depth);

}

static void set_chroma_coefficient_fallback_soln(aom_equation_system_t *eqns) {

  const double kTolerance = 1e-6;

  const int last = eqns->n - 1;

  // Set all of the AR coefficients to zero, but try to solve for correlation

  // with the luma channel

  memset(eqns->x, 0, sizeof(*eqns->x) * eqns->n);

  if (fabs(eqns->A[last * eqns->n + last]) > kTolerance) {

    eqns->x[last] = eqns->b[last] / eqns->A[last * eqns->n + last];

  }

}

int aom_noise_strength_lut_init(aom_noise_strength_lut_t *lut, int num_points) {

  if (!lut) return 0;

  if (num_points <= 0) return 0;

  lut->num_points = 0;

  lut->points = (double(*)[2])aom_malloc(num_points * sizeof(*lut->points));

  if (!lut->points) return 0;

  lut->num_points = num_points;

  memset(lut->points, 0, sizeof(*lut->points) * num_points);

  return 1;

}

void aom_noise_strength_lut_free(aom_noise_strength_lut_t *lut) {

  if (!lut) return;

  aom_free(lut->points);

  memset(lut, 0, sizeof(*lut));

}

double aom_noise_strength_lut_eval(const aom_noise_strength_lut_t *lut,

                                   double x) {

  int i = 0;

  // Constant extrapolation for x <  x_0.

  if (x < lut->points[0][0]) return lut->points[0][1];

  for (i = 0; i < lut->num_points - 1; ++i) {

    if (x >= lut->points[i][0] && x <= lut->points[i + 1][0]) {

      const double a =

          (x - lut->points[i][0]) / (lut->points[i + 1][0] - lut->points[i][0]);

      return lut->points[i + 1][1] * a + lut->points[i][1] * (1.0 - a);

    }

  }

  // Constant extrapolation for x > x_{n-1}

  return lut->points[lut->num_points - 1][1];

}

static double noise_strength_solver_get_bin_index(

    const aom_noise_strength_solver_t *solver, double value) {

  const double val =

      fclamp(value, solver->min_intensity, solver->max_intensity);

  const double range = solver->max_intensity - solver->min_intensity;

  return (solver->num_bins - 1) * (val - solver->min_intensity) / range;

}

static double noise_strength_solver_get_value(

    const aom_noise_strength_solver_t *solver, double x) {

  const double bin = noise_strength_solver_get_bin_index(solver, x);

  const int bin_i0 = (int)floor(bin);

  const int bin_i1 = AOMMIN(solver->num_bins - 1, bin_i0 + 1);

  const double a = bin - bin_i0;

  return (1.0 - a) * solver->eqns.x[bin_i0] + a * solver->eqns.x[bin_i1];

}

void aom_noise_strength_solver_add_measurement(

    aom_noise_strength_solver_t *solver, double block_mean, double noise_std) {

  const double bin = noise_strength_solver_get_bin_index(solver, block_mean);

  const int bin_i0 = (int)floor(bin);

  const int bin_i1 = AOMMIN(solver->num_bins - 1, bin_i0 + 1);

  const double a = bin - bin_i0;

  const int n = solver->num_bins;

  solver->eqns.A[bin_i0 * n + bin_i0] += (1.0 - a) * (1.0 - a);

  solver->eqns.A[bin_i1 * n + bin_i0] += a * (1.0 - a);

  solver->eqns.A[bin_i1 * n + bin_i1] += a * a;

  solver->eqns.A[bin_i0 * n + bin_i1] += a * (1.0 - a);

  solver->eqns.b[bin_i0] += (1.0 - a) * noise_std;

  solver->eqns.b[bin_i1] += a * noise_std;

  solver->total += noise_std;

  solver->num_equations++;

}

int aom_noise_strength_solver_solve(aom_noise_strength_solver_t *solver) {

  // Add regularization proportional to the number of constraints

  const int n = solver->num_bins;

  const double kAlpha = 2.0 * (double)(solver->num_equations) / n;

  int result = 0;

  double mean = 0;

  // Do this in a non-destructive manner so it is not confusing to the caller

  double *old_A = solver->eqns.A;

  double *A = (double *)aom_malloc(sizeof(*A) * n * n);

  if (!A) {

    fprintf(stderr, "Unable to allocate copy of A\n");

    return 0;

  }

  memcpy(A, old_A, sizeof(*A) * n * n);

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

    const int i_lo = AOMMAX(0, i - 1);

    const int i_hi = AOMMIN(n - 1, i + 1);

    A[i * n + i_lo] -= kAlpha;

    A[i * n + i] += 2 * kAlpha;

    A[i * n + i_hi] -= kAlpha;

  }

  // Small regularization to give average noise strength

  mean = solver->total / solver->num_equations;

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

    A[i * n + i] += 1.0 / 8192.;

    solver->eqns.b[i] += mean / 8192.;

  }

  solver->eqns.A = A;

  result = equation_system_solve(&solver->eqns);

  solver->eqns.A = old_A;

  aom_free(A);

  return result;

}

int aom_noise_strength_solver_init(aom_noise_strength_solver_t *solver,

                                   int num_bins, int bit_depth) {

  if (!solver) return 0;

  memset(solver, 0, sizeof(*solver));

  solver->num_bins = num_bins;

  solver->min_intensity = 0;

  solver->max_intensity = (1 << bit_depth) - 1;

  solver->total = 0;

  solver->num_equations = 0;

  return equation_system_init(&solver->eqns, num_bins);

}

void aom_noise_strength_solver_free(aom_noise_strength_solver_t *solver) {

  if (!solver) return;

  equation_system_free(&solver->eqns);

}

double aom_noise_strength_solver_get_center(

    const aom_noise_strength_solver_t *solver, int i) {

  const double range = solver->max_intensity - solver->min_intensity;

  const int n = solver->num_bins;

  return ((double)i) / (n - 1) * range + solver->min_intensity;

}

// Computes the residual if a point were to be removed from the lut. This is

// calculated as the area between the output of the solver and the line segment

// that would be formed between [x_{i - 1}, x_{i + 1}).

static void update_piecewise_linear_residual(

    const aom_noise_strength_solver_t *solver,

    const aom_noise_strength_lut_t *lut, double *residual, int start, int end) {

  const double dx = 255. / solver->num_bins;

  for (int i = AOMMAX(start, 1); i < AOMMIN(end, lut->num_points - 1); ++i) {

    const int lower = AOMMAX(0, (int)floor(noise_strength_solver_get_bin_index(

                                    solver, lut->points[i - 1][0])));

    const int upper = AOMMIN(solver->num_bins - 1,

                             (int)ceil(noise_strength_solver_get_bin_index(

                                 solver, lut->points[i + 1][0])));

    double r = 0;

    for (int j = lower; j <= upper; ++j) {

      const double x = aom_noise_strength_solver_get_center(solver, j);

      if (x < lut->points[i - 1][0]) continue;

      if (x >= lut->points[i + 1][0]) continue;

      const double y = solver->eqns.x[j];

      const double a = (x - lut->points[i - 1][0]) /

                       (lut->points[i + 1][0] - lut->points[i - 1][0]);

      const double estimate_y =

          lut->points[i - 1][1] * (1.0 - a) + lut->points[i + 1][1] * a;

      r += fabs(y - estimate_y);

    }

    residual[i] = r * dx;

  }

}

int aom_noise_strength_solver_fit_piecewise(

    const aom_noise_strength_solver_t *solver, int max_output_points,

    aom_noise_strength_lut_t *lut) {

  // The tolerance is normalized to be give consistent results between

  // different bit-depths.

  const double kTolerance = solver->max_intensity * 0.00625 / 255.0;

  if (!aom_noise_strength_lut_init(lut, solver->num_bins)) {

    fprintf(stderr, "Failed to init lut\n");

    return 0;

  }

  for (int i = 0; i < solver->num_bins; ++i) {

    lut->points[i][0] = aom_noise_strength_solver_get_center(solver, i);

    lut->points[i][1] = solver->eqns.x[i];

  }

  if (max_output_points < 0) {

    max_output_points = solver->num_bins;

  }

  double *residual = (double *)aom_malloc(solver->num_bins * sizeof(*residual));

  if (!residual) {

    aom_noise_strength_lut_free(lut);

    return 0;

  }

  memset(residual, 0, sizeof(*residual) * solver->num_bins);

  update_piecewise_linear_residual(solver, lut, residual, 0, solver->num_bins);

  // Greedily remove points if there are too many or if it doesn't hurt local

  // approximation (never remove the end points)

  while (lut->num_points > 2) {

    int min_index = 1;

    for (int j = 1; j < lut->num_points - 1; ++j) {

      if (residual[j] < residual[min_index]) {

        min_index = j;

      }

    }

    const double dx =

        lut->points[min_index + 1][0] - lut->points[min_index - 1][0];

    const double avg_residual = residual[min_index] / dx;

    if (lut->num_points <= max_output_points && avg_residual > kTolerance) {

      break;

    }

    const int num_remaining = lut->num_points - min_index - 1;

    memmove(lut->points + min_index, lut->points + min_index + 1,

            sizeof(lut->points[0]) * num_remaining);

    lut->num_points--;

    update_piecewise_linear_residual(solver, lut, residual, min_index - 1,

                                     min_index + 1);

  }

  aom_free(residual);

  return 1;

}

int aom_flat_block_finder_init(aom_flat_block_finder_t *block_finder,

                               int block_size, int bit_depth, int use_highbd) {

  const int n = block_size * block_size;

  aom_equation_system_t eqns;

  double *AtA_inv = 0;

  double *A = 0;

  int x = 0, y = 0, i = 0, j = 0;

  block_finder->A = NULL;

  block_finder->AtA_inv = NULL;

  if (!equation_system_init(&eqns, kLowPolyNumParams)) {

    fprintf(stderr, "Failed to init equation system for block_size=%d\n",

            block_size);

    return 0;

  }

  AtA_inv = (double *)aom_malloc(kLowPolyNumParams * kLowPolyNumParams *

                                 sizeof(*AtA_inv));

  A = (double *)aom_malloc(kLowPolyNumParams * n * sizeof(*A));

  if (AtA_inv == NULL || A == NULL) {

    fprintf(stderr, "Failed to alloc A or AtA_inv for block_size=%d\n",

            block_size);

    aom_free(AtA_inv);

    aom_free(A);

    equation_system_free(&eqns);

    return 0;

  }

  block_finder->A = A;

  block_finder->AtA_inv = AtA_inv;

  block_finder->block_size = block_size;

  block_finder->normalization = (1 << bit_depth) - 1;

  block_finder->use_highbd = use_highbd;

  for (y = 0; y < block_size; ++y) {

    const double yd = ((double)y - block_size / 2.) / (block_size / 2.);

    for (x = 0; x < block_size; ++x) {

      const double xd = ((double)x - block_size / 2.) / (block_size / 2.);

      const double coords[3] = { yd, xd, 1 };

      const int row = y * block_size + x;

      A[kLowPolyNumParams * row + 0] = yd;

      A[kLowPolyNumParams * row + 1] = xd;

      A[kLowPolyNumParams * row + 2] = 1;

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

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

          eqns.A[kLowPolyNumParams * i + j] += coords[i] * coords[j];

        }

      }

    }

  }

  // Lazy inverse using existing equation solver.

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

    memset(eqns.b, 0, sizeof(*eqns.b) * kLowPolyNumParams);

    eqns.b[i] = 1;

    const int ret = equation_system_solve(&eqns);

    if (!ret) return ret;

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

      AtA_inv[j * kLowPolyNumParams + i] = eqns.x[j];

    }

  }

  equation_system_free(&eqns);

  return 1;

}

void aom_flat_block_finder_free(aom_flat_block_finder_t *block_finder) {

  if (!block_finder) return;

  aom_free(block_finder->A);

  aom_free(block_finder->AtA_inv);

  memset(block_finder, 0, sizeof(*block_finder));

}

void aom_flat_block_finder_extract_block(

    const aom_flat_block_finder_t *block_finder, const uint8_t *const data,

    int w, int h, int stride, int offsx, int offsy, double *plane,

    double *block) {

  const int block_size = block_finder->block_size;

  const int n = block_size * block_size;

  const double *A = block_finder->A;

  const double *AtA_inv = block_finder->AtA_inv;

  double plane_coords[kLowPolyNumParams];

  double AtA_inv_b[kLowPolyNumParams];

  int xi, yi, i;

  if (block_finder->use_highbd) {

    const uint16_t *const data16 = (const uint16_t *const)data;

    for (yi = 0; yi < block_size; ++yi) {

      const int y = clamp(offsy + yi, 0, h - 1);

      for (xi = 0; xi < block_size; ++xi) {

        const int x = clamp(offsx + xi, 0, w - 1);

        block[yi * block_size + xi] =

            ((double)data16[y * stride + x]) / block_finder->normalization;

      }

    }

  } else {

    for (yi = 0; yi < block_size; ++yi) {

      const int y = clamp(offsy + yi, 0, h - 1);

      for (xi = 0; xi < block_size; ++xi) {

        const int x = clamp(offsx + xi, 0, w - 1);

        block[yi * block_size + xi] =

            ((double)data[y * stride + x]) / block_finder->normalization;

      }

    }

  }

  multiply_mat(block, A, AtA_inv_b, 1, n, kLowPolyNumParams);

  multiply_mat(AtA_inv, AtA_inv_b, plane_coords, kLowPolyNumParams,

               kLowPolyNumParams, 1);

  multiply_mat(A, plane_coords, plane, n, kLowPolyNumParams, 1);

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

    block[i] -= plane[i];

  }

}

typedef struct {

  int index;

  float score;

} index_and_score_t;

static int compare_scores(const void *a, const void *b) {

  const float diff =

      ((index_and_score_t *)a)->score - ((index_and_score_t *)b)->score;

  if (diff < 0)

    return -1;

  else if (diff > 0)

    return 1;

  return 0;

}

int aom_flat_block_finder_run(const aom_flat_block_finder_t *block_finder,

                              const uint8_t *const data, int w, int h,

                              int stride, uint8_t *flat_blocks) {

  // The gradient-based features used in this code are based on:

  //  A. Kokaram, D. Kelly, H. Denman and A. Crawford, "Measuring noise

  //  correlation for improved video denoising," 2012 19th, ICIP.

  // The thresholds are more lenient to allow for correct grain modeling

  // if extreme cases.

  const int block_size = block_finder->block_size;

  const int n = block_size * block_size;

  const double kTraceThreshold = 0.15 / (32 * 32);

  const double kRatioThreshold = 1.25;

  const double kNormThreshold = 0.08 / (32 * 32);

  const double kVarThreshold = 0.005 / (double)n;

  const int num_blocks_w = (w + block_size - 1) / block_size;

  const int num_blocks_h = (h + block_size - 1) / block_size;

  int num_flat = 0;

  double *plane = (double *)aom_malloc(n * sizeof(*plane));

  double *block = (double *)aom_malloc(n * sizeof(*block));

  index_and_score_t *scores = (index_and_score_t *)aom_malloc(

      num_blocks_w * num_blocks_h * sizeof(*scores));

  if (plane == NULL || block == NULL || scores == NULL) {

    fprintf(stderr, "Failed to allocate memory for block of size %d\n", n);

    aom_free(plane);

    aom_free(block);

    aom_free(scores);

    return -1;

  }

#ifdef NOISE_MODEL_LOG_SCORE

  fprintf(stderr, "score = [");

#endif

  for (int by = 0; by < num_blocks_h; ++by) {

    for (int bx = 0; bx < num_blocks_w; ++bx) {

      // Compute gradient covariance matrix.

      aom_flat_block_finder_extract_block(block_finder, data, w, h, stride,

                                          bx * block_size, by * block_size,

                                          plane, block);

      double Gxx = 0, Gxy = 0, Gyy = 0;

      double mean = 0;

      double var = 0;

      for (int yi = 1; yi < block_size - 1; ++yi) {

        for (int xi = 1; xi < block_size - 1; ++xi) {

          const double gx = (block[yi * block_size + xi + 1] -

                             block[yi * block_size + xi - 1]) /

                            2;

          const double gy = (block[yi * block_size + xi + block_size] -

                             block[yi * block_size + xi - block_size]) /

                            2;

          Gxx += gx * gx;

          Gxy += gx * gy;

          Gyy += gy * gy;

          const double value = block[yi * block_size + xi];

          mean += value;

          var += value * value;

        }

      }

      mean /= (block_size - 2) * (block_size - 2);

      // Normalize gradients by block_size.

      Gxx /= ((block_size - 2) * (block_size - 2));

      Gxy /= ((block_size - 2) * (block_size - 2));

      Gyy /= ((block_size - 2) * (block_size - 2));

      var = var / ((block_size - 2) * (block_size - 2)) - mean * mean;

      {

        const double trace = Gxx + Gyy;

        const double det = Gxx * Gyy - Gxy * Gxy;

        const double e1 = (trace + sqrt(trace * trace - 4 * det)) / 2.;

        const double e2 = (trace - sqrt(trace * trace - 4 * det)) / 2.;

        const double norm = e1;  // Spectral norm

        const double ratio = (e1 / AOMMAX(e2, 1e-6));

        const int is_flat = (trace < kTraceThreshold) &&

                            (ratio < kRatioThreshold) &&

                            (norm < kNormThreshold) && (var > kVarThreshold);

        // The following weights are used to combine the above features to give

        // a sigmoid score for flatness. If the input was normalized to [0,100]

        // the magnitude of these values would be close to 1 (e.g., weights

        // corresponding to variance would be a factor of 10000x smaller).

        // The weights are given in the following order:

        //    [{var}, {ratio}, {trace}, {norm}, offset]

        // with one of the most discriminative being simply the variance.

        const double weights[5] = { -6682, -0.2056, 13087, -12434, 2.5694 };

        double sum_weights = weights[0] * var + weights[1] * ratio +

                             weights[2] * trace + weights[3] * norm +

                             weights[4];

        // clamp the value to [-25.0, 100.0] to prevent overflow

        sum_weights = fclamp(sum_weights, -25.0, 100.0);

        const float score = (float)(1.0 / (1 + exp(-sum_weights)));

        flat_blocks[by * num_blocks_w + bx] = is_flat ? 255 : 0;

        scores[by * num_blocks_w + bx].score = var > kVarThreshold ? score : 0;

        scores[by * num_blocks_w + bx].index = by * num_blocks_w + bx;

#ifdef NOISE_MODEL_LOG_SCORE

        fprintf(stderr, "%g %g %g %g %g %d ", score, var, ratio, trace, norm,

                is_flat);

#endif

        num_flat += is_flat;

      }

    }

#ifdef NOISE_MODEL_LOG_SCORE

    fprintf(stderr, "\n");

#endif

  }

#ifdef NOISE_MODEL_LOG_SCORE

  fprintf(stderr, "];\n");

#endif

  // Find the top-scored blocks (most likely to be flat) and set the flat blocks

  // be the union of the thresholded results and the top 10th percentile of the

  // scored results.

  qsort(scores, num_blocks_w * num_blocks_h, sizeof(*scores), &compare_scores);

  const int top_nth_percentile = num_blocks_w * num_blocks_h * 90 / 100;

  const float score_threshold = scores[top_nth_percentile].score;

  for (int i = 0; i < num_blocks_w * num_blocks_h; ++i) {

    if (scores[i].score >= score_threshold) {

      num_flat += flat_blocks[scores[i].index] == 0;

      flat_blocks[scores[i].index] |= 1;

    }

  }

  aom_free(block);

  aom_free(plane);

  aom_free(scores);

  return num_flat;

}

int aom_noise_model_init(aom_noise_model_t *model,

                         const aom_noise_model_params_t params) {

  const int n = num_coeffs(params);

  const int lag = params.lag;

  const int bit_depth = params.bit_depth;

  int x = 0, y = 0, i = 0, c = 0;

  memset(model, 0, sizeof(*model));

  if (params.lag < 1) {

    fprintf(stderr, "Invalid noise param: lag = %d must be >= 1\n", params.lag);

    return 0;

  }

  if (params.lag > kMaxLag) {

    fprintf(stderr, "Invalid noise param: lag = %d must be <= %d\n", params.lag,

            kMaxLag);

    return 0;

  }

  if (!(params.bit_depth == 8 || params.bit_depth == 10 ||

        params.bit_depth == 12)) {

    return 0;

  }

  model->params = params;

  for (c = 0; c < 3; ++c) {

    if (!noise_state_init(&model->combined_state[c], n + (c > 0), bit_depth)) {

      fprintf(stderr, "Failed to allocate noise state for channel %d\n", c);

      aom_noise_model_free(model);

      return 0;

    }

    if (!noise_state_init(&model->latest_state[c], n + (c > 0), bit_depth)) {

      fprintf(stderr, "Failed to allocate noise state for channel %d\n", c);

      aom_noise_model_free(model);

      return 0;

    }

  }

  model->n = n;

  model->coords = (int(*)[2])aom_malloc(sizeof(*model->coords) * n);

  if (!model->coords) {

    aom_noise_model_free(model);

    return 0;

  }

  for (y = -lag; y <= 0; ++y) {

    const int max_x = y == 0 ? -1 : lag;

    for (x = -lag; x <= max_x; ++x) {

      switch (params.shape) {

        case AOM_NOISE_SHAPE_DIAMOND:

          if (abs(x) <= y + lag) {

            model->coords[i][0] = x;

            model->coords[i][1] = y;

            ++i;

          }

          break;

        case AOM_NOISE_SHAPE_SQUARE:

          model->coords[i][0] = x;

          model->coords[i][1] = y;

          ++i;

          break;

        default:

          fprintf(stderr, "Invalid shape\n");

          aom_noise_model_free(model);

          return 0;

      }

    }

  }

  assert(i == n);

  return 1;

}

void aom_noise_model_free(aom_noise_model_t *model) {

  int c = 0;

  if (!model) return;

  aom_free(model->coords);

  for (c = 0; c < 3; ++c) {

    equation_system_free(&model->latest_state[c].eqns);

    equation_system_free(&model->combined_state[c].eqns);

    equation_system_free(&model->latest_state[c].strength_solver.eqns);

    equation_system_free(&model->combined_state[c].strength_solver.eqns);

  }

  memset(model, 0, sizeof(*model));

}

// Extracts the neighborhood defined by coords around point (x, y) from

// the difference between the data and denoised images. Also extracts the

// entry (possibly downsampled) for (x, y) in the alt_data (e.g., luma).

#define EXTRACT_AR_ROW(INT_TYPE, suffix)                                   \

  static double extract_ar_row_##suffix(                                   \

      int(*coords)[2], int num_coords, const INT_TYPE *const data,         \

      const INT_TYPE *const denoised, int stride, int sub_log2[2],         \

      const INT_TYPE *const alt_data, const INT_TYPE *const alt_denoised,  \

      int alt_stride, int x, int y, double *buffer) {                      \

    for (int i = 0; i < num_coords; ++i) {                                 \

      const int x_i = x + coords[i][0], y_i = y + coords[i][1];            \

      buffer[i] =                                                          \

          (double)data[y_i * stride + x_i] - denoised[y_i * stride + x_i]; \

    }                                                                      \

    const double val =                                                     \

        (double)data[y * stride + x] - denoised[y * stride + x];           \

                                                                           \

    if (alt_data && alt_denoised) {                                        \

      double avg_data = 0, avg_denoised = 0;                               \

      int num_samples = 0;                                                 \

      for (int dy_i = 0; dy_i < (1 << sub_log2[1]); dy_i++) {              \

        const int y_up = (y << sub_log2[1]) + dy_i;                        \

        for (int dx_i = 0; dx_i < (1 << sub_log2[0]); dx_i++) {            \

          const int x_up = (x << sub_log2[0]) + dx_i;                      \

          avg_data += alt_data[y_up * alt_stride + x_up];                  \

          avg_denoised += alt_denoised[y_up * alt_stride + x_up];          \

          num_samples++;                                                   \

        }                                                                  \

      }                                                                    \

      buffer[num_coords] = (avg_data - avg_denoised) / num_samples;        \

    }                                                                      \

    return val;                                                            \

  }

Unexecuted instantiation: noise_model.c:extract_ar_row_highbd

Unexecuted instantiation: noise_model.c:extract_ar_row_lowbd

EXTRACT_AR_ROW(uint8_t, lowbd)

EXTRACT_AR_ROW(uint16_t, highbd)

static int add_block_observations(

    aom_noise_model_t *noise_model, int c, const uint8_t *const data,

    const uint8_t *const denoised, int w, int h, int stride, int sub_log2[2],

    const uint8_t *const alt_data, const uint8_t *const alt_denoised,

    int alt_stride, const uint8_t *const flat_blocks, int block_size,

    int num_blocks_w, int num_blocks_h) {

  const int lag = noise_model->params.lag;

  const int num_coords = noise_model->n;

  const double normalization = (1 << noise_model->params.bit_depth) - 1;

  double *A = noise_model->latest_state[c].eqns.A;

  double *b = noise_model->latest_state[c].eqns.b;

  double *buffer = (double *)aom_malloc(sizeof(*buffer) * (num_coords + 1));

  const int n = noise_model->latest_state[c].eqns.n;

  if (!buffer) {

    fprintf(stderr, "Unable to allocate buffer of size %d\n", num_coords + 1);

    return 0;

  }

  for (int by = 0; by < num_blocks_h; ++by) {

    const int y_o = by * (block_size >> sub_log2[1]);

    for (int bx = 0; bx < num_blocks_w; ++bx) {

      const int x_o = bx * (block_size >> sub_log2[0]);

      if (!flat_blocks[by * num_blocks_w + bx]) {

        continue;

      }

      int y_start =

          (by > 0 && flat_blocks[(by - 1) * num_blocks_w + bx]) ? 0 : lag;

      int x_start =

          (bx > 0 && flat_blocks[by * num_blocks_w + bx - 1]) ? 0 : lag;

      int y_end = AOMMIN((h >> sub_log2[1]) - by * (block_size >> sub_log2[1]),

                         block_size >> sub_log2[1]);

      int x_end = AOMMIN(

          (w >> sub_log2[0]) - bx * (block_size >> sub_log2[0]) - lag,

          (bx + 1 < num_blocks_w && flat_blocks[by * num_blocks_w + bx + 1])

              ? (block_size >> sub_log2[0])

              : ((block_size >> sub_log2[0]) - lag));

      for (int y = y_start; y < y_end; ++y) {

        for (int x = x_start; x < x_end; ++x) {

          const double val =

              noise_model->params.use_highbd

                  ? extract_ar_row_highbd(noise_model->coords, num_coords,

                                          (const uint16_t *const)data,

                                          (const uint16_t *const)denoised,

                                          stride, sub_log2,

                                          (const uint16_t *const)alt_data,

                                          (const uint16_t *const)alt_denoised,

                                          alt_stride, x + x_o, y + y_o, buffer)

                  : extract_ar_row_lowbd(noise_model->coords, num_coords, data,

                                         denoised, stride, sub_log2, alt_data,