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

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

namespace vsetvl_test {
namespace {

using namespace test_v_helpers;

const uint64_t VLEN = 512u;

#define MIN(X, Y) ((X) < (Y) ? (X) : (Y))

static uint32_t calculate_vl(uint32_t sew, uint32_t avl, float lmul) {
  uint32_t vlmax = (uint32_t)(VLEN * lmul / sew);
  return MIN(avl, vlmax);
}

static void test_vsetvl(VSEW sew, VLMUL vlmul, uint32_t width, float lmul) {
  size_t vl = 0;
  for (int i = 0; i < AVL_COUNT; i++) {
    vl = set_vsetvl_intrinsic(sew, vlmul, AVLS[i]);
    EXPECT_EQ(vl, calculate_vl(width, AVLS[i], lmul));
  }
  for (int i = 0; i < AVL_COUNT; i++) {
    vl = set_vsetvl(sew, vlmul, AVLS[i], false, false);
    EXPECT_EQ(vl, calculate_vl(width, AVLS[i], lmul));
  }
  for (int i = 0; i < AVL_COUNT; i++) {
    vl = set_vsetvli(sew, vlmul, AVLS[i]);
    EXPECT_EQ(vl, calculate_vl(width, AVLS[i], lmul));
  }
  vl = set_vsetivli(sew, vlmul);
  EXPECT_EQ(vl, calculate_vl(width, AVL_CONST, lmul));
}

static void test_vsetvlmax(VSEW sew, VLMUL vlmul, uint32_t width, float lmul) {
  EXPECT_EQ(get_vsetvlmax_intrinsic(sew, vlmul), (int)VLEN / width * lmul);
}

TEST(VsetvlTest, vsetvl_e8m1) {
  test_vsetvl(VSEW::SEW_E8, VLMUL::LMUL_M1, 8, 1.0);
}

TEST(VsetvlTest, vsetvl_e16m1) {
  test_vsetvl(VSEW::SEW_E16, VLMUL::LMUL_M1, 16, 1.0);
}

TEST(VsetvlTest, vsetvl_e32m1) {
  test_vsetvl(VSEW::SEW_E32, VLMUL::LMUL_M1, 32, 1.0);
}

TEST(VsetvlTest, vsetvl_e8m2) {
  test_vsetvl(VSEW::SEW_E8, VLMUL::LMUL_M2, 8, 2.0);
}

TEST(VsetvlTest, vsetvl_e16m2) {
  test_vsetvl(VSEW::SEW_E16, VLMUL::LMUL_M2, 16, 2.0);
}

TEST(VsetvlTest, vsetvl_e32m2) {
  test_vsetvl(VSEW::SEW_E32, VLMUL::LMUL_M2, 32, 2.0);
}

TEST(VsetvlTest, vsetvl_e8m4) {
  test_vsetvl(VSEW::SEW_E8, VLMUL::LMUL_M4, 8, 4.0);
}

TEST(VsetvlTest, vsetvl_e16m4) {
  test_vsetvl(VSEW::SEW_E16, VLMUL::LMUL_M4, 16, 4.0);
}

TEST(VsetvlTest, vsetvl_e32m4) {
  test_vsetvl(VSEW::SEW_E32, VLMUL::LMUL_M4, 32, 4.0);
}

TEST(VsetvlTest, vsetvl_e8m8) {
  test_vsetvl(VSEW::SEW_E8, VLMUL::LMUL_M8, 8, 8.0);
}

TEST(VsetvlTest, vsetvl_e16m8) {
  test_vsetvl(VSEW::SEW_E16, VLMUL::LMUL_M8, 16, 8);
}

TEST(VsetvlTest, vsetvl_e32m8) {
  test_vsetvl(VSEW::SEW_E32, VLMUL::LMUL_M8, 32, 8.0);
}

TEST(VsetvlTest, vsetvlmax_e8m1) {
  test_vsetvlmax(VSEW::SEW_E8, VLMUL::LMUL_M1, 8, 1.0);
}

TEST(VsetvlTest, vsetvlmax_e8m2) {
  test_vsetvlmax(VSEW::SEW_E8, VLMUL::LMUL_M2, 8, 2.0);
}

TEST(VsetvlTest, vsetvlmax_e8m4) {
  test_vsetvlmax(VSEW::SEW_E8, VLMUL::LMUL_M4, 8, 4.0);
}

TEST(VsetvlTest, vsetvlmax_e8m8) {
  test_vsetvlmax(VSEW::SEW_E8, VLMUL::LMUL_M8, 8, 8.0);
}

TEST(VsetvlTest, vsetvlmax_e16m1) {
  test_vsetvlmax(VSEW::SEW_E16, VLMUL::LMUL_M1, 16, 1.0);
}

TEST(VsetvlTest, vsetvlmax_e16m2) {
  test_vsetvlmax(VSEW::SEW_E16, VLMUL::LMUL_M2, 16, 2.0);
}

TEST(VsetvlTest, vsetvlmax_e16m4) {
  test_vsetvlmax(VSEW::SEW_E16, VLMUL::LMUL_M4, 16, 4.0);
}

TEST(VsetvlTest, vsetvlmax_e16m8) {
  test_vsetvlmax(VSEW::SEW_E16, VLMUL::LMUL_M8, 16, 8.0);
}

TEST(VsetvlTest, vsetvlmax_e32m1) {
  test_vsetvlmax(VSEW::SEW_E32, VLMUL::LMUL_M1, 32, 1.0);
}

TEST(VsetvlTest, vsetvlmax_e32m2) {
  test_vsetvlmax(VSEW::SEW_E32, VLMUL::LMUL_M2, 32, 2.0);
}

TEST(VsetvlTest, vsetvlmax_e32m4) {
  test_vsetvlmax(VSEW::SEW_E32, VLMUL::LMUL_M4, 32, 4.0);
}

TEST(VsetvlTest, vsetvlmax_e32m8) {
  test_vsetvlmax(VSEW::SEW_E32, VLMUL::LMUL_M8, 32, 8.0);
}

}  // namespace
}  // namespace vsetvl_test