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

#include <arm_neon.h>
#include <arm_sve.h>

#include <assert.h>
#include <stdint.h>

#include "aom_dsp/arm/aom_neon_sve_bridge.h"
#include "aom_dsp/arm/mem_neon.h"
#include "aom_dsp/arm/sum_neon.h"
#include "aom_dsp/arm/transpose_neon.h"
#include "av1/encoder/arm/pickrst_neon.h"
#include "av1/encoder/arm/pickrst_sve.h"
#include "av1/encoder/pickrst.h"

static inline uint16_t highbd_find_average_sve(const uint16_t *src,
                                               int src_stride, int width,
                                               int height) {
  uint64x2_t avg_u64 = vdupq_n_u64(0);
  uint16x8_t ones = vdupq_n_u16(1);

  // Use a predicate to compute the last columns.
  svbool_t pattern = svwhilelt_b16_u32(0, width % 8 == 0 ? 8 : width % 8);

  int h = height;
  do {
    int j = width;
    const uint16_t *src_ptr = src;
    while (j > 8) {
      uint16x8_t s = vld1q_u16(src_ptr);
      avg_u64 = aom_udotq_u16(avg_u64, s, ones);

      j -= 8;
      src_ptr += 8;
    }
    uint16x8_t s_end = svget_neonq_u16(svld1_u16(pattern, src_ptr));
    avg_u64 = aom_udotq_u16(avg_u64, s_end, ones);

    src += src_stride;
  } while (--h != 0);
  return (uint16_t)(vaddvq_u64(avg_u64) / (width * height));
}

static inline void sub_avg_block_highbd_sve(const uint16_t *buf, int buf_stride,
                                            int16_t avg, int width, int height,
                                            int16_t *buf_avg,
                                            int buf_avg_stride) {
  uint16x8_t avg_u16 = vdupq_n_u16(avg);

  // Use a predicate to compute the last columns.
  svbool_t pattern = svwhilelt_b16_u32(0, width % 8 == 0 ? 8 : width % 8);

  uint16x8_t avg_end = svget_neonq_u16(svdup_n_u16_z(pattern, avg));

  do {
    int j = width;
    const uint16_t *buf_ptr = buf;
    int16_t *buf_avg_ptr = buf_avg;
    while (j > 8) {
      uint16x8_t d = vld1q_u16(buf_ptr);
      vst1q_s16(buf_avg_ptr, vreinterpretq_s16_u16(vsubq_u16(d, avg_u16)));

      j -= 8;
      buf_ptr += 8;
      buf_avg_ptr += 8;
    }
    uint16x8_t d_end = svget_neonq_u16(svld1_u16(pattern, buf_ptr));
    vst1q_s16(buf_avg_ptr, vreinterpretq_s16_u16(vsubq_u16(d_end, avg_end)));

    buf += buf_stride;
    buf_avg += buf_avg_stride;
  } while (--height > 0);
}

void av1_compute_stats_highbd_sve(int32_t wiener_win, const uint8_t *dgd8,
                                  const uint8_t *src8, int16_t *dgd_avg,
                                  int16_t *src_avg, int32_t h_start,
                                  int32_t h_end, int32_t v_start, int32_t v_end,
                                  int32_t dgd_stride, int32_t src_stride,
                                  int64_t *M, int64_t *H,
                                  aom_bit_depth_t bit_depth) {
  const int32_t wiener_win2 = wiener_win * wiener_win;
  const int32_t wiener_halfwin = (wiener_win >> 1);
  const uint16_t *src = CONVERT_TO_SHORTPTR(src8);
  const uint16_t *dgd = CONVERT_TO_SHORTPTR(dgd8);
  const int32_t width = h_end - h_start;
  const int32_t height = v_end - v_start;
  const int32_t d_stride = (width + 2 * wiener_halfwin + 15) & ~15;
  const int32_t s_stride = (width + 15) & ~15;

  const uint16_t *dgd_start = dgd + h_start + v_start * dgd_stride;
  const uint16_t *src_start = src + h_start + v_start * src_stride;
  const uint16_t avg =
      highbd_find_average_sve(dgd_start, dgd_stride, width, height);

  sub_avg_block_highbd_sve(src_start, src_stride, avg, width, height, src_avg,
                           s_stride);
  sub_avg_block_highbd_sve(
      dgd + (v_start - wiener_halfwin) * dgd_stride + h_start - wiener_halfwin,
      dgd_stride, avg, width + 2 * wiener_halfwin, height + 2 * wiener_halfwin,
      dgd_avg, d_stride);

  if (wiener_win == WIENER_WIN) {
    compute_stats_win7_sve(dgd_avg, d_stride, src_avg, s_stride, width, height,
                           M, H);
  } else {
    assert(wiener_win == WIENER_WIN_CHROMA);
    compute_stats_win5_sve(dgd_avg, d_stride, src_avg, s_stride, width, height,
                           M, H);
  }

  // H is a symmetric matrix, so we only need to fill out the upper triangle.
  // We can copy it down to the lower triangle outside the (i, j) loops.
  if (bit_depth == AOM_BITS_8) {
    diagonal_copy_stats_neon(wiener_win2, H);
  } else if (bit_depth == AOM_BITS_10) {  // bit_depth == EB_TEN_BIT
    const int32_t k4 = wiener_win2 & ~3;

    int32_t k = 0;
    do {
      int64x2_t dst = div4_neon(vld1q_s64(M + k));
      vst1q_s64(M + k, dst);
      dst = div4_neon(vld1q_s64(M + k + 2));
      vst1q_s64(M + k + 2, dst);
      H[k * wiener_win2 + k] /= 4;
      k += 4;
    } while (k < k4);

    H[k * wiener_win2 + k] /= 4;

    for (; k < wiener_win2; ++k) {
      M[k] /= 4;
    }

    div4_diagonal_copy_stats_neon(wiener_win2, H);
  } else {  // bit_depth == AOM_BITS_12
    const int32_t k4 = wiener_win2 & ~3;

    int32_t k = 0;
    do {
      int64x2_t dst = div16_neon(vld1q_s64(M + k));
      vst1q_s64(M + k, dst);
      dst = div16_neon(vld1q_s64(M + k + 2));
      vst1q_s64(M + k + 2, dst);
      H[k * wiener_win2 + k] /= 16;
      k += 4;
    } while (k < k4);

    H[k * wiener_win2 + k] /= 16;

    for (; k < wiener_win2; ++k) {
      M[k] /= 16;
    }

    div16_diagonal_copy_stats_neon(wiener_win2, H);
  }
}