#include "simd_ops.h"
#include <iostream>
#include <algorithm> // For std::sort, std::min, std::max
#include <cstring>   // For std::memcpy

#if defined(__x86_64__) || defined(_M_X64)
#include <immintrin.h>
#endif

namespace simd
{

// Static variable to control SIMD usage
static bool simd_enabled = true;

// Enable or disable SIMD acceleration
void enable_simd(bool enable)
{
    simd_enabled = enable;
}

// Check if SIMD acceleration is enabled
bool is_simd_enabled()
{
    return simd_enabled;
}

// Runtime CPU feature detection
SimdLevel get_simd_level()
{
    // If SIMD is disabled, always return None
    if (!simd_enabled)
    {
        return SimdLevel::None;
    }

    static SimdLevel detected_level = SimdLevel::None;
    static bool initialized = false;

    if (initialized)
    {
        return detected_level;
    }

#if defined(__x86_64__) || defined(_M_X64)
    // x86_64 CPU feature detection
#if defined(_MSC_VER)
    // MSVC implementation
    int cpu_info[4];
    __cpuid(cpu_info, 1);
    bool has_sse41 = (cpu_info[2] & (1 << 19)) != 0;

    __cpuid(cpu_info, 7);
    bool has_avx = (cpu_info[1] & (1 << 5)) != 0;
    bool has_avx2 = (cpu_info[1] & (1 << 5)) != 0;
#else
    // GCC/Clang implementation
    bool has_sse41 = __builtin_cpu_supports("sse4.1");
    bool has_avx = __builtin_cpu_supports("avx");
    bool has_avx2 = __builtin_cpu_supports("avx2");
#endif

    if (has_avx2)
    {
        detected_level = SimdLevel::AVX2;
    }
    else if (has_avx)
    {
        detected_level = SimdLevel::AVX;
    }
    else if (has_sse41)
    {
        detected_level = SimdLevel::SSE41;
    }
    else
    {
        detected_level = SimdLevel::None;
    }
#elif defined(__ARM_NEON) || defined(__ARM_NEON__)
    // ARM NEON detection
    detected_level = SimdLevel::NEON;
#else
    detected_level = SimdLevel::None;
#endif

    initialized = true;
    return detected_level;
}

// Dispatcher functions that select the appropriate implementation based on runtime detection

// 8-bit operations
void add_u8(uint8_t* a, const uint8_t* b, size_t n)
{
    SimdLevel level = get_simd_level();

    switch (level)
    {
#ifdef HAVE_AVX2
        case SimdLevel::AVX2:
            add_u8_avx2(a, b, n);
            return;
#endif
#ifdef HAVE_AVX
        case SimdLevel::AVX:
            add_u8_avx(a, b, n);
            return;
#endif
#ifdef HAVE_SSE
        case SimdLevel::SSE41:
            add_u8_sse(a, b, n);
            return;
#endif
#ifdef HAVE_NEON
        case SimdLevel::NEON:
            add_u8_neon(a, b, n);
            return;
#endif
        default:
            add_u8_generic(a, b, n);
            return;
    }
}

void subtract_u8(uint8_t* a, const uint8_t* b, size_t n)
{
    SimdLevel level = get_simd_level();

    switch (level)
    {
#ifdef HAVE_AVX2
        case SimdLevel::AVX2:
            subtract_u8_avx2(a, b, n);
            return;
#endif
#ifdef HAVE_AVX
        case SimdLevel::AVX:
            subtract_u8_avx(a, b, n);
            return;
#endif
#ifdef HAVE_SSE
        case SimdLevel::SSE41:
            subtract_u8_sse(a, b, n);
            return;
#endif
#ifdef HAVE_NEON
        case SimdLevel::NEON:
            subtract_u8_neon(a, b, n);
            return;
#endif
        default:
            subtract_u8_generic(a, b, n);
            return;
    }
}

void multiply_u8(uint8_t* a, float factor, size_t n)
{
    SimdLevel level = get_simd_level();

    switch (level)
    {
#ifdef HAVE_AVX2
        case SimdLevel::AVX2:
            multiply_u8_avx2(a, factor, n);
            return;
#endif
#ifdef HAVE_AVX
        case SimdLevel::AVX:
            multiply_u8_avx(a, factor, n);
            return;
#endif
#ifdef HAVE_SSE
        case SimdLevel::SSE41:
            multiply_u8_sse(a, factor, n);
            return;
#endif
#ifdef HAVE_NEON
        case SimdLevel::NEON:
            multiply_u8_neon(a, factor, n);
            return;
#endif
        default:
            multiply_u8_generic(a, factor, n);
            return;
    }
}

// 16-bit operations
void add_u16(uint16_t* a, const uint16_t* b, size_t n)
{
    SimdLevel level = get_simd_level();

    switch (level)
    {
#ifdef HAVE_AVX2
        case SimdLevel::AVX2:
            add_u16_avx2(a, b, n);
            return;
#endif
#ifdef HAVE_AVX
        case SimdLevel::AVX:
            add_u16_avx(a, b, n);
            return;
#endif
#ifdef HAVE_SSE
        case SimdLevel::SSE41:
            add_u16_sse(a, b, n);
            return;
#endif
#ifdef HAVE_NEON
        case SimdLevel::NEON:
            add_u16_neon(a, b, n);
            return;
#endif
        default:
            add_u16_generic(a, b, n);
            return;
    }
}

void subtract_u16(uint16_t* a, const uint16_t* b, size_t n)
{
    SimdLevel level = get_simd_level();

    switch (level)
    {
#ifdef HAVE_AVX2
        case SimdLevel::AVX2:
            subtract_u16_avx2(a, b, n);
            return;
#endif
#ifdef HAVE_AVX
        case SimdLevel::AVX:
            subtract_u16_avx(a, b, n);
            return;
#endif
#ifdef HAVE_SSE
        case SimdLevel::SSE41:
            subtract_u16_sse(a, b, n);
            return;
#endif
#ifdef HAVE_NEON
        case SimdLevel::NEON:
            subtract_u16_neon(a, b, n);
            return;
#endif
        default:
            subtract_u16_generic(a, b, n);
            return;
    }
}

void multiply_u16(uint16_t* a, float factor, size_t n)
{
    SimdLevel level = get_simd_level();

    switch (level)
    {
#ifdef HAVE_AVX2
        case SimdLevel::AVX2:
            multiply_u16_avx2(a, factor, n);
            return;
#endif
#ifdef HAVE_AVX
        case SimdLevel::AVX:
            multiply_u16_avx(a, factor, n);
            return;
#endif
#ifdef HAVE_SSE
        case SimdLevel::SSE41:
            multiply_u16_sse(a, factor, n);
            return;
#endif
#ifdef HAVE_NEON
        case SimdLevel::NEON:
            multiply_u16_neon(a, factor, n);
            return;
#endif
        default:
            multiply_u16_generic(a, factor, n);
            return;
    }
}

// Pixel statistics operations
float mean_u8(const uint8_t* data, size_t n)
{
    SimdLevel level = get_simd_level();

    switch (level)
    {
#ifdef HAVE_AVX2
        case SimdLevel::AVX2:
            return mean_u8_avx2(data, n);
#endif
#ifdef HAVE_AVX
        case SimdLevel::AVX:
            return mean_u8_avx(data, n);
#endif
#ifdef HAVE_SSE
        case SimdLevel::SSE41:
            return mean_u8_sse(data, n);
#endif
#ifdef HAVE_NEON
        case SimdLevel::NEON:
            return mean_u8_neon(data, n);
#endif
        default:
            return mean_u8_generic(data, n);
    }
}

float mean_u16(const uint16_t* data, size_t n)
{
    SimdLevel level = get_simd_level();

    switch (level)
    {
#ifdef HAVE_AVX2
        case SimdLevel::AVX2:
            return mean_u16_avx2(data, n);
#endif
#ifdef HAVE_AVX
        case SimdLevel::AVX:
            return mean_u16_avx(data, n);
#endif
#ifdef HAVE_SSE
        case SimdLevel::SSE41:
            return mean_u16_sse(data, n);
#endif
#ifdef HAVE_NEON
        case SimdLevel::NEON:
            return mean_u16_neon(data, n);
#endif
        default:
            return mean_u16_generic(data, n);
    }
}

// Generic implementations are now in simd_generic.cpp

// Dispatcher functions for the new operations
float variance_u8(const uint8_t* data, size_t n)
{
    SimdLevel level = get_simd_level();

    switch (level)
    {
#ifdef HAVE_AVX2
        case SimdLevel::AVX2:
            return variance_u8_avx2(data, n);
#endif
#ifdef HAVE_AVX
        case SimdLevel::AVX:
            return variance_u8_avx(data, n);
#endif
#ifdef HAVE_SSE
        case SimdLevel::SSE41:
            return variance_u8_sse(data, n);
#endif
#ifdef HAVE_NEON
        case SimdLevel::NEON:
            return variance_u8_neon(data, n);
#endif
        default:
            return variance_u8_generic(data, n);
    }
}

float variance_u16(const uint16_t* data, size_t n)
{
    SimdLevel level = get_simd_level();

    switch (level)
    {
#ifdef HAVE_AVX2
        case SimdLevel::AVX2:
            return variance_u16_avx2(data, n);
#endif
#ifdef HAVE_AVX
        case SimdLevel::AVX:
            return variance_u16_avx(data, n);
#endif
#ifdef HAVE_SSE
        case SimdLevel::SSE41:
            return variance_u16_sse(data, n);
#endif
#ifdef HAVE_NEON
        case SimdLevel::NEON:
            return variance_u16_neon(data, n);
#endif
        default:
            return variance_u16_generic(data, n);
    }
}

uint8_t median_u8(uint8_t* data, size_t n)
{
    SimdLevel level = get_simd_level();

    switch (level)
    {
#ifdef HAVE_SSE
        case SimdLevel::SSE41:
            return median_u8_sse(data, n);
#endif
#ifdef HAVE_NEON
        case SimdLevel::NEON:
            return median_u8_neon(data, n);
#endif
        default:
            return median_u8_generic(data, n);
    }
}

uint16_t median_u16(uint16_t* data, size_t n)
{
    SimdLevel level = get_simd_level();

    switch (level)
    {
#ifdef HAVE_SSE
        case SimdLevel::SSE41:
            return median_u16_sse(data, n);
#endif
#ifdef HAVE_NEON
        case SimdLevel::NEON:
            return median_u16_neon(data, n);
#endif
        default:
            return median_u16_generic(data, n);
    }
}

void min_max_u8(const uint8_t* data, size_t n, uint8_t &min_val, uint8_t &max_val)
{
    SimdLevel level = get_simd_level();

    switch (level)
    {
#ifdef HAVE_AVX2
        case SimdLevel::AVX2:
            min_max_u8_avx2(data, n, min_val, max_val);
            return;
#endif
#ifdef HAVE_AVX
        case SimdLevel::AVX:
            min_max_u8_avx(data, n, min_val, max_val);
            return;
#endif
#ifdef HAVE_SSE
        case SimdLevel::SSE41:
            min_max_u8_sse(data, n, min_val, max_val);
            return;
#endif
#ifdef HAVE_NEON
        case SimdLevel::NEON:
            min_max_u8_neon(data, n, min_val, max_val);
            return;
#endif
        default:
            min_max_u8_generic(data, n, min_val, max_val);
            return;
    }
}

void min_max_u16(const uint16_t* data, size_t n, uint16_t &min_val, uint16_t &max_val)
{
    SimdLevel level = get_simd_level();

    switch (level)
    {
#ifdef HAVE_AVX2
        case SimdLevel::AVX2:
            min_max_u16_avx2(data, n, min_val, max_val);
            return;
#endif
#ifdef HAVE_AVX
        case SimdLevel::AVX:
            min_max_u16_avx(data, n, min_val, max_val);
            return;
#endif
#ifdef HAVE_SSE
        case SimdLevel::SSE41:
            min_max_u16_sse(data, n, min_val, max_val);
            return;
#endif
#ifdef HAVE_NEON
        case SimdLevel::NEON:
            min_max_u16_neon(data, n, min_val, max_val);
            return;
#endif
        default:
            min_max_u16_generic(data, n, min_val, max_val);
            return;
    }
}

void histogram_u8(const uint8_t* data, size_t n, uint32_t* histogram, size_t bins)
{
    SimdLevel level = get_simd_level();

    switch (level)
    {
#ifdef HAVE_AVX2
        case SimdLevel::AVX2:
            histogram_u8_avx2(data, n, histogram, bins);
            return;
#endif
#ifdef HAVE_AVX
        case SimdLevel::AVX:
            histogram_u8_avx(data, n, histogram, bins);
            return;
#endif
#ifdef HAVE_SSE
        case SimdLevel::SSE41:
            histogram_u8_sse(data, n, histogram, bins);
            return;
#endif
#ifdef HAVE_NEON
        case SimdLevel::NEON:
            histogram_u8_neon(data, n, histogram, bins);
            return;
#endif
        default:
            histogram_u8_generic(data, n, histogram, bins);
            return;
    }
}

void histogram_u16(const uint16_t* data, size_t n, uint32_t* histogram, size_t bins)
{
    // For 16-bit histograms, we'll use the generic implementation for now
    // as SIMD implementations would be more complex due to the large number of bins
    histogram_u16_generic(data, n, histogram, bins);
}

void threshold_u8(uint8_t* data, size_t n, uint8_t threshold, uint8_t low_value, uint8_t high_value)
{
    SimdLevel level = get_simd_level();

    switch (level)
    {
#ifdef HAVE_AVX2
        case SimdLevel::AVX2:
            threshold_u8_avx2(data, n, threshold, low_value, high_value);
            return;
#endif
#ifdef HAVE_AVX
        case SimdLevel::AVX:
            threshold_u8_avx(data, n, threshold, low_value, high_value);
            return;
#endif
#ifdef HAVE_SSE
        case SimdLevel::SSE41:
            threshold_u8_sse(data, n, threshold, low_value, high_value);
            return;
#endif
#ifdef HAVE_NEON
        case SimdLevel::NEON:
            threshold_u8_neon(data, n, threshold, low_value, high_value);
            return;
#endif
        default:
            threshold_u8_generic(data, n, threshold, low_value, high_value);
            return;
    }
}

void threshold_u16(uint16_t* data, size_t n, uint16_t threshold, uint16_t low_value, uint16_t high_value)
{
    SimdLevel level = get_simd_level();

    switch (level)
    {
#ifdef HAVE_AVX2
        case SimdLevel::AVX2:
            threshold_u16_avx2(data, n, threshold, low_value, high_value);
            return;
#endif
#ifdef HAVE_AVX
        case SimdLevel::AVX:
            threshold_u16_avx(data, n, threshold, low_value, high_value);
            return;
#endif
#ifdef HAVE_SSE
        case SimdLevel::SSE41:
            threshold_u16_sse(data, n, threshold, low_value, high_value);
            return;
#endif
#ifdef HAVE_NEON
        case SimdLevel::NEON:
            threshold_u16_neon(data, n, threshold, low_value, high_value);
            return;
#endif
        default:
            threshold_u16_generic(data, n, threshold, low_value, high_value);
            return;
    }
}

void gaussian_blur_u8(uint8_t* dst, const uint8_t* src, int width, int height, int channels, float sigma)
{
    SimdLevel level = get_simd_level();

    switch (level)
    {
#ifdef HAVE_AVX2
        case SimdLevel::AVX2:
            gaussian_blur_u8_avx2(dst, src, width, height, channels, sigma);
            return;
#endif
#ifdef HAVE_AVX
        case SimdLevel::AVX:
            gaussian_blur_u8_avx(dst, src, width, height, channels, sigma);
            return;
#endif
#ifdef HAVE_SSE
        case SimdLevel::SSE41:
            gaussian_blur_u8_sse(dst, src, width, height, channels, sigma);
            return;
#endif
#ifdef HAVE_NEON
        case SimdLevel::NEON:
            gaussian_blur_u8_neon(dst, src, width, height, channels, sigma);
            return;
#endif
        default:
            gaussian_blur_u8_generic(dst, src, width, height, channels, sigma);
            return;
    }
}

void median_filter_u8(uint8_t* dst, const uint8_t* src, int width, int height, int channels, int kernel_size)
{
    SimdLevel level = get_simd_level();

    switch (level)
    {
#ifdef HAVE_AVX2
        case SimdLevel::AVX2:
            median_filter_u8_avx2(dst, src, width, height, channels, kernel_size);
            return;
#endif
#ifdef HAVE_SSE
        case SimdLevel::SSE41:
            median_filter_u8_sse(dst, src, width, height, channels, kernel_size);
            return;
#endif
#ifdef HAVE_NEON
        case SimdLevel::NEON:
            median_filter_u8_neon(dst, src, width, height, channels, kernel_size);
            return;
#endif
        default:
            median_filter_u8_generic(dst, src, width, height, channels, kernel_size);
            return;
    }
}


} // namespace simd
