#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 vmv_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];


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

  // Initialize test_vector1 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;
  }
  T *ptr_vec_1 = reinterpret_cast<T *>(test_vector_1);
  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);
  for (int idx = 0; idx < vl; idx++) {
    ptr_vec_1[idx] = idx;
  }
  return std::make_tuple(vlmax, vl);
}

template <typename T>
static void vmv_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 VmvTest : public ::testing::Test {
 protected:
  void SetUp() override { zero_vector_registers(); }
  void TearDown() override { zero_vector_registers(); }
};

TEST_F(VmvTest, vmv_demo) {
  for (int i = 0; i < AVL_COUNT; i++) {
    int32_t avl = AVLS[i];
    int vlmax;
    int vl;
    std::tie(vlmax, vl) = vmv_test_setup<uint8_t>(VLMUL::LMUL_M1, avl);
    if (avl > vlmax) {
      continue;
    }

    __asm__ volatile("vle8.v v0, (%0)" : : "r"(test_vector_1));
    __asm__ volatile("vmv.v.v v1, v0");
    __asm__ volatile("vse8.v v1, (%0)" : : "r"(test_vector_2));
    vmv_check<uint8_t>(vlmax);
  }
}

// TODO(henryherman): clang vle intrinsic uses vs1r.v (unsupported in renode)
TEST_F(VmvTest, DISABLED_intrinsic_vmv_demo) {
  for (int i = 0; i < AVL_COUNT; i++) {
    int32_t avl = AVLS[i];
    int vlmax;
    int vl;
    std::tie(vlmax, vl) = vmv_test_setup<uint8_t>(VLMUL::LMUL_M1, avl);
    if (avl > vlmax) {
      continue;
    }
    vint8m1_t vec1 = vle8_v_i8m1(reinterpret_cast<int8_t *>(test_vector_1), vl);
    vint8m1_t vec2 = vmv_v_v_i8m1(vec1, vl);
    int8_t *ptr_vec_2 = reinterpret_cast<int8_t *>(test_vector_2);
    vse8_v_i8m1(ptr_vec_2, vec2, vl);
    vmv_check<uint8_t>(vlmax);
  }
}

#define DEFINE_TEST_VMV_V_V_I_INTRINSIC(_SEW_, _LMUL_)                         \
  TEST_F(VmvTest, DISABLED_intrinsic_vmv_v_v_i##_SEW_##m##_LMUL_) {            \
    for (int i = 0; i < AVL_COUNT; i++) {                                      \
      int32_t avl = AVLS[i];                                                   \
      int vlmax;                                                               \
      int vl;                                                                  \
      std::tie(vlmax, vl) =                                                    \
          vmv_test_setup<uint##_SEW_##_t>(VLMUL::LMUL_M##_LMUL_, avl);         \
      if (avl > vlmax) {                                                       \
        continue;                                                              \
      }                                                                        \
      vint##_SEW_##m##_LMUL_##_t vec1 = vle##_SEW_##_v_i##_SEW_##m##_LMUL_(    \
          reinterpret_cast<int##_SEW_##_t *>(test_vector_1), vl);              \
      vint##_SEW_##m##_LMUL_##_t vec2 = vmv_v_v_i##_SEW_##m##_LMUL_(vec1, vl); \
      int##_SEW_##_t *ptr_vec_2 =                                              \
          reinterpret_cast<int##_SEW_##_t *>(test_vector_2);                   \
      vse##_SEW_##_v_i##_SEW_##m##_LMUL_(ptr_vec_2, vec2, vl);                 \
      vmv_check<uint##_SEW_##_t>(vlmax);                                       \
    }                                                                          \
  }

DEFINE_TEST_VMV_V_V_I_INTRINSIC(8, 1)
DEFINE_TEST_VMV_V_V_I_INTRINSIC(8, 2)
DEFINE_TEST_VMV_V_V_I_INTRINSIC(8, 4)
DEFINE_TEST_VMV_V_V_I_INTRINSIC(8, 8)

DEFINE_TEST_VMV_V_V_I_INTRINSIC(16, 1)
DEFINE_TEST_VMV_V_V_I_INTRINSIC(16, 2)
DEFINE_TEST_VMV_V_V_I_INTRINSIC(16, 4)
DEFINE_TEST_VMV_V_V_I_INTRINSIC(16, 8)

DEFINE_TEST_VMV_V_V_I_INTRINSIC(32, 1)
DEFINE_TEST_VMV_V_V_I_INTRINSIC(32, 2)
DEFINE_TEST_VMV_V_V_I_INTRINSIC(32, 4)
DEFINE_TEST_VMV_V_V_I_INTRINSIC(32, 8)

#define DEFINE_TEST_VMV_V_V_I(_SEW_, _LMUL_)                           \
  TEST_F(VmvTest, vmv_v_v_i##_SEW_##m##_LMUL_) {                       \
    for (int i = 0; i < AVL_COUNT; i++) {                              \
      int32_t avl = AVLS[i];                                           \
      int vlmax;                                                       \
      int vl;                                                          \
      std::tie(vlmax, vl) =                                            \
          vmv_test_setup<uint##_SEW_##_t>(VLMUL::LMUL_M##_LMUL_, avl); \
      if (avl > vlmax) {                                               \
        continue;                                                      \
      }                                                                \
      uint##_SEW_##_t *ptr_vec_1 =                                     \
          reinterpret_cast<uint##_SEW_##_t *>(test_vector_1);          \
      uint##_SEW_##_t *ptr_vec_2 =                                     \
          reinterpret_cast<uint##_SEW_##_t *>(test_vector_2);          \
      __asm__ volatile("vle" #_SEW_ ".v v0, (%0)" : : "r"(ptr_vec_1)); \
      __asm__ volatile("vmv.v.v v8, v0");                              \
      __asm__ volatile("vse" #_SEW_ ".v v8, (%0)" : : "r"(ptr_vec_2)); \
      vmv_check<uint##_SEW_##_t>(vlmax);                               \
    }                                                                  \
  }

DEFINE_TEST_VMV_V_V_I(8, 1)
DEFINE_TEST_VMV_V_V_I(8, 2)
DEFINE_TEST_VMV_V_V_I(8, 4)
DEFINE_TEST_VMV_V_V_I(8, 8)

DEFINE_TEST_VMV_V_V_I(16, 1)
DEFINE_TEST_VMV_V_V_I(16, 2)
DEFINE_TEST_VMV_V_V_I(16, 4)
DEFINE_TEST_VMV_V_V_I(16, 8)

DEFINE_TEST_VMV_V_V_I(32, 1)
DEFINE_TEST_VMV_V_V_I(32, 2)
DEFINE_TEST_VMV_V_V_I(32, 4)
DEFINE_TEST_VMV_V_V_I(32, 8)

TEST_F(VmvTest, vmv_v_x_demo) {
  for (int i = 0; i < AVL_COUNT; i++) {
    int32_t avl = AVLS[i];
    int vlmax;
    int vl;
    std::tie(vlmax, vl) = vmv_test_setup<uint8_t>(VLMUL::LMUL_M1, avl);
    if (avl > vlmax) {
      continue;
    }
    uint8_t *ptr_vec_1 = reinterpret_cast<uint8_t *>(test_vector_1);
    uint8_t *ptr_vec_2 = reinterpret_cast<uint8_t *>(test_vector_2);
    uint8_t test_val = 0xAB;
    __asm__ volatile("vmv.v.x v8, %[RS1]" ::[RS1] "r"(test_val));
    for (int i = 0; i < vl; i++) {
      ptr_vec_1[i] = test_val;
    }
    __asm__ volatile("vse8.v v8, (%0)" : : "r"(ptr_vec_2));
    vmv_check<uint8_t>(vlmax);
  }
}

#define DEFINE_TEST_VMV_V_X_I(_SEW_, _LMUL_, TEST_VAL)                 \
  TEST_F(VmvTest, vmv_v_x_e##_SEW_##m##_LMUL_) {                       \
    for (int i = 0; i < AVL_COUNT; i++) {                              \
      int32_t avl = AVLS[i];                                           \
      int vlmax;                                                       \
      int vl;                                                          \
      std::tie(vlmax, vl) =                                            \
          vmv_test_setup<uint##_SEW_##_t>(VLMUL::LMUL_M##_LMUL_, avl); \
      if (avl > vlmax) {                                               \
        continue;                                                      \
      }                                                                \
      uint##_SEW_##_t *ptr_vec_1 =                                     \
          reinterpret_cast<uint##_SEW_##_t *>(test_vector_1);          \
      uint##_SEW_##_t *ptr_vec_2 =                                     \
          reinterpret_cast<uint##_SEW_##_t *>(test_vector_2);          \
      const uint##_SEW_##_t test_val = TEST_VAL;                       \
      __asm__ volatile("vmv.v.x v8, %[RS1]" ::[RS1] "r"(test_val));    \
      for (int i = 0; i < vl; i++) {                                   \
        ptr_vec_1[i] = test_val;                                       \
      }                                                                \
      __asm__ volatile("vse" #_SEW_ ".v v8, (%0)" : : "r"(ptr_vec_2)); \
      vmv_check<uint##_SEW_##_t>(vlmax);                               \
    }                                                                  \
  }

DEFINE_TEST_VMV_V_X_I(8, 1, 0xab)
DEFINE_TEST_VMV_V_X_I(8, 2, 0xac)
DEFINE_TEST_VMV_V_X_I(8, 4, 0xad)
DEFINE_TEST_VMV_V_X_I(8, 8, 0xae)

DEFINE_TEST_VMV_V_X_I(16, 1, 0xabc1)
DEFINE_TEST_VMV_V_X_I(16, 2, 0xabc2)
DEFINE_TEST_VMV_V_X_I(16, 4, 0xabc3)
DEFINE_TEST_VMV_V_X_I(16, 8, 0xabc4)

DEFINE_TEST_VMV_V_X_I(32, 1, 0xabcdef12)
DEFINE_TEST_VMV_V_X_I(32, 2, 0xabcdef13)
DEFINE_TEST_VMV_V_X_I(32, 4, 0xabcdef14)
DEFINE_TEST_VMV_V_X_I(32, 8, 0xabcdef15)

TEST_F(VmvTest, vmv_v_i_demo) {
  for (int i = 0; i < AVL_COUNT; i++) {
    int32_t avl = AVLS[i];
    int vlmax;
    int vl;
    std::tie(vlmax, vl) = vmv_test_setup<int8_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);
    int8_t test_val = -12;
    __asm__ volatile("vmv.v.i v8, -12" ::);
    for (int i = 0; i < vl; i++) {
      ptr_vec_1[i] = test_val;
    }
    __asm__ volatile("vse8.v v8, (%0)" : : "r"(ptr_vec_2));
    vmv_check<uint8_t>(vlmax);
  }
}

// TODO(gkielian): Allow mechanism for multiple tests for same sew,lmul pair
#define DEFINE_TEST_VMV_V_I_I(_SEW_, _LMUL_, TEST_VAL)                 \
  TEST_F(VmvTest, vmv_v_i_e##_SEW_##m##_LMUL_) {                       \
    for (int i = 0; i < AVL_COUNT; i++) {                              \
      int32_t avl = AVLS[i];                                           \
      int vlmax;                                                       \
      int vl;                                                          \
      std::tie(vlmax, vl) =                                            \
          vmv_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);           \
      int##_SEW_##_t test_val = TEST_VAL;                              \
      __asm__ volatile("vmv.v.i v8, " #TEST_VAL);                      \
      for (int i = 0; i < vl; i++) {                                   \
        ptr_vec_1[i] = test_val;                                       \
      }                                                                \
      __asm__ volatile("vse" #_SEW_ ".v v8, (%0)" : : "r"(ptr_vec_2)); \
      vmv_check<int##_SEW_##_t>(vlmax);                                \
    }                                                                  \
  }

DEFINE_TEST_VMV_V_I_I(8, 1, -11)
DEFINE_TEST_VMV_V_I_I(8, 2, -12)
DEFINE_TEST_VMV_V_I_I(8, 4, -13)
DEFINE_TEST_VMV_V_I_I(8, 8, -14)

DEFINE_TEST_VMV_V_I_I(16, 1, -10)
DEFINE_TEST_VMV_V_I_I(16, 2, -9)
DEFINE_TEST_VMV_V_I_I(16, 4, -8)
DEFINE_TEST_VMV_V_I_I(16, 8, -7)

DEFINE_TEST_VMV_V_I_I(32, 1, -2)
DEFINE_TEST_VMV_V_I_I(32, 2, -3)
DEFINE_TEST_VMV_V_I_I(32, 4, -4)
DEFINE_TEST_VMV_V_I_I(32, 8, -5)

}  // namespace
}  // namespace vmv_test