#pragma once

#if defined(__cplusplus)
#include <cstring>
#endif

struct __half_raw {
  unsigned short x;
};

struct __half2_raw {
  unsigned short x;
  unsigned short y;
};

#if defined(__cplusplus)
struct __half;

__half __float2half(float);
float __half2float(__half);

// BEGIN STRUCT __HALF
struct __half {
 protected:
  unsigned short __x;

 public:
  // CREATORS
  __half() = default;
  __half(const __half_raw& x) : __x{x.x} {}
#if !defined(__HIP_NO_HALF_CONVERSIONS__)
  __half(float x) : __x{__float2half(x).__x} {}
  __half(double x) : __x{__float2half(x).__x} {}
#endif
  __half(const __half&) = default;
  __half(__half&&) = default;
  ~__half() = default;

  // MANIPULATORS
  __half& operator=(const __half&) = default;
  __half& operator=(__half&&) = default;
  __half& operator=(const __half_raw& x) {
    __x = x.x;
    return *this;
  }
#if !defined(__HIP_NO_HALF_CONVERSIONS__)
  __half& operator=(float x) {
    __x = __float2half(x).__x;
    return *this;
  }
  __half& operator=(double x) { return *this = static_cast<float>(x); }
#endif

  // ACCESSORS
  operator float() const { return __half2float(*this); }
  operator __half_raw() const { return __half_raw{__x}; }
};
// END STRUCT __HALF

// BEGIN STRUCT __HALF2
struct __half2 {
 public:
  __half x;
  __half y;

  // CREATORS
  __half2() = default;
  __half2(const __half2_raw& ix)
      : x{reinterpret_cast<const __half&>(ix.x)}, y{reinterpret_cast<const __half&>(ix.y)} {}
  __half2(const __half& ix, const __half& iy) : x{ix}, y{iy} {}
  __half2(const __half2&) = default;
  __half2(__half2&&) = default;
  ~__half2() = default;

  // MANIPULATORS
  __half2& operator=(const __half2&) = default;
  __half2& operator=(__half2&&) = default;
  __half2& operator=(const __half2_raw& ix) {
    x = reinterpret_cast<const __half_raw&>(ix.x);
    y = reinterpret_cast<const __half_raw&>(ix.y);
    return *this;
  }

  // ACCESSORS
  operator __half2_raw() const {
    return __half2_raw{reinterpret_cast<const unsigned short&>(x),
                       reinterpret_cast<const unsigned short&>(y)};
  }
};
// END STRUCT __HALF2

inline unsigned short __internal_float2half(float flt, unsigned int& sgn, unsigned int& rem) {
  unsigned int x{};
  std::memcpy(&x, &flt, sizeof(flt));

  unsigned int u = (x & 0x7fffffffU);
  sgn = ((x >> 16) & 0x8000U);

  // NaN/+Inf/-Inf
  if (u >= 0x7f800000U) {
    rem = 0;
    return static_cast<unsigned short>((u == 0x7f800000U) ? (sgn | 0x7c00U) : 0x7fffU);
  }
  // Overflows
  if (u > 0x477fefffU) {
    rem = 0x80000000U;
    return static_cast<unsigned short>(sgn | 0x7bffU);
  }
  // Normal numbers
  if (u >= 0x38800000U) {
    rem = u << 19;
    u -= 0x38000000U;
    return static_cast<unsigned short>(sgn | (u >> 13));
  }
  // +0/-0
  if (u < 0x33000001U) {
    rem = u;
    return static_cast<unsigned short>(sgn);
  }
  // Denormal numbers
  unsigned int exponent = u >> 23;
  unsigned int mantissa = (u & 0x7fffffU);
  unsigned int shift = 0x7eU - exponent;
  mantissa |= 0x800000U;
  rem = mantissa << (32 - shift);
  return static_cast<unsigned short>(sgn | (mantissa >> shift));
}

inline __half __float2half(float x) {
  __half_raw r;
  unsigned int sgn{};
  unsigned int rem{};
  r.x = __internal_float2half(x, sgn, rem);
  if (rem > 0x80000000U || (rem == 0x80000000U && (r.x & 0x1))) ++r.x;

  return r;
}

inline __half __float2half_rn(float x) { return __float2half(x); }

inline __half __float2half_rz(float x) {
  __half_raw r;
  unsigned int sgn{};
  unsigned int rem{};
  r.x = __internal_float2half(x, sgn, rem);

  return r;
}

inline __half __float2half_rd(float x) {
  __half_raw r;
  unsigned int sgn{};
  unsigned int rem{};
  r.x = __internal_float2half(x, sgn, rem);
  if (rem && sgn) ++r.x;

  return r;
}

inline __half __float2half_ru(float x) {
  __half_raw r;
  unsigned int sgn{};
  unsigned int rem{};
  r.x = __internal_float2half(x, sgn, rem);
  if (rem && !sgn) ++r.x;

  return r;
}

inline __half2 __float2half2_rn(float x) { return __half2{__float2half_rn(x), __float2half_rn(x)}; }

inline __half2 __floats2half2_rn(float x, float y) {
  return __half2{__float2half_rn(x), __float2half_rn(y)};
}

inline float __internal_half2float(unsigned short x) {
  unsigned int sign = ((x >> 15) & 1);
  unsigned int exponent = ((x >> 10) & 0x1f);
  unsigned int mantissa = ((x & 0x3ff) << 13);

  if (exponent == 0x1fU) { /* NaN or Inf */
    mantissa = (mantissa ? (sign = 0, 0x7fffffU) : 0);
    exponent = 0xffU;
  } else if (!exponent) { /* Denorm or Zero */
    if (mantissa) {
      unsigned int msb;
      exponent = 0x71U;
      do {
        msb = (mantissa & 0x400000U);
        mantissa <<= 1; /* normalize */
        --exponent;
      } while (!msb);
      mantissa &= 0x7fffffU; /* 1.mantissa is implicit */
    }
  } else {
    exponent += 0x70U;
  }
  unsigned int u = ((sign << 31) | (exponent << 23) | mantissa);
  float f;
  memcpy(&f, &u, sizeof(u));

  return f;
}

inline float __half2float(__half x) { return __internal_half2float(static_cast<__half_raw>(x).x); }
inline float2 __half22float2(__half2 x) {
  return float2{__internal_half2float(static_cast<__half2_raw>(x).x),
                __internal_half2float(static_cast<__half2_raw>(x).x)};
}

inline float __low2float(__half2 x) { return __internal_half2float(static_cast<__half2_raw>(x).x); }

inline float __high2float(__half2 x) {
  return __internal_half2float(static_cast<__half2_raw>(x).y);
}

#if !defined(HIP_NO_HALF)
using half = __half;
using half2 = __half2;
#endif
#endif  // defined(__cplusplus)
