// Copyright 2021 The libgav1 Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#include "src/dsp/loop_restoration.h"

#include <algorithm>
#include <cstdint>
#include <cstring>
#include <string>

#include "absl/strings/match.h"
#include "absl/time/clock.h"
#include "absl/time/time.h"
#include "gtest/gtest.h"
#include "src/dsp/common.h"
#include "src/dsp/dsp.h"
#include "src/utils/common.h"
#include "src/utils/constants.h"
#include "src/utils/cpu.h"
#include "src/utils/memory.h"
#include "tests/block_utils.h"
#include "tests/third_party/libvpx/acm_random.h"
#include "tests/utils.h"

namespace libgav1 {
namespace dsp {
namespace {

// in unit of Pixel.
constexpr int kBorder = 16;
constexpr int kWidth = 256;
constexpr int kHeight = 255;
constexpr int kStride = kWidth + 2 * kBorder;
constexpr int kOffset = kBorder * kStride + kBorder;
constexpr int kMaxBlockSize = 288 * kStride;
constexpr int kUnitWidths[] = {32, 64, 128, 256};

constexpr int kNumRadiusTypes = 3;
constexpr int kNumWienerOrders = 4;
constexpr int kWienerOrders[] = {7, 5, 3, 1};
constexpr int kWienerOrderIdLookup[] = {0, 3, 0, 2, 0, 1, 0, 0};

template <int bitdepth, typename Pixel>
class SelfGuidedFilterTest : public testing::TestWithParam<int>,
                             public test_utils::MaxAlignedAllocable {
 public:
  SelfGuidedFilterTest() = default;
  SelfGuidedFilterTest(const SelfGuidedFilterTest&) = delete;
  SelfGuidedFilterTest& operator=(const SelfGuidedFilterTest&) = delete;
  ~SelfGuidedFilterTest() override = default;

  void SetUp() override {
    test_utils::ResetDspTable(bitdepth);
    LoopRestorationInit_C();
    const testing::TestInfo* const test_info =
        testing::UnitTest::GetInstance()->current_test_info();
    const char* const test_case = test_info->test_suite_name();
    if (absl::StartsWith(test_case, "C/")) {
    } else if (absl::StartsWith(test_case, "AVX2/")) {
      if ((GetCpuInfo() & kAVX2) != 0) {
        LoopRestorationInit_AVX2();
#if LIBGAV1_MAX_BITDEPTH >= 10
        LoopRestorationInit10bpp_AVX2();
#endif
      }
    } else if (absl::StartsWith(test_case, "SSE41/")) {
      if ((GetCpuInfo() & kSSE4_1) != 0) {
        LoopRestorationInit_SSE4_1();
#if LIBGAV1_MAX_BITDEPTH >= 10
        LoopRestorationInit10bpp_SSE4_1();
#endif
      }
    } else if (absl::StartsWith(test_case, "NEON/")) {
      LoopRestorationInit_NEON();
    } else {
      FAIL() << "Unrecognized architecture prefix in test case name: "
             << test_case;
    }
    const Dsp* const dsp = GetDspTable(bitdepth);
    ASSERT_NE(dsp, nullptr);
    target_self_guided_filter_func_ = dsp->loop_restorations[1];
    restoration_info_.type = kLoopRestorationTypeSgrProj;
    memset(dst_, 0, sizeof(dst_));
  }

  void SetInputData(int type, Pixel value, int radius_index,
                    libvpx_test::ACMRandom* rnd);
  void TestFixedValues(int test_index, Pixel value);
  void TestRandomValues(bool speed);

 protected:
  const int unit_width_ = GetParam();
  const int unit_height_ = kRestorationUnitHeight;

 private:
  alignas(kMaxAlignment) Pixel src_[kMaxBlockSize];
  alignas(kMaxAlignment) Pixel dst_[kMaxBlockSize];
  RestorationUnitInfo restoration_info_;
  RestorationBuffer restoration_buffer_;
  LoopRestorationFunc target_self_guided_filter_func_;
};

template <int bitdepth, typename Pixel>
void SelfGuidedFilterTest<bitdepth, Pixel>::SetInputData(
    int type, Pixel value, int radius_index,
    libvpx_test::ACMRandom* const rnd) {
  const int mask = (1 << bitdepth) - 1;
  if (type == 0) {  // Set fixed values
    for (auto& s : src_) s = value;
  } else {  // Set random values
    for (auto& s : src_) s = rnd->Rand16() & mask;
  }
  for (auto& d : dst_) d = rnd->Rand16() & mask;
  restoration_info_.sgr_proj_info.multiplier[0] =
      kSgrProjMultiplierMin[0] +
      rnd->PseudoUniform(kSgrProjMultiplierMax[0] - kSgrProjMultiplierMin[0] +
                         1);
  restoration_info_.sgr_proj_info.multiplier[1] =
      kSgrProjMultiplierMin[1] +
      rnd->PseudoUniform(kSgrProjMultiplierMax[1] - kSgrProjMultiplierMin[1] +
                         1);
  // regulate multiplier so that it matches libaom.
  // Valid self-guided filter doesn't allow r0 and r1 to be 0 at the same time.
  // When r0 or r1 is zero, its corresponding multiplier is set to zero in
  // libaom.
  int index;
  if (radius_index == 0) {
    index = 0;  // r0 = 2, r1 = 1
  } else if (radius_index == 1) {
    index = 10;  // r0 = 0, r1 = 1
  } else /* if (radius_index == 2) */ {
    index = 14;  // r0 = 2, r1 = 0
  }
  const uint8_t r0 = kSgrProjParams[index][0];
  const uint8_t r1 = kSgrProjParams[index][2];
  static constexpr int kMultiplier[2] = {0, 95};
  restoration_info_.sgr_proj_info.index = index;
  if (r0 == 0) {
    restoration_info_.sgr_proj_info.multiplier[0] = kMultiplier[0];
  } else if (r1 == 0) {
    restoration_info_.sgr_proj_info.multiplier[1] = kMultiplier[1];
  }
}

template <int bitdepth, typename Pixel>
void SelfGuidedFilterTest<bitdepth, Pixel>::TestFixedValues(int test_index,
                                                            Pixel value) {
  static const char* const kDigest[][2][kNumRadiusTypes] = {
      {{"7b78783ff4f03625a50c2ebfd574adca", "4faa0810639016f11a9f761ce28c38b0",
        "a03314fc210bee68c7adbb44d2bbdac7"},
       {"fce031d1339cfef5016e76a643538a71", "d439e1060de3f07b5b29c9b0b7c08e54",
        "a6583fe9359877f4a259c81d900fc4fb"}},
      {{"948ea16a90c4cefef87ce5b0ee105fc6", "76740629877b721432b84dbbdb4e352a",
        "27100f37b3e42a5f2a051e1566edb6f8"},
       {"dd320de3bc82f4ba69738b2190ea9f85", "bf82f271e30a1aca91e53b086e133fb3",
        "69c274ac59c99999e1bfbf2fc4586ebd"}},
      {{"9fbf1b246011250f38532a543cc6dd74", "d5c1e0142390ebb51b075c49f8ee9ff4",
        "92f31086ba2f9e1508983b22d93a4e5c"},
       {"2198321e6b95e7199738e60f5ddc6966", "34f74626027ffca010c824ddf0942b13",
        "43dd7df2c2a601262c68cd8af1c61b82"}},
      {{"42364ff8dbdbd6706fa3b8855a4258be", "a7843fdfd4d3c0d80ba812b353b4d6b4",
        "f8a6a025827f29f857bed3e28ba3ea33"},
       {"b83c1f8d7712e37f9b21b033822e37ed", "589daf2e3e6f8715873920515cfc1b42",
        "20dcbe8e317a4373bebf11d56adc5f02"}}};
  if (target_self_guided_filter_func_ == nullptr) return;
  ASSERT_LT(value, 1 << bitdepth);
  constexpr int bd_index = (bitdepth == 8) ? 0 : 1;
  libvpx_test::ACMRandom rnd(libvpx_test::ACMRandom::DeterministicSeed());
  const Pixel* const src = src_ + kOffset;
  Pixel* const dst = dst_ + kOffset;
  for (int radius_index = 0; radius_index < kNumRadiusTypes; ++radius_index) {
    SetInputData(0, value, radius_index, &rnd);
    const absl::Time start = absl::Now();
    for (int y = 0; y < kHeight; y += unit_height_) {
      const int height = std::min(unit_height_, kHeight - y);
      for (int x = 0; x < kWidth; x += unit_width_) {
        const int width = std::min(unit_width_, kWidth - x);
        const Pixel* const source = src + y * kStride + x;
        target_self_guided_filter_func_(
            restoration_info_, source, kStride,
            source - kRestorationVerticalBorder * kStride, kStride,
            source + height * kStride, kStride, width, height,
            &restoration_buffer_, dst + y * kStride + x);
      }
    }
    const absl::Duration elapsed_time = absl::Now() - start;
    test_utils::CheckMd5Digest(
        "kLoopRestorationTypeSgrProj", std::to_string(GetParam()).c_str(),
        kDigest[test_index][bd_index][radius_index], dst_ + kBorder * kStride,
        kHeight * kStride * sizeof(*dst_), elapsed_time);
  }
}

template <int bitdepth, typename Pixel>
void SelfGuidedFilterTest<bitdepth, Pixel>::TestRandomValues(bool speed) {
  static const char* const kDigest[][2][kNumRadiusTypes] = {
      {{"9f8358ed820943fa0abe3a8ebb5887db", "fb5d48870165522341843bcbfa8674fb",
        "ca67159cd29475ac5d52ca4a0df3ea10"},
       {"a78641886ea0cf8757057d1d91e01434", "1b95172a5f2f9c514c78afa4cf8e5678",
        "a8ba988283d9e1ad1f0dcdbf6bbdaade"}},
      {{"f219b445e5c80ffb5dd0359cc2cb4dd4", "699b2c9ddca1cbb0d4fc24cbcbe951e9",
        "a4005899fa8d3c3c4669910f93ff1290"},
       {"10a75cab3c78b891c8c6d92d55f685d1", "d46f158f57c628136f6f298ee8ca6e0e",
        "07203ad761775d5d317f2b7884afd9fe"}},
      {{"000d4e382be4003b514c9135893d0a37", "8fb082dca975be363bfc9c2d317ae084",
        "475bcb6a58f87da7723f6227bc2aca0e"},
       {"4d589683f69ccc5b416149dcc5c835d5", "986b6832df1f6020d50be61ae121e42f",
        "7cb5c5dbdb3d1c54cfa00def450842dc"}},
      {{"fd43bfe34d63614554dd29fb24b12173", "5c1ba74ba3062c769d5c3c86a85ac9b9",
        "f1eda6d15b37172199d9949c2315832f"},
       {"a11be3117fb77e8fe113581b06f98bd1", "df94d12b774ad5cf744c871e707c36c8",
        "b23dc0b54c3500248d53377030428a61"}},
      {{"f3079b3b21d8dc6fce7bb1fd104be359", "c6fcbc686cfb97ab3a64f445d73aad36",
        "23966cba3e0e7803eeb951905861e0dd"},
       {"7210391a6fe26e5ca5ea205bc38aa035", "4c3e6eccad3ea152d320ecd1077169de",
        "dcee48f94126a2132963e86e93dd4903"}}};
  if (target_self_guided_filter_func_ == nullptr) return;
  constexpr int bd_index = (bitdepth == 8) ? 0 : 1;
  const int num_inputs = speed ? 1 : 5;
  const int num_tests = speed ? 20000 : 1;
  libvpx_test::ACMRandom rnd(libvpx_test::ACMRandom::DeterministicSeed());
  const Pixel* const src = src_ + kOffset;
  Pixel* const dst = dst_ + kOffset;
  for (int i = 0; i < num_inputs; ++i) {
    for (int radius_index = 0; radius_index < kNumRadiusTypes; ++radius_index) {
      SetInputData(1, 0, radius_index, &rnd);
      const absl::Time start = absl::Now();
      for (int k = 0; k < num_tests; ++k) {
        for (int y = 0; y < kHeight; y += unit_height_) {
          const int height = std::min(unit_height_, kHeight - y);
          for (int x = 0; x < kWidth; x += unit_width_) {
            const int width = std::min(unit_width_, kWidth - x);
            const Pixel* const source = src + y * kStride + x;
            target_self_guided_filter_func_(
                restoration_info_, source, kStride,
                source - kRestorationVerticalBorder * kStride, kStride,
                source + height * kStride, kStride, width, height,
                &restoration_buffer_, dst + y * kStride + x);
          }
        }
      }
      const absl::Duration elapsed_time = absl::Now() - start;
      test_utils::CheckMd5Digest(
          "kLoopRestorationTypeSgrProj", std::to_string(GetParam()).c_str(),
          kDigest[i][bd_index][radius_index], dst_ + kBorder * kStride,
          kHeight * kStride * sizeof(*dst_), elapsed_time);
    }
  }
}

using SelfGuidedFilterTest8bpp = SelfGuidedFilterTest<8, uint8_t>;

TEST_P(SelfGuidedFilterTest8bpp, Correctness) {
  TestFixedValues(0, 0);
  TestFixedValues(1, 1);
  TestFixedValues(2, 128);
  TestFixedValues(3, 255);
  TestRandomValues(false);
}

TEST_P(SelfGuidedFilterTest8bpp, DISABLED_Speed) { TestRandomValues(true); }

INSTANTIATE_TEST_SUITE_P(C, SelfGuidedFilterTest8bpp,
                         testing::ValuesIn(kUnitWidths));
#if LIBGAV1_ENABLE_AVX2
INSTANTIATE_TEST_SUITE_P(AVX2, SelfGuidedFilterTest8bpp,
                         testing::ValuesIn(kUnitWidths));
#endif
#if LIBGAV1_ENABLE_SSE4_1
INSTANTIATE_TEST_SUITE_P(SSE41, SelfGuidedFilterTest8bpp,
                         testing::ValuesIn(kUnitWidths));
#endif
#if LIBGAV1_ENABLE_NEON
INSTANTIATE_TEST_SUITE_P(NEON, SelfGuidedFilterTest8bpp,
                         testing::ValuesIn(kUnitWidths));
#endif

#if LIBGAV1_MAX_BITDEPTH >= 10
using SelfGuidedFilterTest10bpp = SelfGuidedFilterTest<10, uint16_t>;

TEST_P(SelfGuidedFilterTest10bpp, Correctness) {
  TestFixedValues(0, 0);
  TestFixedValues(1, 1);
  TestFixedValues(2, 512);
  TestFixedValues(3, 1023);
  TestRandomValues(false);
}

TEST_P(SelfGuidedFilterTest10bpp, DISABLED_Speed) { TestRandomValues(true); }

INSTANTIATE_TEST_SUITE_P(C, SelfGuidedFilterTest10bpp,
                         testing::ValuesIn(kUnitWidths));

#if LIBGAV1_ENABLE_AVX2
INSTANTIATE_TEST_SUITE_P(AVX2, SelfGuidedFilterTest10bpp,
                         testing::ValuesIn(kUnitWidths));
#endif
#if LIBGAV1_ENABLE_SSE4_1
INSTANTIATE_TEST_SUITE_P(SSE41, SelfGuidedFilterTest10bpp,
                         testing::ValuesIn(kUnitWidths));
#endif

#endif  // LIBGAV1_MAX_BITDEPTH >= 10

template <int bitdepth, typename Pixel>
class WienerFilterTest : public testing::TestWithParam<int>,
                         public test_utils::MaxAlignedAllocable {
 public:
  WienerFilterTest() = default;
  WienerFilterTest(const WienerFilterTest&) = delete;
  WienerFilterTest& operator=(const WienerFilterTest&) = delete;
  ~WienerFilterTest() override = default;

  void SetUp() override {
    test_utils::ResetDspTable(bitdepth);
    LoopRestorationInit_C();
    const Dsp* const dsp = GetDspTable(bitdepth);
    ASSERT_NE(dsp, nullptr);
    base_wiener_filter_func_ = dsp->loop_restorations[0];
    const testing::TestInfo* const test_info =
        testing::UnitTest::GetInstance()->current_test_info();
    const char* const test_case = test_info->test_suite_name();
    if (absl::StartsWith(test_case, "C/")) {
    } else if (absl::StartsWith(test_case, "AVX2/")) {
      if ((GetCpuInfo() & kAVX2) != 0) {
        LoopRestorationInit_AVX2();
#if LIBGAV1_MAX_BITDEPTH >= 10
        LoopRestorationInit10bpp_AVX2();
#endif
      }
    } else if (absl::StartsWith(test_case, "SSE41/")) {
      if ((GetCpuInfo() & kSSE4_1) != 0) {
        LoopRestorationInit_SSE4_1();
#if LIBGAV1_MAX_BITDEPTH >= 10
        LoopRestorationInit10bpp_SSE4_1();
#endif
      }
    } else if (absl::StartsWith(test_case, "NEON/")) {
      LoopRestorationInit_NEON();
    } else {
      FAIL() << "Unrecognized architecture prefix in test case name: "
             << test_case;
    }
    target_wiener_filter_func_ = dsp->loop_restorations[0];
    restoration_info_.type = kLoopRestorationTypeWiener;
    memset(dst_, 0, sizeof(dst_));
    memset(tmp_, 0, sizeof(tmp_));
    memset(buffer_, 0, sizeof(buffer_));
  }

  static void CleanFilterByOrder(const int order,
                                 int16_t filter[kWienerFilterTaps]) {
    if (order <= 5) filter[0] = 0;
    if (order <= 3) filter[1] = 0;
    if (order <= 1) filter[2] = 0;
  }

  void SetInputData(int type, Pixel value, int vertical_order,
                    int horizontal_order);
  void TestFixedValues(int digest_id, Pixel value);
  void TestRandomValues(bool speed);
  void TestCompare2C();

 protected:
  const int unit_width_ = GetParam();
  const int unit_height_ = kRestorationUnitHeight;

 private:
  alignas(kMaxAlignment)
      uint16_t buffer_[(kRestorationUnitWidth + kWienerFilterTaps - 1) *
                       kRestorationUnitHeight];
  alignas(kMaxAlignment) Pixel src_[kMaxBlockSize];
  alignas(kMaxAlignment) Pixel dst_[kMaxBlockSize];
  alignas(kMaxAlignment) Pixel tmp_[kMaxBlockSize];
  RestorationUnitInfo restoration_info_;
  RestorationBuffer restoration_buffer_;
  LoopRestorationFunc base_wiener_filter_func_;
  LoopRestorationFunc target_wiener_filter_func_;
};

template <int bitdepth, typename Pixel>
void WienerFilterTest<bitdepth, Pixel>::SetInputData(
    int type, Pixel value, const int vertical_order,
    const int horizontal_order) {
  const int mask = (1 << bitdepth) - 1;
  libvpx_test::ACMRandom rnd(libvpx_test::ACMRandom::DeterministicSeed());
  if (type == 0) {
    for (auto& s : src_) s = value;
  } else {
    for (auto& s : src_) s = rnd.Rand16() & mask;
  }
  int order = vertical_order;
  for (int i = WienerInfo::kVertical; i <= WienerInfo::kHorizontal; ++i) {
    auto& filter = restoration_info_.wiener_info.filter[i];
    filter[3] = 128;
    for (int j = 0; j < 3; ++j) {
      filter[j] = kWienerTapsMin[j] +
                  rnd.PseudoUniform(kWienerTapsMax[j] - kWienerTapsMin[j] + 1);
    }
    CleanFilterByOrder(order, filter);
    filter[3] -= 2 * (filter[0] + filter[1] + filter[2]);
    restoration_info_.wiener_info.number_leading_zero_coefficients[i] =
        (kWienerFilterTaps - order) / 2;
    order = horizontal_order;
  }
}

template <int bitdepth, typename Pixel>
void WienerFilterTest<bitdepth, Pixel>::TestFixedValues(int digest_id,
                                                        Pixel value) {
  static const char* const kDigest[2][4] = {
      {"74fc90760a14b13340cb718f200ba350", "5bacaca0128cd36f4805330b3787771d",
       "1109e17545cc4fbd5810b8b77e19fc36", "e7f914ec9d065aba92338016e17a526c"},
      {"c8cc38790ceb0bea1eb989686755e1e5", "70f573b7e8875262c638a68d2f317916",
       "193b19065899c835cb513149eb36d135", "f1dff65e3e53558b303ef0a2e3f3ba98"}};
  if (target_wiener_filter_func_ == nullptr) return;
  ASSERT_LT(value, 1 << bitdepth);
  constexpr int bd_index = (bitdepth == 8) ? 0 : 1;
  const Pixel* const src = src_ + kOffset;
  Pixel* const dst = dst_ + kOffset;
  for (const auto vertical_order : kWienerOrders) {
    for (const auto horizontal_order : kWienerOrders) {
      SetInputData(0, value, vertical_order, horizontal_order);
      memset(dst_, 0, sizeof(dst_));
      const absl::Time start = absl::Now();
      for (int y = 0; y < kHeight; y += unit_height_) {
        const int height = std::min(unit_height_, kHeight - y);
        for (int x = 0; x < kWidth; x += unit_width_) {
          const int width = std::min(unit_width_, kWidth - x);
          const Pixel* const source = src + y * kStride + x;
          target_wiener_filter_func_(
              restoration_info_, source, kStride,
              source - kRestorationVerticalBorder * kStride, kStride,
              source + height * kStride, kStride, width, height,
              &restoration_buffer_, dst + y * kStride + x);
        }
      }
      const absl::Duration elapsed_time = absl::Now() - start;
      test_utils::CheckMd5Digest(
          "kLoopRestorationTypeWiener", std::to_string(GetParam()).c_str(),
          kDigest[bd_index][digest_id], dst_, sizeof(dst_), elapsed_time);
    }
  }
}

template <int bitdepth, typename Pixel>
void WienerFilterTest<bitdepth, Pixel>::TestRandomValues(bool speed) {
  static const char* const kDigest[2][kNumWienerOrders][kNumWienerOrders] = {
      {{"40d0cf56d2ffb4f581e68b0fc97f547f", "5c04745209b684ba98004ebb0f64e70b",
        "545ed7d3f7e7ca3b86b4ada31f7aaee7", "0d6b2967f1bd1d99b720e563fe0cf03f"},
       {"44b37076f0cf27f6eb506aca50c1d3e4", "e927d64dc9249e05a65e10ee75baa7d9",
        "6136ecb4e29b17c9566504148943fd47", "c5ee2da81d44dc8cb2ac8021f724eb7a"},
       {"125cbb227313ec91a2683f26e6f049d1", "77671b6529c806d23b749f304b548f59",
        "28d53a1b486881895b8f73fa64486df1", "f5e32165bafe575d7ee7a6fbae75f36d"},
       {"e832c41f2566ab542b32abba9d4f27bd", "ab1336ee6b85cba651f35ee5d3b3cc5c",
        "52a673b6d14fbdca5ebdb1a34ee3326f",
        "ebb42c7c9111f2e39f21e2158e801d9e"}},
      {{"8cd9c6bd9983bd49564a58ed4af9098a", "f71f333c9d71237ed4e46f0ef2283196",
        "375b43abc1d6682d62f91c1841b8b0fc", "71e2444822ae9c697ddfc96e07c6e8a1"},
       {"d9ed3a66ceef405c08c87f6e91b71059", "c171fcff5fb7bb919f13ead7a4917a4c",
        "8fbd1edb82fcd78d4d286886f65a700a", "fe14a143e6b261c5bb07b179d40be5a2"},
       {"1c995f4e7f117857de73211b81093bd0", "5ab1ee3bb14adcd66d66802d58bee068",
        "d77430783e173ebd1b30e5d9336c8b69", "e159a3620747458dff7ed3d20da1a4b7"},
       {"5346fa07d195c257548a332753b057a3", "c77674bc0a638abc4d38d58e494fc7cf",
        "7cbc1562a9dd08e1973b3b9ac1afc765",
        "3c91bf1a34672cd40bf261c5820d3ec3"}}};
  if (target_wiener_filter_func_ == nullptr) return;
  constexpr int bd_index = (bitdepth == 8) ? 0 : 1;
  const int num_tests = speed ? 100000 : 1;
  const Pixel* const src = src_ + kOffset;
  Pixel* const dst = dst_ + kOffset;
  for (const auto vertical_order : kWienerOrders) {
    for (const auto horizontal_order : kWienerOrders) {
      SetInputData(1, (1 << bitdepth) - 1, vertical_order, horizontal_order);
      memset(dst_, 0, sizeof(dst_));
      const absl::Time start = absl::Now();
      for (int i = 0; i < num_tests; ++i) {
        for (int y = 0; y < kHeight; y += unit_height_) {
          const int height = std::min(unit_height_, kHeight - y);
          for (int x = 0; x < kWidth; x += unit_width_) {
            const int width = std::min(unit_width_, kWidth - x);
            const Pixel* const source = src + y * kStride + x;
            target_wiener_filter_func_(
                restoration_info_, source, kStride,
                source - kRestorationVerticalBorder * kStride, kStride,
                source + height * kStride, kStride, width, height,
                &restoration_buffer_, dst + y * kStride + x);
          }
        }
      }
      const absl::Duration elapsed_time = absl::Now() - start;
      test_utils::CheckMd5Digest(
          "kLoopRestorationTypeWiener", std::to_string(GetParam()).c_str(),
          kDigest[bd_index][kWienerOrderIdLookup[vertical_order]]
                 [kWienerOrderIdLookup[horizontal_order]],
          dst_, sizeof(dst_), elapsed_time);
    }
  }
}

template <int bitdepth, typename Pixel>
void WienerFilterTest<bitdepth, Pixel>::TestCompare2C() {
  if (base_wiener_filter_func_ == nullptr) return;
  if (target_wiener_filter_func_ == nullptr) return;
  if (base_wiener_filter_func_ == target_wiener_filter_func_) return;
  const Pixel* const src = src_ + kOffset;
  Pixel* const dst = dst_ + kOffset;
  Pixel* const tmp = tmp_ + kOffset;
  for (const auto vertical_order : kWienerOrders) {
    for (const auto horizontal_order : kWienerOrders) {
      SetInputData(1, (1 << bitdepth) - 1, vertical_order, horizontal_order);
      for (int x = 0; x < 2; ++x) {
        // Prepare min/max filter coefficients.
        int order = vertical_order;
        for (int i = WienerInfo::kVertical; i <= WienerInfo::kHorizontal; ++i) {
          auto& filter = restoration_info_.wiener_info.filter[i];
          for (int j = 0; j < 3; ++j) {
            filter[j] = (x == 0) ? kWienerTapsMin[j] : kWienerTapsMax[j];
          }
          CleanFilterByOrder(order, filter);
          filter[3] = 128 - 2 * (filter[0] + filter[1] + filter[2]);
          restoration_info_.wiener_info.number_leading_zero_coefficients[i] =
              (kWienerFilterTaps - order) / 2;
          order = horizontal_order;
        }
        base_wiener_filter_func_(restoration_info_, src, kStride,
                                 src - kRestorationVerticalBorder * kStride,
                                 kStride, src + unit_height_ * kStride, kStride,
                                 unit_width_, unit_height_,
                                 &restoration_buffer_, dst);
        target_wiener_filter_func_(restoration_info_, src, kStride,
                                   src - kRestorationVerticalBorder * kStride,
                                   kStride, src + unit_height_ * kStride,
                                   kStride, unit_width_, unit_height_,
                                   &restoration_buffer_, tmp);
        if (!test_utils::CompareBlocks(dst, tmp, unit_width_, unit_height_,
                                       kStride, kStride, false, false)) {
          ADD_FAILURE() << "Mismatch -- wiener taps min/max";
        }
      }
    }
  }
}

using WienerFilterTest8bpp = WienerFilterTest<8, uint8_t>;

TEST_P(WienerFilterTest8bpp, Correctness) {
  TestFixedValues(0, 0);
  TestFixedValues(1, 1);
  TestFixedValues(2, 128);
  TestFixedValues(3, 255);
  TestRandomValues(false);
}

TEST_P(WienerFilterTest8bpp, DISABLED_Speed) { TestRandomValues(true); }

TEST_P(WienerFilterTest8bpp, TestCompare2C) { TestCompare2C(); }

INSTANTIATE_TEST_SUITE_P(C, WienerFilterTest8bpp,
                         testing::ValuesIn(kUnitWidths));
#if LIBGAV1_ENABLE_AVX2
INSTANTIATE_TEST_SUITE_P(AVX2, WienerFilterTest8bpp,
                         testing::ValuesIn(kUnitWidths));
#endif
#if LIBGAV1_ENABLE_SSE4_1
INSTANTIATE_TEST_SUITE_P(SSE41, WienerFilterTest8bpp,
                         testing::ValuesIn(kUnitWidths));
#endif
#if LIBGAV1_ENABLE_NEON
INSTANTIATE_TEST_SUITE_P(NEON, WienerFilterTest8bpp,
                         testing::ValuesIn(kUnitWidths));
#endif

#if LIBGAV1_MAX_BITDEPTH >= 10
using WienerFilterTest10bpp = WienerFilterTest<10, uint16_t>;

TEST_P(WienerFilterTest10bpp, Correctness) {
  TestFixedValues(0, 0);
  TestFixedValues(1, 1);
  TestFixedValues(2, 512);
  TestFixedValues(3, 1023);
  TestRandomValues(false);
}

TEST_P(WienerFilterTest10bpp, DISABLED_Speed) { TestRandomValues(true); }

TEST_P(WienerFilterTest10bpp, TestCompare2C) { TestCompare2C(); }

INSTANTIATE_TEST_SUITE_P(C, WienerFilterTest10bpp,
                         testing::ValuesIn(kUnitWidths));

#if LIBGAV1_ENABLE_AVX2
INSTANTIATE_TEST_SUITE_P(AVX2, WienerFilterTest10bpp,
                         testing::ValuesIn(kUnitWidths));
#endif
#if LIBGAV1_ENABLE_SSE4_1
INSTANTIATE_TEST_SUITE_P(SSE41, WienerFilterTest10bpp,
                         testing::ValuesIn(kUnitWidths));
#endif

#endif  // LIBGAV1_MAX_BITDEPTH >= 10

}  // namespace
}  // namespace dsp
}  // namespace libgav1