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

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

namespace test_vsetvl {
namespace {

using namespace test_v_helpers;

const uint64_t VLEN = 512u;
const uint64_t ELEN = 32u;

uint32_t AVLS[] = {1,    2,    3,     4,     5,     8,    16,   17,
                   32,   36,   55,    64,    100,   128,  256,  321,
                   512,  623,  1024,  1100,  1543,  2048, 3052, 4096,
                   5555, 8192, 10241, 16384, 24325, 32768};
const int32_t AVL_COUNT = sizeof(AVLS) / sizeof(AVLS[0]);

#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) {
  for (int i = 0; i < AVL_COUNT; i++) {
    size_t 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++) {
    uint32_t vl = set_vsetvl(sew, vlmul, AVLS[i], false, false);
    EXPECT_EQ(vl, calculate_vl(width, AVLS[i], 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 test_vsetvl