/*

 * 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 https://www.aomedia.org/license/software-license. 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 https://www.aomedia.org/license/patent-license.

 */

#include "EbConfigMacros.h"

#if CONFIG_ENABLE_FILM_GRAIN

#include <math.h>

#include <stdio.h>

#include <stdlib.h>

#include "noise_model.h"

#include "definitions.h"

#include "sequence_control_set.h"

#include "noise_util.h"

#include "mathutils.h"

#include "svt_log.h"

#include "aom_dsp_rtcd.h"

#include "pic_operators.h"

static const int32_t k_max_lag = 4;

void svt_aom_un_pack2d(uint16_t* in16_bit_buffer, uint32_t in_stride, uint8_t* out8_bit_buffer, uint32_t out8_stride,

                       uint8_t* outn_bit_buffer, uint32_t outn_stride, uint32_t width, uint32_t height);

void svt_aom_pack2d_src(uint8_t* in8_bit_buffer, uint32_t in8_stride, uint8_t* inn_bit_buffer, uint32_t inn_stride,

                        uint16_t* out16_bit_buffer, uint32_t out_stride, uint32_t width, uint32_t height);

void svt_aom_compressed_pack_sb(uint8_t* in8_bit_buffer, uint32_t in8_stride, uint8_t* inn_bit_buffer,

                                uint32_t inn_stride, uint16_t* out16_bit_buffer, uint32_t out_stride, uint32_t width,

                                uint32_t height);

// 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, int32_t w, int32_t h, int32_t stride, int32_t x_o, int32_t y_o, int32_t block_size) { \

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

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

        double        block_mean = 0;                                                                               \

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

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

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

            }                                                                                                       \

        }                                                                                                           \

        return block_mean / (max_w * max_h);                                                                        \

    }

GET_BLOCK_MEAN(uint8_t, lowbd);

GET_BLOCK_MEAN(uint16_t, highbd);

static INLINE double get_block_mean(const uint8_t* data, int32_t w, int32_t h, int32_t stride, int32_t x_o, int32_t y_o,

                                    int32_t block_size, int32_t 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,                                                  \

                                         int32_t         stride,                                                    \

                                         int32_t         w,                                                         \

                                         int32_t         h,                                                         \

                                         int32_t         x_o,                                                       \

                                         int32_t         y_o,                                                       \

                                         int32_t         block_size_x,                                              \

                                         int32_t         block_size_y) {                                                    \

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

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

        double        noise_var  = 0;                                                                               \

        double        noise_mean = 0;                                                                               \

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

            for (int32_t 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;                                               \

    }

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, int32_t w, int32_t h, int32_t stride,

                                   int32_t x_o, int32_t y_o, int32_t block_size_x, int32_t block_size_y,

                                   int32_t 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_free(AomEquationSystem* eqns) {

    if (!eqns) {

        return;

    }

    free(eqns->A);

    eqns->A = NULL;

    free(eqns->b);

    eqns->b = NULL;

    free(eqns->x);

    eqns->x = NULL;

    eqns->n = 0;

}

static void equation_system_clear(AomEquationSystem* eqns) {

    const int32_t 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(AomEquationSystem* dst, const AomEquationSystem* src) {

    const int32_t n = dst->n;

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

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

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

}

static int32_t equation_system_init(AomEquationSystem* eqns, int32_t n) {

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

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

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

    eqns->n = n;

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

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

        equation_system_free(eqns);

        return 0;

    }

    equation_system_clear(eqns);

    return 1;

}

static int32_t equation_system_solve(AomEquationSystem* eqns) {

    const int32_t n   = eqns->n;

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

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

    int32_t       ret = 0;

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

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

        free(b);

        free(A);

        return 0;

    }

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

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

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

    free(b);

    free(A);

    if (ret == 0) {

        return 0;

    }

    return 1;

}

static void noise_strength_solver_clear(AomNoiseStrengthSolver* solver) {

    equation_system_clear(&solver->eqns);

    solver->num_equations = 0;

    solver->total         = 0;

}

static void noise_strength_solver_copy(AomNoiseStrengthSolver* dest, AomNoiseStrengthSolver* src) {

    equation_system_copy(&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 int32_t num_coeffs(const AomNoiseModelParams params) {

    const int32_t 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 int32_t noise_state_init(AomNoiseState* state, int32_t n, int32_t bit_depth) {

    const int32_t k_num_bins = 20;

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

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

        return 0;

    }

    state->ar_gain          = 1.0;

    state->num_observations = 0;

    return svt_aom_noise_strength_solver_init(&state->strength_solver, k_num_bins, bit_depth);

}

static void set_chroma_coefficient_fallback_soln(AomEquationSystem* eqns) {

    const double  k_tolerance = 1e-6;

    const int32_t 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]) > k_tolerance) {

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

    }

}

int32_t svt_aom_noise_strength_lut_init(AomNoiseStrengthLut* lut, int32_t num_points) {

    if (!lut) {

        return 0;

    }

    lut->points = (double (*)[2])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 svt_aom_noise_strength_lut_free(AomNoiseStrengthLut* lut) {

    if (!lut) {

        return;

    }

    free(lut->points);

    lut->points     = NULL;

    lut->num_points = 0;

}

static double noise_strength_solver_get_bin_index(const AomNoiseStrengthSolver* 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 AomNoiseStrengthSolver* solver, double x) {

    const double  bin    = noise_strength_solver_get_bin_index(solver, x);

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

    const int32_t 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 svt_aom_noise_strength_solver_add_measurement(AomNoiseStrengthSolver* solver, double block_mean,

                                                   double noise_std) {

    const double  bin    = noise_strength_solver_get_bin_index(solver, block_mean);

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

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

    const double  a      = bin - bin_i0;

    const int32_t 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++;

}

int32_t svt_aom_noise_strength_solver_solve(AomNoiseStrengthSolver* solver) {

    // Add regularization proportional to the number of constraints

    const int32_t n       = solver->num_bins;

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

    int32_t       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*)malloc(sizeof(*A) * n * n);

    if (!A) {

        SVT_ERROR("Unable to allocate copy of A\n");

        return 0;

    }

    SVT_MEMCPY(A, old_a, sizeof(*A) * n * n);

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

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

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

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

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

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

    }

    // Small regularization to give average noise strength

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

    for (int32_t 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;

    free(A);

    return result;

}

int32_t svt_aom_noise_strength_solver_init(AomNoiseStrengthSolver* solver, int32_t num_bins, int32_t 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);

}

double svt_aom_noise_strength_solver_get_center(const AomNoiseStrengthSolver* solver, int32_t i) {

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

    const int32_t 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 AomNoiseStrengthSolver* solver, const AomNoiseStrengthLut* lut,

                                             double* residual, int32_t start, int32_t end) {

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

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

        const int32_t lower = AOMMAX(

            0, (int32_t)floor(noise_strength_solver_get_bin_index(solver, lut->points[i - 1][0])));

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

                                     (int32_t)ceil(noise_strength_solver_get_bin_index(solver, lut->points[i + 1][0])));

        double        r     = 0;

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

            const double x = svt_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;

    }

}

int32_t svt_aom_noise_strength_solver_fit_piecewise(const AomNoiseStrengthSolver* solver, int32_t max_output_points,

                                                    AomNoiseStrengthLut* lut) {

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

    // different bit-depths.

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

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

        SVT_ERROR("Failed to init lut\n");

        return 0;

    }

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

        lut->points[i][0] = svt_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 = malloc(solver->num_bins * sizeof(*residual));

    ASSERT(residual != NULL);

    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) {

        int32_t min_index = 1;

        for (int32_t 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 > k_tolerance) {

            break;

        }

        const int32_t 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);

    }

    free(residual);

    return 1;

}

int32_t svt_aom_flat_block_finder_init(AomFlatBlockFinder* block_finder, int32_t block_size, int32_t bit_depth,

                                       int32_t use_highbd) {

    const int32_t     n = block_size * block_size;

    AomEquationSystem eqns;

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

        SVT_ERROR("Failed to init equation system for block_size=%d\n", block_size);

        return 0;

    }

    double* at_a_inv = (double*)malloc(kLowPolyNumParams * kLowPolyNumParams * sizeof(*at_a_inv));

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

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

        SVT_ERROR("Failed to alloc A or at_a_inv for block_size=%d\n", block_size);

        free(at_a_inv);

        free(A);

        equation_system_free(&eqns);

        return 0;

    }

    block_finder->A             = A;

    block_finder->at_a_inv      = at_a_inv;

    block_finder->block_size    = block_size;

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

    block_finder->use_highbd    = use_highbd;

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

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

        for (int32_t 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 int32_t row              = y * block_size + x;

            A[kLowPolyNumParams * row + 0] = yd;

            A[kLowPolyNumParams * row + 1] = xd;

            A[kLowPolyNumParams * row + 2] = 1;

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

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

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

                }

            }

        }

    }

    // Lazy inverse using existing equation solver.

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

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

        eqns.b[i] = 1;

        equation_system_solve(&eqns);

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

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

        }

    }

    equation_system_free(&eqns);

    return 1;

}

void svt_aom_flat_block_finder_free(AomFlatBlockFinder* block_finder) {

    if (!block_finder) {

        return;

    }

    free(block_finder->A);

    free(block_finder->at_a_inv);

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

}

void svt_aom_flat_block_finder_extract_block_c(const AomFlatBlockFinder* block_finder, const uint8_t* const data,

                                               int32_t w, int32_t h, int32_t stride, int32_t offsx, int32_t offsy,

                                               double* plane, double* block) {

    const int32_t block_size = block_finder->block_size;

    const int32_t n          = block_size * block_size;

    const double* A          = block_finder->A;

    const double* at_a_inv   = block_finder->at_a_inv;

    const double  recp_norm  = 1 / block_finder->normalization;

    double        plane_coords[kLowPolyNumParams];

    double        at_a_inv__b[kLowPolyNumParams];

    int32_t       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 int32_t y = clamp(offsy + yi, 0, h - 1);

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

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

                block[yi * block_size + xi] = ((double)data16[y * stride + x]) * recp_norm;

            }

        }

    } else {

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

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

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

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

                block[yi * block_size + xi] = ((double)data[y * stride + x]) * recp_norm;

            }

        }

    }

#if (kLowPolyNumParams == 3)

    multiply_mat_1_n_3(block, A, at_a_inv__b, n);

    multiply_mat_3_3_1(at_a_inv, at_a_inv__b, plane_coords);

    multiply_mat_n_3_1(A, plane_coords, plane, n);

#else

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

    multiply_mat(at_a_inv, at_a_inv__b, plane_coords, kLowPolyNumParams, kLowPolyNumParams, 1);

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

#endif

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

        block[i] -= plane[i];

    }

}

typedef struct {

    int32_t index;

    float   score;

} IndexAndscore;

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

    const float diff = ((IndexAndscore*)a)->score - ((IndexAndscore*)b)->score;

    return diff < 0 ? -1 : diff > 0;

}

int32_t svt_aom_flat_block_finder_run(const AomFlatBlockFinder* block_finder, const uint8_t* const data, int32_t w,

                                      int32_t h, int32_t 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 int32_t  block_size        = block_finder->block_size;

    const int32_t  n                 = block_size * block_size;

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

    const double   k_ratio_threshold = 1.25;

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

    const double   k_var_threshold   = 0.005 / (double)n;

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

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

    int32_t        num_flat          = 0;

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

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

    IndexAndscore* scores            = (IndexAndscore*)malloc(num_blocks_w * num_blocks_h * sizeof(*scores));

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

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

        free(plane);

        free(block);

        free(scores);

        return -1;

    }

#ifdef NOISE_MODEL_LOG_SCORE

    SVT_ERROR("score = [");

#endif

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

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

            // Compute gradient covariance matrix.

            double g_xx = 0, g_xy = 0, g_yy = 0;

            double var  = 0;

            double mean = 0;

            svt_aom_flat_block_finder_extract_block(

                block_finder, data, w, h, stride, bx * block_size, by * block_size, plane, block);

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

                for (int32_t 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;

                    g_xx += gx * gx;

                    g_xy += gx * gy;

                    g_yy += gy * gy;

                    mean += block[yi * block_size + xi];

                    var += block[yi * block_size + xi] * block[yi * block_size + xi];

                }

            }

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

            // Normalize gradients by BlockSize.

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

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

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

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

            {

                const double  trace   = g_xx + g_yy;

                const double  det     = g_xx * g_yy - g_xy * g_xy;

                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 int32_t is_flat = (trace < k_trace_threshold) && (ratio < k_ratio_threshold) &&

                    (norm < k_norm_threshold) && (var > k_var_threshold);

                // 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};

                const float  score                   = (float)(1.0 /

                                            (1 +

                                             exp(-(weights[0] * var + weights[1] * ratio + weights[2] * trace +

                                                   weights[3] * norm + weights[4]))));

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

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

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

#ifdef NOISE_MODEL_LOG_SCORE

                SVT_ERROR("%g %g %g %g %g %d ", score, var, ratio, trace, norm, is_flat);

#endif

                num_flat += is_flat;

            }

        }

#ifdef NOISE_MODEL_LOG_SCORE

        SVT_ERROR("\n");

#endif

    }

#ifdef NOISE_MODEL_LOG_SCORE

    SVT_ERROR("];\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 int32_t top_nth_percentile = num_blocks_w * num_blocks_h * 90 / 100;

    const float   score_threshold    = scores[top_nth_percentile].score;

    for (int32_t 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;

        }

    }

    free(block);

    free(plane);

    free(scores);

    return num_flat;

}

int32_t svt_aom_noise_model_init(AomNoiseModel* model, const AomNoiseModelParams params) {

    const int32_t n         = num_coeffs(params);

    const int32_t lag       = params.lag;

    const int32_t bit_depth = params.bit_depth;

    int32_t       i         = 0;

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

    if (params.lag < 1) {

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

        return 0;

    }

    if (params.lag > k_max_lag) {

        SVT_ERROR("Invalid noise param: lag = %d must be <= %d\n", params.lag, k_max_lag);

        return 0;

    }

    SVT_MEMCPY(&model->params, &params, sizeof(params));

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

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

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

            svt_aom_noise_model_free(model);

            return 0;

        }

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

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

            svt_aom_noise_model_free(model);

            return 0;

        }

    }

    model->n      = n;

    model->coords = (int32_t (*)[2])malloc(sizeof(*model->coords) * n);

    if (!model->coords) {

        SVT_ERROR("Failed to allocate memory for coords\n");

        svt_aom_noise_model_free(model);

        return 0;

    }

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

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

        for (int32_t 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:

                SVT_ERROR("Invalid shape\n");

                svt_aom_noise_model_free(model);

                return 0;

            }

        }

    }

    assert(i == n);

    return 1;

}

void svt_aom_noise_model_free(AomNoiseModel* model) {

    int32_t c = 0;

    if (!model) {

        return;

    }

    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(int32_t (*coords)[2],                          \

                                          int32_t               num_coords,              \

                                          const INT_TYPE* const data,                    \

                                          const INT_TYPE* const denoised,                \

                                          int32_t               stride,                  \

                                          const int32_t         sub_log2[2],             \

                                          const INT_TYPE* const alt_data,                \

                                          const INT_TYPE* const alt_denoised,            \

                                          int32_t               alt_stride,              \

                                          int32_t               x,                       \

                                          int32_t               y,                       \

                                          double*               buffer) {                              \

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

            const int32_t 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;                                      \

            int32_t num_samples = 0;                                                     \

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

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

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

                    const int32_t 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++;                                                       \

                }                                                                        \

            }                                                                            \

            assert(num_samples > 0);                                                     \

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

        }                                                                                \

        return val;                                                                      \

    }

EXTRACT_AR_ROW(uint8_t, lowbd);

EXTRACT_AR_ROW(uint16_t, highbd);

void svt_av1_add_block_observations_internal_c(uint32_t n, const double val, const double recp_sqr_norm, double* buffer,

                                               double* buffer_norm, double* b, double* A) {

    uint32_t i;

    for (i = 0; i + 8 - 1 < n; i += 8) {

        buffer_norm[i + 0] = buffer[i + 0] * recp_sqr_norm;

        buffer_norm[i + 1] = buffer[i + 1] * recp_sqr_norm;

        buffer_norm[i + 2] = buffer[i + 2] * recp_sqr_norm;

        buffer_norm[i + 3] = buffer[i + 3] * recp_sqr_norm;

        buffer_norm[i + 4] = buffer[i + 4] * recp_sqr_norm;

        buffer_norm[i + 5] = buffer[i + 5] * recp_sqr_norm;

        buffer_norm[i + 6] = buffer[i + 6] * recp_sqr_norm;

        buffer_norm[i + 7] = buffer[i + 7] * recp_sqr_norm;

        b[i + 0] += buffer_norm[i + 0] * val;

        b[i + 1] += buffer_norm[i + 1] * val;

        b[i + 2] += buffer_norm[i + 2] * val;

        b[i + 3] += buffer_norm[i + 3] * val;

        b[i + 4] += buffer_norm[i + 4] * val;

        b[i + 5] += buffer_norm[i + 5] * val;

        b[i + 6] += buffer_norm[i + 6] * val;

        b[i + 7] += buffer_norm[i + 7] * val;

    }

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

        buffer_norm[i] = buffer[i] * recp_sqr_norm;

        b[i] += buffer_norm[i] * val;

    }

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

        uint32_t     j             = 0;

        const double buffer_norm_i = buffer_norm[i];

        for (j = 0; j + 8 - 1 < n; j += 8) {

            A[i * n + j + 0] += (buffer_norm_i * buffer[j + 0]);

            A[i * n + j + 1] += (buffer_norm_i * buffer[j + 1]);

            A[i * n + j + 2] += (buffer_norm_i * buffer[j + 2]);

            A[i * n + j + 3] += (buffer_norm_i * buffer[j + 3]);

            A[i * n + j + 4] += (buffer_norm_i * buffer[j + 4]);

            A[i * n + j + 5] += (buffer_norm_i * buffer[j + 5]);

            A[i * n + j + 6] += (buffer_norm_i * buffer[j + 6]);

            A[i * n + j + 7] += (buffer_norm_i * buffer[j + 7]);

        }

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

            A[i * n + j] += (buffer_norm_i * buffer[j]);

        }

    }

}

static int32_t add_block_observations(AomNoiseModel* noise_model, int32_t c, const uint8_t* const data,

                                      const uint8_t* const denoised, int32_t w, int32_t h, int32_t stride,

                                      int32_t sub_log2[2], const uint8_t* const alt_data,

                                      const uint8_t* const alt_denoised, int32_t alt_stride,

                                      const uint8_t* const flat_blocks, int32_t block_size, int32_t num_blocks_w,

                                      int32_t num_blocks_h) {

    const int32_t lag           = noise_model->params.lag;

    const int32_t num_coords    = noise_model->n;

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

    const double  recp_sqr_norm = 1 / (normalization * normalization);

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

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

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

    double*       buffer;

    double*       buffer_norm;

    EB_MALLOC_ALIGNED(buffer, sizeof(*buffer) * (num_coords + 1));

    if (!buffer) {

        SVT_ERROR("Unable to allocate buffer of size %d\n", sizeof(*buffer) * (num_coords + 1));

        return 0;

    }

    EB_MALLOC_ALIGNED(buffer_norm, sizeof(*buffer_norm) * (n));

    if (!buffer_norm) {

        EB_FREE_ALIGNED(buffer);

        SVT_ERROR("Unable to allocate buffer of size %d\n", sizeof(*buffer_norm) * (n));

        return 0;

    }

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

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

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

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

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

                continue;

            }

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

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

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

            int32_t 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 (int32_t y = y_start; y < y_end; ++y) {

                for (int32_t 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,

                                               alt_denoised,

                                               alt_stride,

                                               x + x_o,

                                               y + y_o,

                                               buffer);

                    svt_av1_add_block_observations_internal(n, val, recp_sqr_norm, buffer, buffer_norm, b, A);

                }

            }

            //There is situation when x_start is greater than x_end,

            //use max() to cap negative result and do not increment num_observations

            noise_model->latest_state[c].num_observations += AOMMAX((y_end - y_start), 0) *

                AOMMAX((x_end - x_start), 0);

        }

    }

    EB_FREE_ALIGNED(buffer);

    EB_FREE_ALIGNED(buffer_norm);

    return 1;

}

static void add_noise_std_observations(AomNoiseModel* noise_model, int32_t c, const double* coeffs,

                                       const uint8_t* const data, const uint8_t* const denoised, int32_t w, int32_t h,

                                       int32_t stride, const int32_t sub_log2[2], const uint8_t* const alt_data,

                                       int32_t alt_stride, const uint8_t* const flat_blocks, int32_t block_size,

                                       int32_t num_blocks_w, int32_t num_blocks_h) {

    const int32_t           num_coords            = noise_model->n;

    AomNoiseStrengthSolver* noise_strength_solver = &noise_model->latest_state[c].strength_solver;

    const AomNoiseStrengthSolver* noise_strength_luma = &noise_model->latest_state[0].strength_solver;

    const double                  luma_gain           = noise_model->latest_state[0].ar_gain;

    const double                  noise_gain          = noise_model->latest_state[c].ar_gain;

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

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

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

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

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

                continue;

            }

            const int32_t num_samples_h = AOMMIN((h >> sub_log2[1]) - by * (block_size >> sub_log2[1]),

                                                 block_size >> sub_log2[1]);

            const int32_t num_samples_w = AOMMIN((w >> sub_log2[0]) - bx * (block_size >> sub_log2[0]),

                                                 (block_size >> sub_log2[0]));

            // Make sure that we have a reasonable amount of samples to consider the

            // block

            if (num_samples_w * num_samples_h > block_size) {

                const double block_mean = get_block_mean(alt_data ? alt_data : data,

                                                         w,

                                                         h,

                                                         alt_data ? alt_stride : stride,

                                                         x_o << sub_log2[0],

                                                         y_o << sub_log2[1],

                                                         block_size,

                                                         noise_model->params.use_highbd);

                const double noise_var = get_noise_var(data,

                                                       denoised,

                                                       stride,

                                                       w >> sub_log2[0],

                                                       h >> sub_log2[1],

                                                       x_o,

                                                       y_o,

                                                       block_size >> sub_log2[0],

                                                       block_size >> sub_log2[1],

                                                       noise_model->params.use_highbd);

                // We want to remove the part of the noise that came from being

                // correlated with luma. Note that the noise solver for luma must

                // have already been run.

                const double luma_strength = c > 0

                    ? luma_gain * noise_strength_solver_get_value(noise_strength_luma, block_mean)

                    : 0;

                const double corr          = c > 0 ? coeffs[num_coords] : 0;

                // Chroma noise:

                //    N(0, noise_var) = N(0, uncorr_var) + corr * N(0, luma_strength^2)

                // The uncorrelated component:

                //   uncorr_var = noise_var - (corr * luma_strength)^2

                // But don't allow fully correlated noise (hence the max), since the

                // synthesis cannot model it.

                const double uncorr_std = sqrt(AOMMAX(noise_var / 16, noise_var - pow(corr * luma_strength, 2)));

                // After we've removed correlation with luma, undo the gain that will

                // come from running the IIR filter.

                const double adjusted_strength = uncorr_std / noise_gain;

                svt_aom_noise_strength_solver_add_measurement(noise_strength_solver, block_mean, adjusted_strength);

            }

        }

    }

}

static int32_t ar_equation_system_solve(AomNoiseState* state, int32_t is_chroma) {

    const int32_t ret = equation_system_solve(&state->eqns);

    state->ar_gain    = 1.0;

    if (!ret) {

        return ret;

    }

    // Update the AR gain from the equation system as it will be used to fit

    // the noise strength as a function of intensity.  In the Yule-Walker

    // equations, the diagonal should be the variance of the correlated noise.

    // In the case of the least squares estimate, there will be some variability

    // in the diagonal. So use the mean of the diagonal as the estimate of