#include <riscv_vector.h>
#include <springbok.h>
#include <stdio.h>
#include <stdlib.h>

#include <bit>
#include <tuple>

#include "pw_unit_test/framework.h"
#include "test_v_helpers.h"

namespace vadd_vi_test {
namespace {

const int LMUL_MAX = 8;
const int VLEN = 512;
const int MAXVL_BYTES = VLEN / sizeof(uint8_t) * LMUL_MAX;

using namespace test_v_helpers;

uint8_t test_vector_1[MAXVL_BYTES];
uint8_t test_vector_2[MAXVL_BYTES];

static void zero_vector_registers() {
  // Clear all vector registers
  int vlmax = get_vsetvlmax_intrinsic(VSEW::SEW_E32, VLMUL::LMUL_M8);
  set_vsetvl_intrinsic(VSEW::SEW_E32, VLMUL::LMUL_M8, vlmax);
  __asm__ volatile("vmv.v.i v0, 0");
  __asm__ volatile("vmv.v.i v8, 0");
  __asm__ volatile("vmv.v.i v16, 0");
  __asm__ volatile("vmv.v.i v24, 0");
}

template <typename T>
static std::tuple<int, int> vadd_vi_test_setup(VLMUL lmul, int32_t avl) {
  // Clear all vector registers
  zero_vector_registers();

  // Initialize test_vector_1 and determine vl, vlmax
  uint32_t bw = std::__bit_width(sizeof(T));
  VSEW sew = static_cast<VSEW>(bw - 1);
  int vlmax = get_vsetvlmax_intrinsic(sew, lmul);
  if (avl > vlmax) {
    avl = vlmax;
  }
  memset(test_vector_1, 0, MAXVL_BYTES);
  memset(test_vector_2, 0, MAXVL_BYTES);
  int vl = set_vsetvl_intrinsic(sew, lmul, avl);
  EXPECT_EQ(avl, vl);

  return std::make_tuple(vlmax, vl);
}

template <typename T>
static void vadd_vi_check(int avl) {
  T *ptr_vec_1 = reinterpret_cast<T *>(test_vector_1);
  T *ptr_vec_2 = reinterpret_cast<T *>(test_vector_2);
  for (int idx = 0; idx < avl; idx++) {
    ASSERT_EQ(ptr_vec_1[idx], ptr_vec_2[idx]);
  }
}

class VaddViTest : public ::testing::Test {
 protected:
  void SetUp() override { zero_vector_registers(); }
  void TearDown() override { zero_vector_registers(); }
};

TEST_F(VaddViTest, vadd_vi_demo) {
  for (int i = 0; i < AVL_COUNT; i++) {
    int32_t avl = AVLS[i];
    int vlmax;
    int vl;
    std::tie(vlmax, vl) = vadd_vi_test_setup<uint8_t>(VLMUL::LMUL_M1, avl);
    if (avl > vlmax) {
      continue;
    }
    int8_t *ptr_vec_1 = reinterpret_cast<int8_t *>(test_vector_1);
    int8_t *ptr_vec_2 = reinterpret_cast<int8_t *>(test_vector_2);
    const int8_t test_val_1 = 1;
    const int8_t test_val_2 = -3;
    __asm__ volatile("vadd.vi v8, v8, %[SIMM5]" ::[SIMM5] "n"(test_val_1));
    __asm__ volatile("vadd.vi v8, v8, %[SIMM5]" ::[SIMM5] "n"(test_val_2));
    for (int i = 0; i < vl; i++) {
      ptr_vec_1[i] = test_val_1 + test_val_2;
    }
    __asm__ volatile("vse8.v v8, (%0)" : : "r"(ptr_vec_2));
    vadd_vi_check<int8_t>(vlmax);
  }
}

#define DEFINE_TEST_VADD_VI(_SEW_, _LMUL_, TEST_VAL_1, TEST_VAL_2)            \
  TEST_F(VaddViTest, vadd_vi##_SEW_##m##_LMUL_) {                             \
    for (int i = 0; i < AVL_COUNT; i++) {                                     \
      int32_t avl = AVLS[i];                                                  \
      int vlmax;                                                              \
      int vl;                                                                 \
      std::tie(vlmax, vl) =                                                   \
          vadd_vi_test_setup<int##_SEW_##_t>(VLMUL::LMUL_M##_LMUL_, avl);     \
      if (avl > vlmax) {                                                      \
        continue;                                                             \
      }                                                                       \
      int##_SEW_##_t *ptr_vec_1 =                                             \
          reinterpret_cast<int##_SEW_##_t *>(test_vector_1);                  \
      int##_SEW_##_t *ptr_vec_2 =                                             \
          reinterpret_cast<int##_SEW_##_t *>(test_vector_2);                  \
      const int##_SEW_##_t test_val_1 = TEST_VAL_1;                           \
      const int##_SEW_##_t test_val_2 = TEST_VAL_2;                           \
      __asm__ volatile("vadd.vi v8, v8, %[SIMM5]" ::[SIMM5] "n"(test_val_1)); \
      __asm__ volatile("vadd.vi v8, v8, %[SIMM5]" ::[SIMM5] "n"(test_val_2)); \
      for (int i = 0; i < vl; i++) {                                          \
        ptr_vec_1[i] = test_val_1 + test_val_2;                               \
      }                                                                       \
      __asm__ volatile("vse" #_SEW_ ".v v8, (%0)" : : "r"(ptr_vec_2));        \
      vadd_vi_check<int##_SEW_##_t>(vlmax);                                   \
    }                                                                         \
  }

// TODO(gkielian): modify macro to permit more than one test per sew/lmul pair
DEFINE_TEST_VADD_VI(8, 1, -16, 15)
DEFINE_TEST_VADD_VI(8, 2, -2, 3)
DEFINE_TEST_VADD_VI(8, 4, 15, 15)
DEFINE_TEST_VADD_VI(8, 8, -15, -15)

DEFINE_TEST_VADD_VI(16, 1, -16, 15)
DEFINE_TEST_VADD_VI(16, 2, -2, 3)
DEFINE_TEST_VADD_VI(16, 4, 15, 15)
DEFINE_TEST_VADD_VI(16, 8, -15, -15)

DEFINE_TEST_VADD_VI(32, 1, -16, 15)
DEFINE_TEST_VADD_VI(32, 2, -2, 3)
DEFINE_TEST_VADD_VI(32, 4, 15, 15)
DEFINE_TEST_VADD_VI(32, 8, -15, -15)

}  // namespace
}  // namespace vadd_vi_test