#include "../buffer/imagebuffer.h"
#include <limits>
#include <algorithm>

#ifdef HAVE_SSE
#include <immintrin.h>

// SIMD optimized implementations for x86_64 architecture

namespace simd
{

// Add two 8-bit buffers with SSE
void add_u8_sse(uint8_t* a, const uint8_t* b, size_t n)
{
    size_t i = 0;

    // Process 16 elements at a time using SSE
    if (n >= 16)
    {
        for (; i <= n - 16; i += 16)
        {
            __m128i va = _mm_loadu_si128(reinterpret_cast<const __m128i*>(a + i));
            __m128i vb = _mm_loadu_si128(reinterpret_cast<const __m128i*>(b + i));
            __m128i result = _mm_adds_epu8(va, vb); // Saturated add
            _mm_storeu_si128(reinterpret_cast<__m128i*>(a + i), result);
        }
    }

    // Process remaining elements
    for (; i < n; ++i)
    {
        int sum = static_cast<int>(a[i]) + static_cast<int>(b[i]);
        a[i] = static_cast<uint8_t>(std::min(255, sum));
    }
}

// Subtract two 8-bit buffers with SSE
void subtract_u8_sse(uint8_t* a, const uint8_t* b, size_t n)
{
    size_t i = 0;

    // Process 16 elements at a time using SSE
    if (n >= 16)
    {
        for (; i <= n - 16; i += 16)
        {
            __m128i va = _mm_loadu_si128(reinterpret_cast<const __m128i*>(a + i));
            __m128i vb = _mm_loadu_si128(reinterpret_cast<const __m128i*>(b + i));
            __m128i result = _mm_subs_epu8(va, vb); // Saturated subtract
            _mm_storeu_si128(reinterpret_cast<__m128i*>(a + i), result);
        }
    }

    // Process remaining elements
    for (; i < n; ++i)
    {
        int diff = static_cast<int>(a[i]) - static_cast<int>(b[i]);
        a[i] = static_cast<uint8_t>(std::max(0, diff));
    }
}

// Multiply 8-bit buffer by a scalar with SSE
void multiply_u8_sse(uint8_t* a, float factor, size_t n)
{
    size_t i = 0;

    // Create a vector with the factor replicated 4 times
    __m128 vfactor = _mm_set1_ps(factor);

    // Process 16 elements at a time using SSE
    if (n >= 16)
    {
        for (; i <= n - 16; i += 16)
        {
            // Load 16 bytes
            __m128i v_int = _mm_loadu_si128(reinterpret_cast<const __m128i*>(a + i));

            // Process first 8 bytes
            __m128i v_int_lo = _mm_unpacklo_epi8(v_int, _mm_setzero_si128());
            __m128i v_int_lo_lo = _mm_unpacklo_epi16(v_int_lo, _mm_setzero_si128());
            __m128i v_int_lo_hi = _mm_unpackhi_epi16(v_int_lo, _mm_setzero_si128());

            // Convert to float
            __m128 v_float_lo_lo = _mm_cvtepi32_ps(v_int_lo_lo);
            __m128 v_float_lo_hi = _mm_cvtepi32_ps(v_int_lo_hi);

            // Multiply by factor
            v_float_lo_lo = _mm_mul_ps(v_float_lo_lo, vfactor);
            v_float_lo_hi = _mm_mul_ps(v_float_lo_hi, vfactor);

            // Convert back to int
            v_int_lo_lo = _mm_cvtps_epi32(v_float_lo_lo);
            v_int_lo_hi = _mm_cvtps_epi32(v_float_lo_hi);

            // Process second 8 bytes
            __m128i v_int_hi = _mm_unpackhi_epi8(v_int, _mm_setzero_si128());
            __m128i v_int_hi_lo = _mm_unpacklo_epi16(v_int_hi, _mm_setzero_si128());
            __m128i v_int_hi_hi = _mm_unpackhi_epi16(v_int_hi, _mm_setzero_si128());

            // Convert to float
            __m128 v_float_hi_lo = _mm_cvtepi32_ps(v_int_hi_lo);
            __m128 v_float_hi_hi = _mm_cvtepi32_ps(v_int_hi_hi);

            // Multiply by factor
            v_float_hi_lo = _mm_mul_ps(v_float_hi_lo, vfactor);
            v_float_hi_hi = _mm_mul_ps(v_float_hi_hi, vfactor);

            // Convert back to int
            v_int_hi_lo = _mm_cvtps_epi32(v_float_hi_lo);
            v_int_hi_hi = _mm_cvtps_epi32(v_float_hi_hi);

            // Pack back to 16-bit
            v_int_lo = _mm_packs_epi32(v_int_lo_lo, v_int_lo_hi);
            v_int_hi = _mm_packs_epi32(v_int_hi_lo, v_int_hi_hi);

            // Pack back to 8-bit
            v_int = _mm_packus_epi16(v_int_lo, v_int_hi);

            // Store the result
            _mm_storeu_si128(reinterpret_cast<__m128i*>(a + i), v_int);
        }
    }

    // Process remaining elements
    for (; i < n; ++i)
    {
        float result = static_cast<float>(a[i]) * factor;
        a[i] = static_cast<uint8_t>(std::min(255.0f, result));
    }
}

// Add two 16-bit buffers with SSE
void add_u16_sse(uint16_t* a, const uint16_t* b, size_t n)
{
    size_t i = 0;

    // Process 8 elements at a time using SSE
    if (n >= 8)
    {
        for (; i <= n - 8; i += 8)
        {
            __m128i va = _mm_loadu_si128(reinterpret_cast<const __m128i*>(a + i));
            __m128i vb = _mm_loadu_si128(reinterpret_cast<const __m128i*>(b + i));

            // There's no saturated add for 16-bit unsigned integers in SSE2,
            // so we need to implement it manually
            __m128i sum = _mm_add_epi16(va, vb);

            // Check for overflow: if sum < a or sum < b, then overflow occurred
            __m128i overflow_mask = _mm_or_si128(_mm_cmplt_epi16(sum, va), _mm_cmplt_epi16(sum, vb));

            // Set overflowed elements to max value (0xFFFF)
            __m128i max_val = _mm_set1_epi16(0xFFFF);
            __m128i result = _mm_or_si128(_mm_andnot_si128(overflow_mask, sum),
                                          _mm_and_si128(overflow_mask, max_val));

            _mm_storeu_si128(reinterpret_cast<__m128i*>(a + i), result);
        }
    }

    // Process remaining elements
    for (; i < n; ++i)
    {
        int sum = static_cast<int>(a[i]) + static_cast<int>(b[i]);
        a[i] = static_cast<uint16_t>(std::min(65535, sum));
    }
}

// Subtract two 16-bit buffers with SSE
void subtract_u16_sse(uint16_t* a, const uint16_t* b, size_t n)
{
    size_t i = 0;

    // Process 8 elements at a time using SSE
    if (n >= 8)
    {
        for (; i <= n - 8; i += 8)
        {
            __m128i va = _mm_loadu_si128(reinterpret_cast<const __m128i*>(a + i));
            __m128i vb = _mm_loadu_si128(reinterpret_cast<const __m128i*>(b + i));

            // There's no saturated subtract for 16-bit unsigned integers in SSE2,
            // so we need to implement it manually
            __m128i mask = _mm_cmplt_epi16(va, vb);
            __m128i result = _mm_or_si128(
                                 _mm_and_si128(mask, _mm_setzero_si128()),
                                 _mm_andnot_si128(mask, _mm_sub_epi16(va, vb))
                             );

            _mm_storeu_si128(reinterpret_cast<__m128i*>(a + i), result);
        }
    }

    // Process remaining elements
    for (; i < n; ++i)
    {
        int diff = static_cast<int>(a[i]) - static_cast<int>(b[i]);
        a[i] = static_cast<uint16_t>(std::max(0, diff));
    }
}

// Multiply 16-bit buffer by a scalar with SSE
void multiply_u16_sse(uint16_t* a, float factor, size_t n)
{
    size_t i = 0;

    // Create a vector with the factor replicated 4 times
    __m128 vfactor = _mm_set1_ps(factor);

    // Process 8 elements at a time using SSE
    if (n >= 8)
    {
        for (; i <= n - 8; i += 8)
        {
            // Load 8 16-bit integers
            __m128i v_int = _mm_loadu_si128(reinterpret_cast<const __m128i*>(a + i));

            // Process first 4 16-bit integers
            __m128i v_int_lo = _mm_unpacklo_epi16(v_int, _mm_setzero_si128());

            // Convert to float
            __m128 v_float_lo = _mm_cvtepi32_ps(v_int_lo);

            // Multiply by factor
            v_float_lo = _mm_mul_ps(v_float_lo, vfactor);

            // Convert back to int
            v_int_lo = _mm_cvtps_epi32(v_float_lo);

            // Process second 4 16-bit integers
            __m128i v_int_hi = _mm_unpackhi_epi16(v_int, _mm_setzero_si128());

            // Convert to float
            __m128 v_float_hi = _mm_cvtepi32_ps(v_int_hi);

            // Multiply by factor
            v_float_hi = _mm_mul_ps(v_float_hi, vfactor);

            // Convert back to int
            v_int_hi = _mm_cvtps_epi32(v_float_hi);

            // Pack back to 16-bit with saturation
            v_int = _mm_packus_epi32(v_int_lo, v_int_hi);

            // Store the result
            _mm_storeu_si128(reinterpret_cast<__m128i*>(a + i), v_int);
        }
    }

    // Process remaining elements
    for (; i < n; ++i)
    {
        float result = static_cast<float>(a[i]) * factor;
        a[i] = static_cast<uint16_t>(std::min(65535.0f, result));
    }
}

} // namespace simd

#endif // HAVE_SSE
