/* SPDX-License-Identifier: zlib */
/*
 * This code is cleaned up and modernized version of https://github.com/to-miz/sse_mathfun_extension
 */

#pragma once

#include <emmintrin.h>
#include <xmmintrin.h>

#include <array>
#include <complex>
#include <type_traits>
#include <memory>
#include <utility>

namespace simd_complex_conversions
{
// Before C++20 there is no standard and correct way to implement bit_cast.
// It was verified that MSVC generates the correct code for the current use
// cases
namespace details
{
template <class To, class From>
std::enable_if_t<
   sizeof(To) == sizeof(From) && std::is_trivially_copyable_v<From> &&
      std::is_trivially_copyable_v<To>,
   To>
// constexpr support needs compiler magic
bit_cast(const From& src) noexcept
{
   static_assert(
      std::is_trivially_constructible_v<To>,
      "This implementation additionally requires "
      "destination type to be trivially constructible");

   To dst;
   std::memcpy(&dst, &src, sizeof(To));
   return dst;
}

constexpr float PIF = 3.141592653589793238f;
constexpr float PIO2F = 1.5707963267948966192f;

constexpr float cephes_PIF = 3.141592653589793238f;
constexpr float cephes_PIO2F = 1.5707963267948966192f;
constexpr float cephes_PIO4F = 0.7853981633974483096f;
constexpr float cephes_FOPI = 1.27323954473516f; // 4 / M_PI
constexpr float minus_cephes_DP1 = -0.78515625f;
constexpr float minus_cephes_DP2 = -2.4187564849853515625e-4f;
constexpr float minus_cephes_DP3 = -3.77489497744594108e-8f;
constexpr float sincof_p0 = -1.9515295891e-4f;
constexpr float sincof_p1 = 8.3321608736e-3f;
constexpr float sincof_p2 = -1.6666654611e-1f;
constexpr float coscof_p0 = 2.443315711809948e-005f;
constexpr float coscof_p1 = -1.388731625493765e-003f;
constexpr float coscof_p2 = 4.166664568298827e-002f;

constexpr float atancof_p0 = 8.05374449538e-2f;
constexpr float atancof_p1 = 1.38776856032e-1f;
constexpr float atancof_p2 = 1.99777106478e-1f;
constexpr float atancof_p3 = 3.33329491539e-1f;

static const float sign_mask = bit_cast<float>(0x80000000);
static const float inv_sign_mask = bit_cast<float>(~0x80000000);
} // namespace details

inline __m128 atan_ps(__m128 x)
{
   using namespace details;

   __m128 sign_bit, y;

   sign_bit = x;
   /* take the absolute value */
   x = _mm_and_ps(x, _mm_set1_ps(inv_sign_mask));
   /* extract the sign bit (upper one) */
   sign_bit = _mm_and_ps(sign_bit, _mm_set1_ps(sign_mask));

   /* range reduction, init x and y depending on range */
   /* x > 2.414213562373095 */
   __m128 cmp0 = _mm_cmpgt_ps(x, _mm_set1_ps(2.414213562373095f));
   /* x > 0.4142135623730950 */
   __m128 cmp1 = _mm_cmpgt_ps(x, _mm_set1_ps(0.4142135623730950f));

   /* x > 0.4142135623730950 && !( x > 2.414213562373095 ) */
   __m128 cmp2 = _mm_andnot_ps(cmp0, cmp1);

   /* -( 1.0/x ) */
   __m128 y0 = _mm_and_ps(cmp0, _mm_set1_ps(cephes_PIO2F));
   __m128 x0 = _mm_div_ps(_mm_set1_ps(1.0f), x);
   x0 = _mm_xor_ps(x0, _mm_set1_ps(sign_mask));

   __m128 y1 = _mm_and_ps(cmp2, _mm_set1_ps(cephes_PIO4F));
   /* (x-1.0)/(x+1.0) */
   __m128 x1_o = _mm_sub_ps(x, _mm_set1_ps(1.0f));
   __m128 x1_u = _mm_add_ps(x, _mm_set1_ps(1.0f));
   __m128 x1 = _mm_div_ps(x1_o, x1_u);

   __m128 x2 = _mm_and_ps(cmp2, x1);
   x0 = _mm_and_ps(cmp0, x0);
   x2 = _mm_or_ps(x2, x0);
   cmp1 = _mm_or_ps(cmp0, cmp2);
   x2 = _mm_and_ps(cmp1, x2);
   x = _mm_andnot_ps(cmp1, x);
   x = _mm_or_ps(x2, x);

   y = _mm_or_ps(y0, y1);

   __m128 zz = _mm_mul_ps(x, x);
   __m128 acc = _mm_set1_ps(atancof_p0);
   acc = _mm_mul_ps(acc, zz);
   acc = _mm_sub_ps(acc, _mm_set1_ps(atancof_p1));
   acc = _mm_mul_ps(acc, zz);
   acc = _mm_add_ps(acc, _mm_set1_ps(atancof_p2));
   acc = _mm_mul_ps(acc, zz);
   acc = _mm_sub_ps(acc, _mm_set1_ps(atancof_p3));
   acc = _mm_mul_ps(acc, zz);
   acc = _mm_mul_ps(acc, x);
   acc = _mm_add_ps(acc, x);
   y = _mm_add_ps(y, acc);

   /* update the sign */
   y = _mm_xor_ps(y, sign_bit);

   return y;
}

inline __m128 atan2_ps(__m128 y, __m128 x)
{
   using namespace details;

   __m128 zero = _mm_setzero_ps();
   __m128 x_eq_0 = _mm_cmpeq_ps(x, zero);
   __m128 x_gt_0 = _mm_cmpgt_ps(x, zero);
   __m128 x_le_0 = _mm_cmple_ps(x, zero);
   __m128 y_eq_0 = _mm_cmpeq_ps(y, zero);
   __m128 x_lt_0 = _mm_cmplt_ps(x, zero);
   __m128 y_lt_0 = _mm_cmplt_ps(y, zero);

   __m128 zero_mask = _mm_and_ps(x_eq_0, y_eq_0);
   __m128 zero_mask_other_case = _mm_and_ps(y_eq_0, x_gt_0);
   zero_mask = _mm_or_ps(zero_mask, zero_mask_other_case);

   __m128 pio2_mask = _mm_andnot_ps(y_eq_0, x_eq_0);
   __m128 pio2_mask_sign = _mm_and_ps(y_lt_0, _mm_set1_ps(sign_mask));
   __m128 pio2_result = _mm_set1_ps(cephes_PIO2F);
   pio2_result = _mm_xor_ps(pio2_result, pio2_mask_sign);
   pio2_result = _mm_and_ps(pio2_mask, pio2_result);

   __m128 pi_mask = _mm_and_ps(y_eq_0, x_lt_0);
   __m128 pi = _mm_set1_ps(cephes_PIF);
   __m128 pi_result = _mm_and_ps(pi_mask, pi);

   __m128 swap_sign_mask_offset = _mm_and_ps(x_lt_0, y_lt_0);
   swap_sign_mask_offset =
      _mm_and_ps(swap_sign_mask_offset, _mm_set1_ps(sign_mask));

   __m128 offset0 = _mm_setzero_ps();
   __m128 offset1 = _mm_set1_ps(cephes_PIF);
   offset1 = _mm_xor_ps(offset1, swap_sign_mask_offset);

   __m128 offset = _mm_andnot_ps(x_lt_0, offset0);
   offset = _mm_and_ps(x_lt_0, offset1);

   __m128 arg = _mm_div_ps(y, x);
   __m128 atan_result = atan_ps(arg);
   atan_result = _mm_add_ps(atan_result, offset);

   /* select between zero_result, pio2_result and atan_result */

   __m128 result = _mm_andnot_ps(zero_mask, pio2_result);
   atan_result = _mm_andnot_ps(zero_mask, atan_result);
   atan_result = _mm_andnot_ps(pio2_mask, atan_result);
   result = _mm_or_ps(result, atan_result);
   result = _mm_or_ps(result, pi_result);

   return result;
}

inline std::pair<__m128, __m128> sincos_ps(__m128 x)
{
   using namespace details;
   __m128 xmm1, xmm2, xmm3 = _mm_setzero_ps(), sign_bit_sin, y;
   __m128i emm0, emm2, emm4;

   sign_bit_sin = x;
   /* take the absolute value */
   x = _mm_and_ps(x, _mm_set1_ps(inv_sign_mask));
   /* extract the sign bit (upper one) */
   sign_bit_sin = _mm_and_ps(sign_bit_sin, _mm_set1_ps(sign_mask));

   /* scale by 4/Pi */
   y = _mm_mul_ps(x, _mm_set1_ps(cephes_FOPI));

   /* store the integer part of y in emm2 */
   emm2 = _mm_cvttps_epi32(y);

   /* j=(j+1) & (~1) (see the cephes sources) */
   emm2 = _mm_add_epi32(emm2, _mm_set1_epi32(1));
   emm2 = _mm_and_si128(emm2, _mm_set1_epi32(~1));
   y = _mm_cvtepi32_ps(emm2);

   emm4 = emm2;

   /* get the swap sign flag for the sine */
   emm0 = _mm_and_si128(emm2, _mm_set1_epi32(4));
   emm0 = _mm_slli_epi32(emm0, 29);
   __m128 swap_sign_bit_sin = _mm_castsi128_ps(emm0);

   /* get the polynom selection mask for the sine*/
   emm2 = _mm_and_si128(emm2, _mm_set1_epi32(2));
   emm2 = _mm_cmpeq_epi32(emm2, _mm_setzero_si128());
   __m128 poly_mask = _mm_castsi128_ps(emm2);

   /* The magic pass: "Extended precision modular arithmetic"
      x = ((x - y * DP1) - y * DP2) - y * DP3; */
   xmm1 = _mm_set1_ps(minus_cephes_DP1);
   xmm2 = _mm_set1_ps(minus_cephes_DP2);
   xmm3 = _mm_set1_ps(minus_cephes_DP3);
   xmm1 = _mm_mul_ps(y, xmm1);
   xmm2 = _mm_mul_ps(y, xmm2);
   xmm3 = _mm_mul_ps(y, xmm3);
   x = _mm_add_ps(x, xmm1);
   x = _mm_add_ps(x, xmm2);
   x = _mm_add_ps(x, xmm3);

   emm4 = _mm_sub_epi32(emm4, _mm_set1_epi32(2));
   emm4 = _mm_andnot_si128(emm4, _mm_set1_epi32(4));
   emm4 = _mm_slli_epi32(emm4, 29);
   __m128 sign_bit_cos = _mm_castsi128_ps(emm4);

   sign_bit_sin = _mm_xor_ps(sign_bit_sin, swap_sign_bit_sin);

   /* Evaluate the first polynom  (0 <= x <= Pi/4) */
   __m128 z = _mm_mul_ps(x, x);
   y = _mm_set1_ps(coscof_p0);

   y = _mm_mul_ps(y, z);
   y = _mm_add_ps(y, _mm_set1_ps(coscof_p1));
   y = _mm_mul_ps(y, z);
   y = _mm_add_ps(y, _mm_set1_ps(coscof_p2));
   y = _mm_mul_ps(y, z);
   y = _mm_mul_ps(y, z);
   __m128 tmp = _mm_mul_ps(z, _mm_set1_ps(0.5f));
   y = _mm_sub_ps(y, tmp);
   y = _mm_add_ps(y, _mm_set1_ps(1));

   /* Evaluate the second polynom  (Pi/4 <= x <= 0) */

   __m128 y2 = _mm_set1_ps(sincof_p0);
   y2 = _mm_mul_ps(y2, z);
   y2 = _mm_add_ps(y2, _mm_set1_ps(sincof_p1));
   y2 = _mm_mul_ps(y2, z);
   y2 = _mm_add_ps(y2, _mm_set1_ps(sincof_p2));
   y2 = _mm_mul_ps(y2, z);
   y2 = _mm_mul_ps(y2, x);
   y2 = _mm_add_ps(y2, x);

   /* select the correct result from the two polynoms */
   xmm3 = poly_mask;
   __m128 ysin2 = _mm_and_ps(xmm3, y2);
   __m128 ysin1 = _mm_andnot_ps(xmm3, y);
   y2 = _mm_sub_ps(y2, ysin2);
   y = _mm_sub_ps(y, ysin1);

   xmm1 = _mm_add_ps(ysin1, ysin2);
   xmm2 = _mm_add_ps(y, y2);

   /* update the sign */
   return std::make_pair(
      _mm_xor_ps(xmm1, sign_bit_sin), _mm_xor_ps(xmm2, sign_bit_cos));
}

inline float atan2_ss(float y, float x)
{
   return _mm_cvtss_f32(atan2_ps(_mm_set_ss(y), _mm_set_ss(x)));
}

inline std::pair<float, float> sincos_ss(float angle)
{
   auto res = sincos_ps(_mm_set_ss(angle));
   return std::make_pair(_mm_cvtss_f32(res.first), _mm_cvtss_f32(res.second));
}

inline __m128 norm(__m128 x, __m128 y)
{
   return _mm_add_ps(_mm_mul_ps(x, x), _mm_mul_ps(y, y));
}

inline float sqrt_ss(float x)
{
   __m128 sse_value = _mm_set_ss(x);
   sse_value = _mm_sqrt_ss(sse_value);
   return _mm_cvtss_f32(sse_value);
}

template <typename fnc>
void perform_parallel_simd_aligned(
   const std::complex<float>* input, float* output, int n, const fnc& f)
{
   for (int i = 0; i <= n - 4; i += 4)
   {
      // Safe according to C++ standard
      auto p1 = _mm_load_ps(reinterpret_cast<const float*>(input + i));
      auto p2 = _mm_load_ps(reinterpret_cast<const float*>(input + i + 2));

      // p1 = {real(c1), imag(c1), real(c2), imag(c2)}
      // p2 = {real(c3), imag(c3), real(c4), imag(c4)}

      auto rp = _mm_shuffle_ps(p1, p2, _MM_SHUFFLE(2, 0, 2, 0));
      auto ip = _mm_shuffle_ps(p1, p2, _MM_SHUFFLE(3, 1, 3, 1));

      __m128 out;
      f(rp, ip, out);

      _mm_store_ps(output + i, out);
   }
   // deal with last partial packet
   for (int i = n & (~3); i < n; ++i)
   {
      __m128 out;
      f(_mm_set_ss(real(input[i])), _mm_set_ss(imag(input[i])), out);
      output[i] = _mm_cvtss_f32(out);
   }
}

inline void rotate_parallel_simd_aligned(
   const float* oldPhase, const float* newPhase, std::complex<float>* output,
   int n)
{
   for (int i = 0; i <= n - 4; i += 4)
   {
      auto [sin, cos] = sincos_ps(
         _mm_sub_ps(
            _mm_load_ps(newPhase + i),
            _mm_load_ps(oldPhase + i)));


      // Safe according to C++ standard
      auto p1 = _mm_load_ps(reinterpret_cast<float*>(output + i));
      auto p2 = _mm_load_ps(reinterpret_cast<float*>(output + i + 2));

      // p1 = {real(c1), imag(c1), real(c2), imag(c2)}
      // p2 = {real(c3), imag(c3), real(c4), imag(c4)}

      auto rp = _mm_shuffle_ps(p1, p2, _MM_SHUFFLE(2, 0, 2, 0));
      auto ip = _mm_shuffle_ps(p1, p2, _MM_SHUFFLE(3, 1, 3, 1));

      
      // We need to calculate (rp, ip) * (cos, sin) -> (rp*cos - ip*sin, rp*sin + ip*cos)

      auto out_rp = _mm_sub_ps(_mm_mul_ps(rp, cos), _mm_mul_ps(ip, sin));
      auto out_ip = _mm_add_ps(_mm_mul_ps(rp, sin), _mm_mul_ps(ip, cos));

      p1 = _mm_unpacklo_ps(out_rp, out_ip);
      p2 = _mm_unpackhi_ps(out_rp, out_ip);

      _mm_store_ps(reinterpret_cast<float*>(output + i), p1);
      _mm_store_ps(reinterpret_cast<float*>(output + i + 2), p2);
   }
   // deal with last partial packet
   for (int i = n & (~3); i < n; ++i)
   {
      auto theta = newPhase[i] - oldPhase[i];
      output[i] *= std::complex<float>(cosf(theta), sinf(theta));
   }
}

} // namespace simd_complex_conversions
