#include "hex8_pair_simd.h"

#if defined(FESA_HEX8_PAIR_AVX2) && defined(_MSC_VER) && defined(_M_X64)
#include <cstring>
#include <immintrin.h>
#include <intrin.h>
#include <limits>

namespace fesa::hex8_internal {
namespace {
struct VectorPair {
  __m256d high;
  __m256d low;
};

// TU-local predicates mirror the checked scalar premises without publishing
// AVX2-compiled common math or HighLow COMDAT definitions to baseline callers.
/// @brief Reads binary64 fields without FP classification or comparison.
std::uint64_t ScalarBits(double value) {
  static_assert(sizeof(double) == sizeof(std::uint64_t) &&
                std::numeric_limits<double>::is_iec559 &&
                std::numeric_limits<double>::digits == 53);
  std::uint64_t bits = 0;
  std::memcpy(&bits, &value, sizeof(bits));
  return bits;
}
/// @brief Bitwise finite gate before any scalar FP comparison.
bool FiniteScalar(double value) {
  return (ScalarBits(value) & 0x7ff0000000000000ULL) !=
         0x7ff0000000000000ULL;
}
/// @brief Exact least-bit exponent; caller has checked finite/nonzero first.
int LowestBitExponentScalar(double value) {
  const auto bits = ScalarBits(value);
  const auto encoded_exponent = (bits >> 52U) & 0x7ffU;
  auto significand = bits & 0xfffffffffffffULL;
  const int exponent = encoded_exponent == 0U
                           ? -1074
                           : static_cast<int>(encoded_exponent) - 1075;
  if (encoded_exponent != 0U) significand |= 1ULL << 52U;
  unsigned long trailing = 0;
  _BitScanForward64(&trailing, significand);
  return exponent + static_cast<int>(trailing);
}

/// @brief Proves the exact tail bound using integer exponents only.
/// @pre Both operands are finite and nonzero; false requires the exact scans.
bool TailBoundGuaranteed(double left, double right) {
  const auto left_exponent = (ScalarBits(left) >> 52U) & 0x7ffU;
  const auto right_exponent = (ScalarBits(right) >> 52U) & 0x7ffU;
  const auto left_adjusted = left_exponent == 0 ? 1U : left_exponent;
  const auto right_adjusted = right_exponent == 0 ? 1U : right_exponent;
  // LBE >= adjusted exponent - 1075, including nonzero subnormals.
  return left_adjusted + right_adjusted >= 1076U;
}

/// @brief Selects each lane without evaluating additional FP arithmetic.
__m256d Select(__m256d mask, __m256d selected, __m256d other) {
  return _mm256_blendv_pd(other, selected, mask);
}
/// @brief Produces independent finite masks without signaling NaN compares.
__m256d Finite(__m256d value) {
  const auto exponent = _mm256_set1_epi64x(0x7ff0000000000000LL);
  const auto nonfinite = _mm256_cmpeq_epi64(
      _mm256_and_si256(_mm256_castpd_si256(value), exponent), exponent);
  return _mm256_castsi256_pd(_mm256_xor_si256(
      nonfinite, _mm256_set1_epi64x(-1)));
}
/// @brief Recognizes both signed zero encodings using integer fields.
__m256d Zero(__m256d value) {
  return _mm256_castsi256_pd(_mm256_cmpeq_epi64(
      _mm256_and_si256(_mm256_castpd_si256(value),
                       _mm256_set1_epi64x(0x7fffffffffffffffLL)),
      _mm256_setzero_si256()));
}
/// @brief Preserves finite/canonical-zero/RN64 gates for active lanes only.
/// @note Both sum operands and equality RHS are neutral before FP evaluation.
__m256d Normalized(VectorPair value, __m256d active) {
  const auto finite_active = _mm256_and_pd(active,
      _mm256_and_pd(Finite(value.high), Finite(value.low)));
  const auto both_zero = _mm256_and_pd(Zero(value.high), Zero(value.low));
  const auto canonical_zero = _mm256_castsi256_pd(_mm256_and_si256(
      _mm256_cmpeq_epi64(_mm256_castpd_si256(value.high),
                         _mm256_setzero_si256()),
      _mm256_cmpeq_epi64(_mm256_castpd_si256(value.low),
                         _mm256_setzero_si256())));
  const auto sum_active = _mm256_andnot_pd(both_zero, finite_active);
  const auto high = Select(sum_active, value.high, _mm256_setzero_pd());
  const auto low = Select(sum_active, value.low, _mm256_setzero_pd());
  const auto equality = _mm256_cmp_pd(_mm256_add_pd(high, low), high, _CMP_EQ_OQ);
  return _mm256_and_pd(finite_active,
      Select(both_zero, canonical_zero, equality));
}
/// @brief Retains left-to-right normalized short-circuit for every lane.
bool NormalizedInputs(VectorPair left, VectorPair right) {
  const auto all = _mm256_castsi256_pd(_mm256_set1_epi64x(-1));
  const auto left_valid = Normalized(left, all);
  return _mm256_movemask_pd(Normalized(right, left_valid)) == 15;
}

/// @brief Preserves valid lanes and encodes invalid lanes as quiet NaN/+0.
VectorPair InvalidOutside(VectorPair value, __m256d valid) {
  return {Select(valid, value.high, _mm256_castsi256_pd(
              _mm256_set1_epi64x(0x7ff8000000000000LL))),
          Select(valid, value.low, _mm256_setzero_pd())};
}
/// @brief Loads raw high/low fields in lane order 0..3, with no normalization.
VectorPair Load(const PairPacket& value) {
  return {_mm256_setr_pd(value.values[0].high, value.values[1].high,
                         value.values[2].high, value.values[3].high),
          _mm256_setr_pd(value.values[0].low, value.values[1].low,
                         value.values[2].low, value.values[3].low)};
}
/// @brief Replicates one raw first product without changing either component.
VectorPair Splat(PacketPair value) {
  return {_mm256_set1_pd(value.high), _mm256_set1_pd(value.low)};
}
/// @brief Stores raw high/low fields without rounding or lane reduction.
PairPacket Store(VectorPair value) {
  double high[4];
  double low[4];
  _mm256_storeu_pd(high, value.high);
  _mm256_storeu_pd(low, value.low);
  return {{{high[0], low[0]}, {high[1], low[1]},
           {high[2], low[2]}, {high[3], low[3]}}};
}
/// @brief Retains finite-sum error association; false requests original graph.
/// @note False leaves output unreadable and preserves the necessary sum flags.
/// Prefix flags still require optimized scalar-oracle/codegen validation.
bool TwoSum(__m256d left, __m256d right, VectorPair* output) {
  const auto sum = _mm256_add_pd(left, right);
  // Do not speculate Inf-Inf error arithmetic beyond a failed first sum.
  // The caller retains original inputs for its existing baseline bridge.
  if (_mm256_movemask_pd(Finite(sum)) != 15) return false;
  const auto bridge = _mm256_sub_pd(sum, left);
  const auto restored = _mm256_sub_pd(sum, bridge);
  const auto error = _mm256_add_pd(_mm256_sub_pd(left, restored),
                                  _mm256_sub_pd(right, bridge));
  const auto valid = _mm256_and_pd(_mm256_and_pd(Finite(sum), Finite(bridge)),
      _mm256_and_pd(Finite(restored), Finite(error)));
  const auto exact_zero = _mm256_and_pd(Zero(sum), Zero(error));
  *output = InvalidOutside({Select(exact_zero, _mm256_setzero_pd(), sum),
                            Select(exact_zero, _mm256_setzero_pd(), error)}, valid);
  return true;
}
/// @brief Routes exceptional packets to baseline checked Add in lane order.
VectorPair ScalarAdd(VectorPair left, VectorPair right, bool pair_shortcuts) {
  const auto a = Store(left);
  const auto b = Store(right);
  PairPacket result;
  ScalarPacketAdd(a, b, pair_shortcuts, &result);
  return Load(result);
}
/// @brief Routes exceptional packets to baseline checked MultiplyPair.
VectorPair ScalarMultiply(const PairPacket& left, const PairPacket& right) {
  PairPacket result;
  ScalarPacketMultiply(left, right, &result);
  return Load(result);
}
/// @brief Retains normalized/zero branches and the checked per-lane Add graph.
VectorPair Add(VectorPair left, VectorPair right, bool pair_shortcuts) {
  if (!NormalizedInputs(left, right))
    return ScalarAdd(left, right, pair_shortcuts);
  const int left_zero = _mm256_movemask_pd(
      _mm256_and_pd(Zero(left.high), Zero(left.low)));
  const int right_zero = _mm256_movemask_pd(
      _mm256_and_pd(Zero(right.high), Zero(right.low)));
  if (pair_shortcuts) {
    if (left_zero == 15)
      return {right.high, Select(Zero(right.low), _mm256_setzero_pd(), right.low)};
    if (right_zero == 15)
      return {left.high, Select(Zero(left.low), _mm256_setzero_pd(), left.low)};
    if (left_zero != 0 || right_zero != 0)
      return ScalarAdd(left, right, true);
  }
  VectorPair head;
  VectorPair tail;
  VectorPair middle;
  VectorPair joined;
  if (!TwoSum(left.high, right.high, &head) ||
      !TwoSum(left.low, right.low, &tail) ||
      !TwoSum(head.low, tail.high, &middle) ||
      !TwoSum(head.high, middle.high, &joined))
    return ScalarAdd(left, right, pair_shortcuts);
  const auto remainder = _mm256_add_pd(
      _mm256_add_pd(joined.low, middle.low), tail.low);
  const auto valid = _mm256_and_pd(Finite(remainder), _mm256_and_pd(
      _mm256_and_pd(Finite(head.high), Finite(head.low)), _mm256_and_pd(
      _mm256_and_pd(Finite(tail.high), Finite(tail.low)), _mm256_and_pd(
      _mm256_and_pd(Finite(middle.high), Finite(middle.low)),
      _mm256_and_pd(Finite(joined.high), Finite(joined.low))))));
  // Original Add returns before its last TwoSum when a prior finite gate fails.
  VectorPair result;
  if (!TwoSum(Select(valid, joined.high, _mm256_setzero_pd()),
              Select(valid, remainder, _mm256_setzero_pd()), &result))
    return ScalarAdd(left, right, pair_shortcuts);
  return InvalidOutside(result, valid);
}
/// @brief Checks finite/zero/tail premises before products; neutralizes rejected
/// lanes before the explicit residual FMA and final TwoSum.
VectorPair Product(__m256d left, __m256d right) {
  double a[4];
  double b[4];
  _mm256_storeu_pd(a, left);
  _mm256_storeu_pd(b, right);
  bool special = false;
  for (std::size_t lane = 0; lane < 4; ++lane)
    if (!FiniteScalar(a[lane]) ||
        !FiniteScalar(b[lane]) ||
        a[lane] == 0 || b[lane] == 0 ||
        (!TailBoundGuaranteed(a[lane], b[lane]) &&
         LowestBitExponentScalar(a[lane]) +
             LowestBitExponentScalar(b[lane]) < -1074))
      special = true;
  if (special) {
    PairPacket result;
    ScalarPacketTwoProduct(a, b, &result);
    return Load(result);
  }
  const auto product = _mm256_mul_pd(left, right);
  const auto product_valid = Finite(product);
  const auto one = _mm256_set1_pd(1.0);
  // Invalid products never reach the residual FMA in the original helper.
  const auto error = _mm256_fmsub_pd(Select(product_valid, left, one),
      Select(product_valid, right, one), Select(product_valid, product, one));
  const auto valid = _mm256_and_pd(product_valid, Finite(error));
  VectorPair result;
  if (!TwoSum(Select(valid, product, _mm256_setzero_pd()),
              Select(valid, error, _mm256_setzero_pd()), &result)) {
    PairPacket scalar;
    ScalarPacketTwoProduct(a, b, &scalar);
    return Load(scalar);
  }
  return InvalidOutside(result, valid);
}
/// @brief Uses uniform low-zero branches with original four-product association;
/// mixed or invalid packets preserve the baseline checked path.
VectorPair Multiply(VectorPair left, VectorPair right) {
  if (!NormalizedInputs(left, right))
    return ScalarMultiply(Store(left), Store(right));
  const int left_low_zero = _mm256_movemask_pd(Zero(left.low));
  const int right_low_zero = _mm256_movemask_pd(Zero(right.low));
  if ((left_low_zero != 0 && left_low_zero != 15) ||
      (right_low_zero != 0 && right_low_zero != 15))
    return ScalarMultiply(Store(left), Store(right));
  const int branch = (left_low_zero == 0 ? 1 : 0) |
                     (right_low_zero == 0 ? 2 : 0);
  auto product = Product(left.high, right.high);
  if (branch == 3) {
    product = Add(product, Product(left.high, right.low), false);
    product = Add(product, Product(left.low, right.high), false);
    return Add(product, Product(left.low, right.low), false);
  }
  const VectorPair zero{_mm256_setzero_pd(), _mm256_setzero_pd()};
  if ((branch & 2) != 0)
    product = Add(product, Product(left.high, right.low), true);
  else if ((branch & 1) != 0)
    product = Add(product, zero, true);
  if ((branch & 1) != 0)
    product = Add(product, Product(left.low, right.high), true);
  return Add(product, zero, true);
}
}  // namespace

void EvaluateMaterialPacketAvx2(const MaterialPacketInput& input,
                                PairPacket* output) {
  const auto first = Multiply(Load(input.deviatoric_left),
                              Load(input.deviatoric_right));
  const auto second = Multiply(Load(input.shear_left), Load(input.shear_right));
  const auto third = Multiply(Load(input.tensor_left), Load(input.gradient_dot));
  const auto material = Add(Add(first, second, true), third, true);
  const PairPacket weight{{input.weight, input.weight, input.weight, input.weight}};
  *output = Store(Add(Load(input.accumulated),
                      Multiply(material, Load(weight)), true));
}
void EvaluateGeometryPacketAvx2(const GeometryPacketInput& input,
                                GeometryPacketOutput* output) {
  VectorPair dot{_mm256_setzero_pd(), _mm256_setzero_pd()};
  for (std::size_t i = 0; i < 3; ++i)
    dot = Add(dot, Multiply(Splat(input.gradient_a[i]),
                            Load(input.gradient_b[i])), true);
  VectorPair geo{_mm256_setzero_pd(), _mm256_setzero_pd()};
  for (std::size_t j = 0; j < 3; ++j)
    for (std::size_t l = 0; l < 3; ++l)
      geo = Add(geo, Multiply(Splat(input.gradient_stress[j][l]),
                              Load(input.gradient_b[l])), true);
  output->gradient_dot = Store(dot);
  output->weighted_geometric = Store(Multiply(geo, Splat(input.weight)));
}
}  // namespace fesa::hex8_internal
#else
namespace fesa::hex8_internal {
void EvaluateMaterialPacketAvx2(const MaterialPacketInput& input,
                                PairPacket* output) {
  EvaluateMaterialPacketScalar(input, output);
}
void EvaluateGeometryPacketAvx2(const GeometryPacketInput& input,
                                GeometryPacketOutput* output) {
  EvaluateGeometryPacketScalar(input, output);
}
}  // namespace fesa::hex8_internal
#endif