/*
 * Copyright (c) 2021, Alliance for Open Media. All rights reserved
 *
 * This source code is subject to the terms of the BSD 3-Clause Clear License
 * and the Alliance for Open Media Patent License 1.0. If the BSD 3-Clause Clear
 * License was not distributed with this source code in the LICENSE file, you
 * can obtain it at aomedia.org/license/software-license/bsd-3-c-c/.  If the
 * Alliance for Open Media Patent License 1.0 was not distributed with this
 * source code in the PATENTS file, you can obtain it at
 * aomedia.org/license/patent-license/.
 */

#ifndef AOM_AOM_DSP_NOISE_MODEL_H_
#define AOM_AOM_DSP_NOISE_MODEL_H_

#ifdef __cplusplus
extern "C" {
#endif  // __cplusplus

#include <stdint.h>
#include "aom_dsp/grain_synthesis.h"
#include "aom_scale/yv12config.h"

/*!\brief Wrapper of data required to represent linear system of eqns and soln.
 */
typedef struct {
  double *A;
  double *b;
  double *x;
  int n;
} aom_equation_system_t;

/*!\brief Representation of a piecewise linear curve
 *
 * Holds n points as (x, y) pairs, that store the curve.
 */
typedef struct {
  double (*points)[2];
  int num_points;
} aom_noise_strength_lut_t;

/*!\brief Init the noise strength lut with the given number of points*/
int aom_noise_strength_lut_init(aom_noise_strength_lut_t *lut, int num_points);

/*!\brief Frees the noise strength lut. */
void aom_noise_strength_lut_free(aom_noise_strength_lut_t *lut);

/*!\brief Evaluate the lut at the point x.
 *
 * \param[in] lut  The lut data.
 * \param[in] x    The coordinate to evaluate the lut.
 */
double aom_noise_strength_lut_eval(const aom_noise_strength_lut_t *lut,
                                   double x);

/*!\brief Helper struct to model noise strength as a function of intensity.
 *
 * Internally, this structure holds a representation of a linear system
 * of equations that models noise strength (standard deviation) as a
 * function of intensity. The mapping is initially stored using a
 * piecewise representation with evenly spaced bins that cover the entire
 * domain from [min_intensity, max_intensity]. Each observation (x,y) gives a
 * constraint of the form:
 *   y_{i} (1 - a) + y_{i+1} a = y
 * where y_{i} is the value of bin i and x_{i} <= x <= x_{i+1} and
 * a = x/(x_{i+1} - x{i}). The equation system holds the corresponding
 * normal equations.
 *
 * As there may be missing data, the solution is regularized to get a
 * complete set of values for the bins. A reduced representation after
 * solving can be obtained by getting the corresponding noise_strength_lut_t.
 */
typedef struct {
  aom_equation_system_t eqns;
  double min_intensity;
  double max_intensity;
  int num_bins;
  int num_equations;
  double total;
} aom_noise_strength_solver_t;

/*!\brief Initializes the noise solver with the given number of bins.
 *
 * Returns 0 if initialization fails.
 *
 * \param[in]  solver    The noise solver to be initialized.
 * \param[in]  num_bins  Number of bins to use in the internal representation.
 * \param[in]  bit_depth The bit depth used to derive {min,max}_intensity.
 */
int aom_noise_strength_solver_init(aom_noise_strength_solver_t *solver,
                                   int num_bins, int bit_depth);
void aom_noise_strength_solver_free(aom_noise_strength_solver_t *solver);

/*!\brief Gets the x coordinate of bin i.
 *
 * \param[in]  i  The bin whose coordinate to query.
 */
double aom_noise_strength_solver_get_center(
    const aom_noise_strength_solver_t *solver, int i);

/*!\brief Add an observation of the block mean intensity to its noise strength.
 *
 * \param[in]  block_mean  The average block intensity,
 * \param[in]  noise_std   The observed noise strength.
 */
void aom_noise_strength_solver_add_measurement(
    aom_noise_strength_solver_t *solver, double block_mean, double noise_std);

/*!\brief Solves the current set of equations for the noise strength. */
int aom_noise_strength_solver_solve(aom_noise_strength_solver_t *solver);

/*!\brief Fits a reduced piecewise linear lut to the internal solution
 *
 * \param[in] max_num_points  The maximum number of output points
 * \param[out] lut  The output piecewise linear lut.
 */
int aom_noise_strength_solver_fit_piecewise(
    const aom_noise_strength_solver_t *solver, int max_num_points,
    aom_noise_strength_lut_t *lut);

/*!\brief Helper for holding precomputed data for finding flat blocks.
 *
 * Internally a block is modeled with a low-order polynomial model. A
 * planar model would be a bunch of equations like:
 * <[y_i x_i 1], [a_1, a_2, a_3]>  = b_i
 * for each point in the block. The system matrix A with row i as [y_i x_i 1]
 * is maintained as is the inverse, inv(A'*A), so that the plane parameters
 * can be fit for each block.
 */
typedef struct {
  double *AtA_inv;
  double *A;
  int num_params;  // The number of parameters used for internal low-order model
  int block_size;  // The block size the finder was initialized with
  double normalization;  // Normalization factor (1 / (2^(bit_depth) - 1))
} aom_flat_block_finder_t;

/*!\brief Init the block_finder with the given block size, bit_depth */
int aom_flat_block_finder_init(aom_flat_block_finder_t *block_finder,
                               int block_size, int bit_depth);
void aom_flat_block_finder_free(aom_flat_block_finder_t *block_finder);

/*!\brief Helper to extract a block and low order "planar" model. */
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);

/*!\brief Runs the flat block finder on the input data.
 *
 * Find flat blocks in the input image data. Returns a map of
 * flat_blocks, where the value of flat_blocks map will be non-zero
 * when a block is determined to be flat. A higher value indicates a bigger
 * confidence in the decision.
 */
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 noise shape indicates the allowed coefficients in the AR model.
enum {
  AOM_NOISE_SHAPE_DIAMOND = 0,
  AOM_NOISE_SHAPE_SQUARE = 1
} UENUM1BYTE(aom_noise_shape);

// The parameters of the noise model include the shape type, lag, the
// bit depth of the input images provided, and whether the input images
// will be using uint16 (or uint8) representation.
typedef struct {
  aom_noise_shape shape;
  int lag;
  int bit_depth;
} aom_noise_model_params_t;

/*!\brief State of a noise model estimate for a single channel.
 *
 * This contains a system of equations that can be used to solve
 * for the auto-regressive coefficients as well as a noise strength
 * solver that can be used to model noise strength as a function of
 * intensity.
 */
typedef struct {
  aom_equation_system_t eqns;
  aom_noise_strength_solver_t strength_solver;
  int num_observations;  // The number of observations in the eqn system
  double ar_gain;        // The gain of the current AR filter
} aom_noise_state_t;

/*!\brief Complete model of noise for a planar video
 *
 * This includes a noise model for the latest frame and an aggregated
 * estimate over all previous frames that had similar parameters.
 */
typedef struct {
  aom_noise_model_params_t params;
  aom_noise_state_t combined_state[3];  // Combined state per channel
  aom_noise_state_t latest_state[3];    // Latest state per channel
  int (*coords)[2];  // Offsets (x,y) of the coefficient samples
  int n;             // Number of parameters (size of coords)
  int bit_depth;
} aom_noise_model_t;

/*!\brief Result of a noise model update. */
enum {
  AOM_NOISE_STATUS_OK = 0,
  AOM_NOISE_STATUS_INVALID_ARGUMENT,
  AOM_NOISE_STATUS_INSUFFICIENT_FLAT_BLOCKS,
  AOM_NOISE_STATUS_DIFFERENT_NOISE_TYPE,
  AOM_NOISE_STATUS_INTERNAL_ERROR,
} UENUM1BYTE(aom_noise_status_t);

/*!\brief Initializes a noise model with the given parameters.
 *
 * Returns 0 on failure.
 */
int aom_noise_model_init(aom_noise_model_t *model,
                         const aom_noise_model_params_t params);
void aom_noise_model_free(aom_noise_model_t *model);

/*!\brief Updates the noise model with a new frame observation.
 *
 * Updates the noise model with measurements from the given input frame and a
 * denoised variant of it. Noise is sampled from flat blocks using the flat
 * block map.
 *
 * Returns a noise_status indicating if the update was successful. If the
 * Update was successful, the combined_state is updated with measurements from
 * the provided frame. If status is OK or DIFFERENT_NOISE_TYPE, the latest noise
 * state will be updated with measurements from the provided frame.
 *
 * \param[in,out] noise_model     The noise model to be updated
 * \param[in]     data            Raw frame data
 * \param[in]     denoised        Denoised frame data.
 * \param[in]     w               Frame width
 * \param[in]     h               Frame height
 * \param[in]     strides         Stride of the planes
 * \param[in]     chroma_sub_log2 Chroma subsampling for planes != 0.
 * \param[in]     flat_blocks     A map to blocks that have been determined flat
 * \param[in]     block_size      The size of blocks.
 */
aom_noise_status_t aom_noise_model_update(
    aom_noise_model_t *const noise_model, const uint8_t *const data[3],
    const uint8_t *const denoised[3], int w, int h, int strides[3],
    int chroma_sub_log2[2], const uint8_t *const flat_blocks, int block_size);

/*\brief Save the "latest" estimate into the "combined" estimate.
 *
 * This is meant to be called when the noise modeling detected a change
 * in parameters (or for example, if a user wanted to reset estimation at
 * a shot boundary).
 */
void aom_noise_model_save_latest(aom_noise_model_t *noise_model);

/*!\brief Converts the noise_model parameters to the corresponding
 *    grain_parameters.
 *
 * The noise structs in this file are suitable for estimation (e.g., using
 * floats), but the grain parameters in the bitstream are quantized. This
 * function does the conversion by selecting the correct quantization levels.
 */
int aom_noise_model_get_grain_parameters(aom_noise_model_t *const noise_model,
                                         aom_film_grain_t *film_grain);

/*!\brief Perform a Wiener filter denoising in 2D using the provided noise psd.
 *
 * \param[in]     data            Raw frame data
 * \param[out]    denoised        Denoised frame data
 * \param[in]     w               Frame width
 * \param[in]     h               Frame height
 * \param[in]     stride          Stride of the planes
 * \param[in]     chroma_sub_log2 Chroma subsampling for planes != 0.
 * \param[in]     noise_psd       The power spectral density of the noise
 * \param[in]     block_size      The size of blocks
 * \param[in]     bit_depth       Bit depth of the image
 */
int aom_wiener_denoise_2d(const uint8_t *const data[3], uint8_t *denoised[3],
                          int w, int h, int stride[3], int chroma_sub_log2[2],
                          float *noise_psd[3], int block_size, int bit_depth);

struct aom_denoise_and_model_t;

/*!\brief Denoise the buffer and model the residual noise.
 *
 * This is meant to be called sequentially on input frames. The input buffer
 * is denoised and the residual noise is modelled. The current noise estimate
 * is populated in film_grain. Returns true on success. The grain.apply_grain
 * parameter will be true when the input buffer was successfully denoised and
 * grain was modelled. Returns false on error.
 *
 * \param[in]      ctx   Struct allocated with aom_denoise_and_model_alloc
 *                       that holds some buffers for denoising and the current
 *                       noise estimate.
 * \param[in/out]   buf  The raw input buffer to be denoised.
 * \param[out]    grain  Output film grain parameters
 */
int aom_denoise_and_model_run(struct aom_denoise_and_model_t *ctx,
                              YV12_BUFFER_CONFIG *buf, aom_film_grain_t *grain);

/*!\brief Allocates a context that can be used for denoising and noise modeling.
 *
 * \param[in]  bit_depth   Bit depth of buffers this will be run on.
 * \param[in]  block_size  Block size for noise modeling and flat block
 *                         estimation
 * \param[in]  noise_level The noise_level (2.5 for moderate noise, and 5 for
 *                         higher levels of noise)
 */
struct aom_denoise_and_model_t *aom_denoise_and_model_alloc(int bit_depth,
                                                            int block_size,
                                                            float noise_level);

/*!\brief Frees the denoise context allocated with aom_denoise_and_model_alloc
 */
void aom_denoise_and_model_free(struct aom_denoise_and_model_t *denoise_model);

#ifdef __cplusplus
}  // extern "C"
#endif  // __cplusplus
#endif  // AOM_AOM_DSP_NOISE_MODEL_H_