/*
 * Copyright(c) 2019 Netflix, Inc.
 * Copyright (c) 2016, Alliance for Open Media. All rights reserved
 *
 * This source code is subject to the terms of the BSD 2 Clause License and
 * the Alliance for Open Media Patent License 1.0. If the BSD 2 Clause License
 * was not distributed with this source code in the LICENSE file, you can
 * obtain it at 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 <array>
#include <cmath>
#include "grainSynthesis.h"
#include "gtest/gtest.h"
#include "FilmGrainExpectedResult.h"
#include "acm_random.h"
#include "noise_model.h"
#include "aom_dsp_rtcd.h"
#include "pic_buffer_desc.h"

#if CONFIG_ENABLE_FILM_GRAIN

namespace {
constexpr AomFilmGrain film_grain_test_vectors[3] = {
    /* Test 1 */
    {
        1 /* apply_grain */,
        1 /* update_parameters */,
        {{16, 0},
         {25, 136},
         {33, 144},
         {41, 160},
         {48, 168},
         {56, 136},
         {67, 128},
         {82, 144},
         {97, 152},
         {113, 144},
         {128, 176},
         {143, 168},
         {158, 176},
         {178, 184}},
        14 /* num_points_y */,
        {{16, 0},
         {20, 64},
         {28, 88},
         {60, 104},
         {90, 136},
         {105, 160},
         {134, 168},
         {168, 208}},
        8 /* num_cb_points */,
        {{16, 0},
         {28, 96},
         {56, 80},
         {66, 96},
         {80, 104},
         {108, 96},
         {122, 112},
         {137, 112},
         {169, 176}},
        9 /* num_cr_points */,
        11 /* scaling_shift */,
        2 /* ar_coeff_lag */,
        {0, 0, -58, 0, 0, 0, -76, 100, -43, 0, -51, 82},
        {0, 0, -49, 0, 0, 0, -36, 22, -30, 0, -38, 7, 39},
        {0, 0, -47, 0, 0, 0, -31, 31, -25, 0, -32, 13, -100},
        8 /* ar_coeff_shift */,
        247 /* cb_mult */,
        192 /* cb_luma_mult */,
        18 /* cb_offset */,
        229 /* cr_mult */,
        192 /* cr_luma_mult */,
        54 /* cr_offset */,
        0 /* overlap_flag */,
        1 /* clip_to_restricted_range */,
        8 /* bit_depth */,
        0 /* chroma_scaling_from_luma*/,
        0 /* grain_scale_shift*/,
        45231 /* random_seed */,
        0 /* ignore_ref */
    },
    /* Test 2 */
    {
        1 /* apply_grain */,
        0 /* update_parameters */,
        {{16, 0},
         {25, 136},
         {33, 144},
         {41, 160},
         {48, 168},
         {56, 136},
         {67, 128},
         {82, 144},
         {97, 152},
         {113, 144},
         {128, 176},
         {143, 168},
         {158, 176},
         {178, 184}},
        14 /* num_points_y */,
        {{16, 0},
         {20, 64},
         {28, 88},
         {60, 104},
         {90, 136},
         {105, 160},
         {134, 168},
         {168, 208}},
        8 /* num_cb_points */,
        {{16, 0},
         {28, 96},
         {56, 80},
         {66, 96},
         {80, 104},
         {108, 96},
         {122, 112},
         {137, 112},
         {169, 176}},
        9 /* num_cr_points */,
        11 /* scaling_shift */,
        2 /* ar_coeff_lag */,
        {0, 0, -58, 0, 0, 0, -76, 100, -43, 0, -51, 82},
        {0, 0, -49, 0, 0, 0, -36, 22, -30, 0, -38, 7, 39},
        {0, 0, -47, 0, 0, 0, -31, 31, -25, 0, -32, 13, -100},
        8 /* ar_coeff_shift */,
        247 /* cb_mult */,
        192 /* cb_luma_mult */,
        18 /* cb_offset */,
        229 /* cr_mult */,
        192 /* cr_luma_mult */,
        54 /* cr_offset */,
        0 /* overlap_flag */,
        1 /* clip_to_restricted_range */,
        8 /* bit_depth */,
        0 /* chroma_scaling_from_luma*/,
        0 /* grain_scale_shift*/,
        36007 /* random_seed */,
        0 /* ignore_ref */
    },

    /* Test 3 */
    {
        1 /* apply_grain */,
        1 /* update_parameters */,
        {{0, 96}, {255, 96}},
        2 /* num_points_y */,
        {{0, 64}, {255, 64}},
        2 /* num_cb_points */,
        {{0, 64}, {255, 64}},
        2 /* num_cr_points */,
        11 /* scaling_shift */,
        3 /* ar_coeff_lag */,
        {
            4, 1,   3, 0,   1,  -3, 8,  -3, 7,  -23, 1, -25,
            0, -10, 6, -17, -4, 53, 36, 5,  -5, -17, 8, 66,
        },
        {
            0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,   0,
            0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 127,
        },
        {
            0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,   0,
            0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 127,
        },
        7 /* ar_coeff_shift */,
        128 /* cb_mult */,
        192 /* cb_luma_mult */,
        256 /* cb_offset */,
        128 /* cr_mult */,
        192 /* cr_luma_mult */,
        256 /* cr_offset */,
        1 /* overlap_flag */,
        0 /* clip_to_restricted_range */,
        8 /* bit_depth */,
        0 /*chroma_scaling_from_luma*/,
        0 /* grain_scale_shift*/,
        45231 /* random_seed */,
        0 /* ignore_ref */
    },
};

TEST(FilmGrain, parameters_equality) {
    /* Film grain parameters equality should not depend on random_seed and
     * update_parameters values */
    EXPECT_EQ(svt_aom_film_grain_params_equal(film_grain_test_vectors,
                                              film_grain_test_vectors + 1),
              1);

    /* These two instances of film grain parameters are different */
    EXPECT_EQ(svt_aom_film_grain_params_equal(film_grain_test_vectors,
                                              film_grain_test_vectors + 2),
              0);
}

class AddFilmGrainTest : public ::testing::Test {
  protected:
    void init_data() {
        luma_.fill(0);
        cb_.fill(0);
        cr_.fill(0);
    }

    void check_output(int idx) {
        int fail_y = 0, fail_cb = 0, fail_cr = 0;
        for (int i = 0; i < luma_size; ++i)
            if (luma_[i] != add_grain_expected_luma[idx][i])
                ++fail_y;

        for (int i = 0; i < chroma_size; ++i)
            if (cb_[i] != add_grain_expected_cb[idx][i])
                ++fail_cb;

        for (int i = 0; i < chroma_size; ++i)
            if (cr_[i] != add_grain_expected_cr[idx][i])
                ++fail_cr;

        EXPECT_EQ(fail_y, 0);
        EXPECT_EQ(fail_cb, 0);
        EXPECT_EQ(fail_cr, 0);
    }

  protected:
    static constexpr int kWidth = 128;
    static constexpr int kHeight = 128;
    static constexpr int luma_size = kWidth * kHeight;
    static constexpr int chroma_size = luma_size >> 2;

    alignas(16) std::array<uint8_t, luma_size> luma_{};
    alignas(16) std::array<uint8_t, chroma_size> cb_{};
    alignas(16) std::array<uint8_t, chroma_size> cr_{};
};

TEST_F(AddFilmGrainTest, MatchTest) {
    for (int i = 0; i < 3; ++i) {
        init_data();
        svt_av1_add_film_grain_run(film_grain_test_vectors + i,
                                   luma_.data(),
                                   cb_.data(),
                                   cr_.data(),
                                   kHeight,
                                   kWidth,
                                   kWidth,     /* luma stride */
                                   kWidth / 2, /* chroma stride */
                                   0,
                                   1,
                                   1);
        check_output(i);
        EXPECT_FALSE(HasFailure());
    }
}

// Return normally distrbuted values with standard deviation of sigma.
double randn(libaom_test::ACMRandom &random, double sigma) {
    while (1) {
        const double u = 2.0 * ((double)random.Rand31() /
                                testing::internal::Random::kMaxRange) -
                         1.0;
        const double v = 2.0 * ((double)random.Rand31() /
                                testing::internal::Random::kMaxRange) -
                         1.0;
        const double s = u * u + v * v;
        if (s > 0 && s < 1) {
            return sigma * (u * std::sqrt(-2.0 * std::log(s) / s));
        }
    }
    return 0;
}

/* clang-format off */
constexpr AomFilmGrain expected_film_grain = {
    1 /* apply_grain */,
    1 /* update_parameters */,
    {{0, 33}, {255, 34}, {0, 0}, {0, 0}, },
    2 /* num_y_points */,
    {{0, 21}, {94, 22}, {228, 23}, {255, 23}, {0, 0}, },
    4 /* num_cb_points */,
    {{0, 21}, {94, 22}, {228, 23}, {255, 23}, {0, 0}, },
    4 /* num_cr_points */,
    11,
    3,
    {-3, -5, -6, -5, -4, -2, -2, -3, -7, -4, -7, -6, -4, -1, -5, -6, -7, -15, -5, -1, -3, -2, -6, -15},
    {-3, -1, -10, -7, -6, -3, -3, -4, -9, -17, -13, -12, -4, 1, -7, -12, -26, -29, -17, -5, 2, 1, -9, -28, 111},
    {-3, -1, -10, -7, -6, -3, -3, -4, -9, -17, -13, -12, -4, 1, -7, -12, -26, -29, -17, -5, 2, 1, -9, -28, 111},
    7 /* ar_coeff_shift */,
    128 /* cb_mult */,
    192 /* cb_luma_mult */,
    256 /* cb_offset */,
    128 /* cr_mult */,
    192 /* cr_luma_mult */,
    256 /* cr_offset */,
    1 /* overlap_flag */,
    0 /* clip_to_restricted_range */,
    8 /* bit_depth */,
    0 /* chroma_scaling_from_luma */,
    0 /* grain_scale_shift */,
    0 /* random_seed */,
    0 /* ignore_ref */
};
/* clang-format on */

class DenoiseModelRunTest : public ::testing::Test {
  public:
    static constexpr int width_ = 128;
    static constexpr int height_ = 128;

    ~DenoiseModelRunTest() override {
        if (in_pic_.dctor) {
            in_pic_.dctor(&in_pic_);
        }
        if (noise_model.dctor) {
            noise_model.dctor(&noise_model);
        }
    }

    void SetUp() override {
#if defined(ARCH_X86_64) || defined(ARCH_AARCH64)
        EbCpuFlags cpu_flags = svt_aom_get_cpu_flags_to_use();
#else
        EbCpuFlags cpu_flags = 0;
#endif
        svt_aom_setup_rtcd_internal(cpu_flags);

        // create EbPictureBufferDesc
        EbPictureBufferDescInitData pbd_init_data;
        pbd_init_data.max_width = width_;
        pbd_init_data.max_height = height_;
        pbd_init_data.bit_depth = EB_EIGHT_BIT;
        // allocate all the components
        pbd_init_data.buffer_enable_mask = PICTURE_BUFFER_DESC_FULL_MASK;
        pbd_init_data.border = 0;
        pbd_init_data.color_format = EB_YUV420;
        pbd_init_data.split_mode = false;
        pbd_init_data.is_16bit_pipeline = false;

        const int subsampling_x_ =
            (pbd_init_data.color_format == EB_YUV444 ? 0 : 1);

        EbErrorType err =
            svt_picture_buffer_desc_ctor(&in_pic_, &pbd_init_data);
        EXPECT_EQ(err, EB_ErrorNone) << "create input pic fail";
        data_ptr_[0] = in_pic_.y_buffer;
        data_ptr_[1] = in_pic_.u_buffer;
        data_ptr_[2] = in_pic_.v_buffer;

        // create the denoise and noise model
        DenoiseAndModelInitData fg_init_data;
        fg_init_data.adaptive_film_grain = true;
        fg_init_data.encoder_bit_depth = EB_EIGHT_BIT;
        fg_init_data.encoder_color_format = EB_YUV420;
        fg_init_data.noise_level = 4;  // TODO: check the range;
        fg_init_data.denoise_apply = false;
        fg_init_data.width = width_;
        fg_init_data.height = height_;
        fg_init_data.y_stride = width_;
        fg_init_data.u_stride = fg_init_data.v_stride =
            fg_init_data.y_stride >> subsampling_x_;

        err = svt_aom_denoise_and_model_ctor(&noise_model, &fg_init_data);
        EXPECT_EQ(err, EB_ErrorNone) << "svt_aom_denoise_and_model_ctor fail";
    }

    void init_data() {
        for (int y = 0; y < height_; ++y) {
            for (int x = 0; x < width_; ++x) {
                data_ptr_[0][y * width_ + x] =
                    int(64 + y + randn(this->random_, 1));
                // Make the chroma planes completely correlated with the Y plane
                for (int c = 1; c < 3; ++c) {
                    data_ptr_[c][(y >> 1) * (width_ >> 1) + (x >> 1)] =
                        data_ptr_[0][y * width_ + x];
                }
            }
        }
    }

    void check_filmgrain() {
        EXPECT_EQ(svt_aom_film_grain_params_equal(&output_film_grain,
                                                  &expected_film_grain),
                  1);
    }

    void run_test() {
        init_data();

        svt_aom_denoise_and_model_run(
            &noise_model, &in_pic_, &output_film_grain, 0);
    }

  protected:
    EbPictureBufferDesc in_pic_{};
    AomDenoiseAndModel noise_model{};
    AomFilmGrain output_film_grain{};
    libaom_test::ACMRandom random_{100171};
    uint8_t *data_ptr_[3]{};
    uint8_t *denoised_ptr_[3]{};
};

TEST_F(DenoiseModelRunTest, OutputFilmGrainCheck) {
    run_test();
    check_filmgrain();
    EXPECT_FALSE(HasFailure());
}

}  // namespace

#endif
