/*
 * Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
 *
 * SPDX-License-Identifier: MIT
 */

#pragma once

#include <hip/hip_runtime_api.h>

template <typename T> struct vec4_struct { using type = void; };

#define DEFINE_VEC4_OVERLOAD(base_type, vec_type)                                                  \
  template <> struct vec4_struct<base_type> { using type = vec_type; }

DEFINE_VEC4_OVERLOAD(char, char4);
DEFINE_VEC4_OVERLOAD(short, short4);
DEFINE_VEC4_OVERLOAD(int, int4);
DEFINE_VEC4_OVERLOAD(long, long4);
DEFINE_VEC4_OVERLOAD(long long, longlong4);

DEFINE_VEC4_OVERLOAD(unsigned char, uchar4);
DEFINE_VEC4_OVERLOAD(unsigned short, ushort4);
DEFINE_VEC4_OVERLOAD(unsigned int, uint4);
DEFINE_VEC4_OVERLOAD(unsigned long, ulong4);
DEFINE_VEC4_OVERLOAD(unsigned long long, ulonglong4);

DEFINE_VEC4_OVERLOAD(float, float4);
DEFINE_VEC4_OVERLOAD(double, float4);

template <typename T> using vec4 = typename vec4_struct<T>::type;

template <typename T> inline void SetVec4(vec4<T>& vec, const T val) {
  vec.x = val;
  vec.y = val;
  vec.z = val;
  vec.w = val;
}

template <typename T>
inline void SetVec4(vec4<T>& vec, const T x, const T y, const T z, const T w) {
  vec.x = x;
  vec.y = y;
  vec.z = z;
  vec.w = w;
}

template <typename T> inline auto MakeVec4(const T val) {
  vec4<T> vec;
  SetVec4(vec, val);

  return vec;
}

template <typename T> inline void MakeVec4(const T x, const T y, const T z, const T w) {
  vec4<T> vec;
  SetVec4(vec, x, y, z, w);

  return vec;
}

template <typename T> inline vec4<float> Vec4Map(const T& vec) {
  vec4<float> ret;
  ret.x = NormalizeInteger(vec.x);
  ret.y = NormalizeInteger(vec.y);
  ret.z = NormalizeInteger(vec.z);
  ret.w = NormalizeInteger(vec.w);
  return ret;
}

template <typename T> inline __host__ __device__ auto Vec4Scale(float s, const T& vec) {
  T ret;
  ret.x = s * vec.x;
  ret.y = s * vec.y;
  ret.z = s * vec.z;
  ret.w = s * vec.w;

  return ret;
}

template <typename T> inline __host__ __device__ auto Vec4Add(const T& vec1, const T& vec2) {
  T ret;
  ret.x = vec1.x + vec2.x;
  ret.y = vec1.y + vec2.y;
  ret.z = vec1.z + vec2.z;
  ret.w = vec1.w + vec2.w;

  return ret;
}