// This file is part of Eigen, a lightweight C++ template library
// for linear algebra.
//
// Copyright (C) 2025 The Eigen Authors
//
// This Source Code Form is subject to the terms of the Mozilla
// Public License v. 2.0. If a copy of the MPL was not distributed
// with this file, You can obtain one at http://mozilla.org/MPL/2.0/.
// SPDX-License-Identifier: MPL-2.0

#include "main.h"

template <typename Scalar, typename Packet>
void verify_data(const Scalar* data_in, const Scalar* data_out, const Packet& a, Index begin, Index count) {
  constexpr int PacketSize = internal::unpacket_traits<Packet>::size;
  bool ok = true;
  for (Index i = begin; i < begin + count; i++) {
    ok = ok && numext::equal_strict(data_in[i], data_out[i]);
  }
  if (!ok) {
    std::cout << "begin: " << begin << ", count: " << count << "\n";
    std::cout << "Scalar type: " << type_name(Scalar()) << " x " << PacketSize << "\n";
    std::cout << "data in:  {";
    for (Index i = 0; i < PacketSize; i++) {
      if (i > 0) std::cout << ",";
      if (i < begin || i >= begin + count) {
        std::cout << "MASK";
      } else {
        std::cout << data_in[i];
      }
    }
    std::cout << "}\n";
    std::cout << "data out: {";
    for (Index i = 0; i < PacketSize; i++) {
      if (i > 0) std::cout << ",";
      if (i < begin || i >= begin + count) {
        std::cout << "MASK";
      } else {
        std::cout << data_out[i];
      }
    }
    std::cout << "}\n";
    std::cout << "packet:   ";
    std::cout << internal::postream(a) << "\n";
  }
  VERIFY(ok);
}

template <typename Scalar, int PacketSize, bool Run = internal::find_packet_by_size<Scalar, PacketSize>::value>
struct packet_segment_test_impl {
  using Packet = typename internal::find_packet_by_size<Scalar, PacketSize>::type;
  static void test_unaligned() {
    // test loading a packet segment from unaligned memory that includes unallocated memory

    // | X   X   X   X | *   *   *   X | X   X   X   X |
    //    begin -> { X | *   *   *   } <- begin + count

    VectorX<Scalar> data_in(PacketSize), data_out(PacketSize);
    data_in.setRandom();
    data_out.setRandom();

    Scalar* unaligned_data_in = data_in.data() - 1;
    Scalar* unaligned_data_out = data_out.data() - 1;

    Index begin = 1;
    Index count = PacketSize - 1;

    Packet a = internal::ploaduSegment<Packet>(unaligned_data_in, begin, count);
    internal::pstoreuSegment<Scalar, Packet>(unaligned_data_out, a, begin, count);

    verify_data(unaligned_data_in, unaligned_data_out, a, begin, count);

    // test loading the entire packet

    data_in.setRandom();
    data_out.setRandom();

    unaligned_data_in = data_in.data();
    unaligned_data_out = data_out.data();

    begin = 0;
    count = PacketSize;

    Packet b = internal::ploaduSegment<Packet>(unaligned_data_in, begin, count);
    internal::pstoreuSegment<Scalar, Packet>(unaligned_data_out, b, begin, count);

    verify_data(unaligned_data_in, unaligned_data_out, b, begin, count);

    // test loading an empty packet segment in unallocated memory
    count = 0;

    for (begin = 0; begin < PacketSize; begin++) {
      data_in.setRandom();
      data_out = data_in;
      Packet c = internal::ploaduSegment<Packet>(data_in.data(), begin, count);
      internal::pstoreuSegment<Scalar, Packet>(data_out.data(), c, begin, count);
      // verify that ploaduSegment / pstoreuSegment did nothing
      VERIFY_IS_CWISE_EQUAL(data_in, data_out);
    }
  }
  static void test_aligned() {
    // test loading a packet segment from aligned memory that includes unallocated memory

    // | X   X   X   X | *   *   *   X | X   X   X   X |
    //        begin -> { *   *   *   X } <- begin + count

    VectorX<Scalar> data_in(PacketSize - 1), data_out(PacketSize - 1);
    data_in.setRandom();
    data_out.setRandom();

    Scalar* aligned_data_in = data_in.data();
    Scalar* aligned_data_out = data_out.data();

    Index begin = 0;
    Index count = PacketSize - 1;

    Packet b = internal::ploadSegment<Packet>(aligned_data_in, begin, count);
    internal::pstoreSegment<Scalar, Packet>(aligned_data_out, b, begin, count);

    verify_data(aligned_data_in, aligned_data_out, b, begin, count);
  }
  static void test_all_ranges() {
    // every (begin, count) with begin + count <= PacketSize: lanes in the range are copied, the rest are untouched
    VectorX<Scalar> data_in(PacketSize), data_out(PacketSize), data_ref(PacketSize);
    for (Index begin = 0; begin <= PacketSize; begin++) {
      for (Index count = 0; begin + count <= PacketSize; count++) {
        data_in.setRandom();
        data_out.setRandom();
        data_ref = data_out;
        // A scalar loop: a segment assignment would run the primitives under test.
        for (Index i = begin; i < begin + count; ++i) data_ref(i) = data_in(i);
        Packet a = internal::ploaduSegment<Packet>(data_in.data(), begin, count);
        internal::pstoreuSegment<Scalar, Packet>(data_out.data(), a, begin, count);
        VERIFY_IS_CWISE_EQUAL(data_out, data_ref);
      }
    }
  }
  static void test_zero_fill() {
    // ploaduSegment zeroes the lanes outside [begin, begin + count), even when the memory there holds nonzero values;
    // the GEMV kernels sum whole segment packets.
    VectorX<Scalar> data_in = VectorX<Scalar>::Constant(PacketSize, Scalar(1));
    VectorX<Scalar> lanes(PacketSize);
    for (Index begin = 0; begin <= PacketSize; begin++) {
      for (Index count = 0; begin + count <= PacketSize; count++) {
        internal::pstoreu(lanes.data(), internal::ploaduSegment<Packet>(data_in.data(), begin, count));
        for (Index i = 0; i < PacketSize; ++i) {
          const bool in_range = i >= begin && i < begin + count;
          VERIFY_IS_EQUAL(lanes(i), in_range ? Scalar(1) : Scalar(0));
        }
      }
    }
  }
  static void verify_invalid_range(Index begin, Index count) {
    VectorX<Scalar> data = VectorX<Scalar>::Zero(PacketSize);
    const Packet a = internal::pset1<Packet>(Scalar(0));
    VERIFY_RAISES_ASSERT(internal::ploaduSegment<Packet>(data.data(), begin, count));
    VERIFY_RAISES_ASSERT((internal::pstoreuSegment<Scalar, Packet>(data.data(), a, begin, count)));
  }
  static void test_invalid_ranges() {
    // Every specialization asserts 0 <= begin, 0 <= count, begin + count <= PacketSize.
    verify_invalid_range(-1, 1);
    verify_invalid_range(0, -1);
    verify_invalid_range(0, PacketSize + 1);
    verify_invalid_range(1, PacketSize);
    verify_invalid_range(PacketSize + 1, 0);
  }
  static void run() {
    test_unaligned();
    test_aligned();
    test_all_ranges();
    test_zero_fill();
    test_invalid_ranges();
  }
};

template <typename Scalar, int PacketSize>
struct packet_segment_test_impl<Scalar, PacketSize, false> {
  static void run() {}
};

template <typename Scalar, int PacketSize>
struct packet_segment_test_driver {
  static void run() {
    packet_segment_test_impl<Scalar, PacketSize>::run();
    packet_segment_test_driver<Scalar, PacketSize / 2>::run();
  }
};

template <typename Scalar>
struct packet_segment_test_driver<Scalar, 1> {
  static void run() {}
};

template <bool Enable = internal::packet_traits<half>::Vectorizable>
void testReverseEdgeCase() {
  // this reversed cast uses a non-zero offset for ploadSegment
  Index size = 16 * internal::packet_traits<half>::size + 1;
  VectorX<half> v1(size);
  VectorX<float> v2(size), v3(size);
  v1.setRandom();
  v2 = v1.reverse().cast<float>();
  v3 = v1.cast<float>().reverse();
  VERIFY_IS_EQUAL(v2, v3);
}

template <>
void testReverseEdgeCase<false>() {}

template <typename Scalar>
void test_packet_segment() {
  packet_segment_test_driver<Scalar, internal::packet_traits<Scalar>::size>::run();
}

EIGEN_DECLARE_TEST(packet_segment) {
#if defined(EIGEN_VECTORIZE_AVX) && !defined(EIGEN_VECTORIZE_AVX512VL)
  std::complex<float> data[2] = {};
  VERIFY_RAISES_ASSERT(internal::ploaduSegment<internal::Packet2cf>(data, 0, 3));
  VERIFY_RAISES_ASSERT(internal::ploaduSegment<internal::Packet2cf>(data, 0, -1));
  VERIFY_RAISES_ASSERT(internal::ploaduSegment<internal::Packet2cf>(data, -1, 0));
  VERIFY_RAISES_ASSERT(internal::ploaduSegment<internal::Packet2cf>(data, 3, 0));
  VERIFY_RAISES_ASSERT(internal::ploaduSegment<internal::Packet2cf>(data, (std::numeric_limits<Index>::max)(), 1));
#endif
  for (int i = 0; i < g_repeat; i++) {
    test_packet_segment<bool>();
    test_packet_segment<int8_t>();
    test_packet_segment<uint8_t>();
    test_packet_segment<int16_t>();
    test_packet_segment<uint16_t>();
    test_packet_segment<int32_t>();
    test_packet_segment<uint32_t>();
    test_packet_segment<int64_t>();
    test_packet_segment<uint64_t>();
    test_packet_segment<bfloat16>();
    test_packet_segment<half>();
    test_packet_segment<float>();
    test_packet_segment<double>();
    test_packet_segment<std::complex<float>>();
    test_packet_segment<std::complex<double>>();
    testReverseEdgeCase();
  }
}