#include "simd_ops.h"
#include <algorithm>
#include <cstring>
#include <numeric>
#include <cmath>

#ifdef HAVE_OPENMP
#include <omp.h>
#endif

namespace simd
{

// Basic arithmetic operations for 8-bit data
void add_u8_generic(uint8_t* a, const uint8_t* b, size_t n)
{
#ifdef HAVE_OPENMP
    #pragma omp parallel for schedule(static) if(n > 10000)
#endif
    for (size_t i = 0; 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));
    }
}

void subtract_u8_generic(uint8_t* a, const uint8_t* b, size_t n)
{
#ifdef HAVE_OPENMP
    #pragma omp parallel for schedule(static) if(n > 10000)
#endif
    for (size_t i = 0; 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));
    }
}

void multiply_u8_generic(uint8_t* a, float factor, size_t n)
{
#ifdef HAVE_OPENMP
    #pragma omp parallel for schedule(static) if(n > 10000)
#endif
    for (size_t i = 0; i < n; ++i)
    {
        float result = static_cast<float>(a[i]) * factor;
        a[i] = static_cast<uint8_t>(std::min(255.0f, result));
    }
}

// Basic arithmetic operations for 16-bit data
void add_u16_generic(uint16_t* a, const uint16_t* b, size_t n)
{
#ifdef HAVE_OPENMP
    #pragma omp parallel for schedule(static) if(n > 10000)
#endif
    for (size_t i = 0; 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));
    }
}

void subtract_u16_generic(uint16_t* a, const uint16_t* b, size_t n)
{
#ifdef HAVE_OPENMP
    #pragma omp parallel for schedule(static) if(n > 10000)
#endif
    for (size_t i = 0; 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));
    }
}

void multiply_u16_generic(uint16_t* a, float factor, size_t n)
{
#ifdef HAVE_OPENMP
    #pragma omp parallel for schedule(static) if(n > 10000)
#endif
    for (size_t i = 0; i < n; ++i)
    {
        float result = static_cast<float>(a[i]) * factor;
        a[i] = static_cast<uint16_t>(std::min(65535.0f, result));
    }
}

// Statistical functions for 8-bit data
float mean_u8_generic(const uint8_t* data, size_t n)
{
    if (n == 0) return 0.0f;

    uint64_t sum = 0;

#ifdef HAVE_OPENMP
    #pragma omp parallel reduction(+:sum) if(n > 10000)
    {
        #pragma omp for schedule(static)
        for (size_t i = 0; i < n; ++i)
        {
            sum += data[i];
        }
    }
#else
    for (size_t i = 0; i < n; ++i)
    {
        sum += data[i];
    }
#endif

    return static_cast<float>(sum) / n;
}

float variance_u8_generic(const uint8_t* data, size_t n)
{
    if (n <= 1) return 0.0f;

    float mean_val = mean_u8_generic(data, n);
    float sum_squared_diff = 0.0f;

#ifdef HAVE_OPENMP
    #pragma omp parallel reduction(+:sum_squared_diff) if(n > 10000)
    {
        #pragma omp for schedule(static)
        for (size_t i = 0; i < n; ++i)
        {
            float diff = data[i] - mean_val;
            sum_squared_diff += diff * diff;
        }
    }
#else
    for (size_t i = 0; i < n; ++i)
    {
        float diff = data[i] - mean_val;
        sum_squared_diff += diff * diff;
    }
#endif

    return sum_squared_diff / (n - 1);
}

uint8_t median_u8_generic(uint8_t* data, size_t n)
{
    if (n == 0) return 0;
    if (n == 1) return data[0];

    // Sort the data (in-place)
    std::sort(data, data + n);

    // Return the median
    if (n % 2 == 0)
    {
        // Even number of elements, average the middle two
        return static_cast<uint8_t>((static_cast<uint16_t>(data[n / 2 - 1]) + data[n / 2]) / 2);
    }
    else
    {
        // Odd number of elements, return the middle one
        return data[n / 2];
    }
}

// Statistical functions for 16-bit data
float mean_u16_generic(const uint16_t* data, size_t n)
{
    if (n == 0) return 0.0f;

    uint64_t sum = 0;

#ifdef HAVE_OPENMP
    #pragma omp parallel reduction(+:sum) if(n > 10000)
    {
        #pragma omp for schedule(static)
        for (size_t i = 0; i < n; ++i)
        {
            sum += data[i];
        }
    }
#else
    for (size_t i = 0; i < n; ++i)
    {
        sum += data[i];
    }
#endif

    return static_cast<float>(sum) / n;
}

float variance_u16_generic(const uint16_t* data, size_t n)
{
    if (n <= 1) return 0.0f;

    float mean_val = mean_u16_generic(data, n);
    float sum_squared_diff = 0.0f;

#ifdef HAVE_OPENMP
    #pragma omp parallel reduction(+:sum_squared_diff) if(n > 10000)
    {
        #pragma omp for schedule(static)
        for (size_t i = 0; i < n; ++i)
        {
            float diff = data[i] - mean_val;
            sum_squared_diff += diff * diff;
        }
    }
#else
    for (size_t i = 0; i < n; ++i)
    {
        float diff = data[i] - mean_val;
        sum_squared_diff += diff * diff;
    }
#endif

    return sum_squared_diff / (n - 1);
}

uint16_t median_u16_generic(uint16_t* data, size_t n)
{
    if (n == 0) return 0;
    if (n == 1) return data[0];

    // Sort the data (in-place)
    std::sort(data, data + n);

    // Return the median
    if (n % 2 == 0)
    {
        // Even number of elements, average the middle two
        return static_cast<uint16_t>((static_cast<uint32_t>(data[n / 2 - 1]) + data[n / 2]) / 2);
    }
    else
    {
        // Odd number of elements, return the middle one
        return data[n / 2];
    }
}

// Min/max functions
void min_max_u8_generic(const uint8_t* data, size_t n, uint8_t &min_val, uint8_t &max_val)
{
    if (n == 0)
    {
        min_val = 0;
        max_val = 0;
        return;
    }

    min_val = 255;
    max_val = 0;

#ifdef HAVE_OPENMP
    if (n > 10000) // Only use parallel version for large data
    {
        uint8_t thread_min_vals[128]; // Support up to 128 threads
        uint8_t thread_max_vals[128];

        for (int i = 0; i < 128; i++)
        {
            thread_min_vals[i] = 255;
            thread_max_vals[i] = 0;
        }

        #pragma omp parallel
        {
            int thread_id = omp_get_thread_num();
            uint8_t local_min = 255;
            uint8_t local_max = 0;

            #pragma omp for schedule(static)
            for (size_t i = 0; i < n; ++i)
            {
                local_min = std::min(local_min, data[i]);
                local_max = std::max(local_max, data[i]);
            }

            thread_min_vals[thread_id] = local_min;
            thread_max_vals[thread_id] = local_max;
        }

        // Combine results from all threads
        int num_threads = omp_get_max_threads();
        for (int t = 0; t < num_threads; ++t)
        {
            min_val = std::min(min_val, thread_min_vals[t]);
            max_val = std::max(max_val, thread_max_vals[t]);
        }

        return;
    }
#endif

    // Sequential version for small data or when OpenMP is not available
    for (size_t i = 0; i < n; ++i)
    {
        min_val = std::min(min_val, data[i]);
        max_val = std::max(max_val, data[i]);
    }
}

void min_max_u16_generic(const uint16_t* data, size_t n, uint16_t &min_val, uint16_t &max_val)
{
    if (n == 0)
    {
        min_val = 0;
        max_val = 0;
        return;
    }

    min_val = 65535;
    max_val = 0;

#ifdef HAVE_OPENMP
    if (n > 10000) // Only use parallel version for large data
    {
        uint16_t thread_min_vals[128]; // Support up to 128 threads
        uint16_t thread_max_vals[128];

        for (int i = 0; i < 128; i++)
        {
            thread_min_vals[i] = 65535;
            thread_max_vals[i] = 0;
        }

        #pragma omp parallel
        {
            int thread_id = omp_get_thread_num();
            uint16_t local_min = 65535;
            uint16_t local_max = 0;

            #pragma omp for schedule(static)
            for (size_t i = 0; i < n; ++i)
            {
                local_min = std::min(local_min, data[i]);
                local_max = std::max(local_max, data[i]);
            }

            thread_min_vals[thread_id] = local_min;
            thread_max_vals[thread_id] = local_max;
        }

        // Combine results from all threads
        int num_threads = omp_get_max_threads();
        for (int t = 0; t < num_threads; ++t)
        {
            min_val = std::min(min_val, thread_min_vals[t]);
            max_val = std::max(max_val, thread_max_vals[t]);
        }

        return;
    }
#endif

    // Sequential version for small data or when OpenMP is not available
    for (size_t i = 0; i < n; ++i)
    {
        min_val = std::min(min_val, data[i]);
        max_val = std::max(max_val, data[i]);
    }
}

// Histogram functions
void histogram_u8_generic(const uint8_t* data, size_t n, uint32_t* histogram, size_t bins)
{
    // Initialize histogram to zeros
    for (size_t i = 0; i < bins; ++i)
    {
        histogram[i] = 0;
    }

#ifdef HAVE_OPENMP
    // Use thread-local histograms to avoid contention
    int num_threads = 1;
    #pragma omp parallel
    {
        #pragma omp single
        num_threads = omp_get_num_threads();
    }

    if (num_threads > 1 && n > 10000) // Only use parallel version for large data
    {
        // Create thread-local histograms
        std::vector<std::vector<uint32_t>> thread_histograms(num_threads, std::vector<uint32_t>(bins, 0));

        // Parallel counting
        #pragma omp parallel
        {
            int thread_id = omp_get_thread_num();
            auto &local_hist = thread_histograms[thread_id];

            #pragma omp for schedule(static)
            for (size_t i = 0; i < n; ++i)
            {
                local_hist[data[i]]++;
            }
        }

        // Combine thread-local histograms
        for (int t = 0; t < num_threads; ++t)
        {
            for (size_t i = 0; i < bins; ++i)
            {
                histogram[i] += thread_histograms[t][i];
            }
        }
        return;
    }
#endif

    // Sequential version for small data or when OpenMP is not available
    for (size_t i = 0; i < n; ++i)
    {
        histogram[data[i]]++;
    }
}

void histogram_u16_generic(const uint16_t* data, size_t n, uint32_t* histogram, size_t bins)
{
    // Initialize histogram to zeros
    for (size_t i = 0; i < bins; ++i)
    {
        histogram[i] = 0;
    }

#ifdef HAVE_OPENMP
    // Use thread-local histograms to avoid contention
    int num_threads = 1;
    #pragma omp parallel
    {
        #pragma omp single
        num_threads = omp_get_num_threads();
    }

    if (num_threads > 1 && n > 10000) // Only use parallel version for large data
    {
        // Create thread-local histograms
        std::vector<std::vector<uint32_t>> thread_histograms(num_threads, std::vector<uint32_t>(bins, 0));

        // For 16-bit data, we might need to scale the values to fit in the histogram
        if (bins < 65536)
        {
            float scale = static_cast<float>(bins) / 65536.0f;

            // Parallel counting
            #pragma omp parallel
            {
                int thread_id = omp_get_thread_num();
                auto &local_hist = thread_histograms[thread_id];

                #pragma omp for schedule(static)
                for (size_t i = 0; i < n; ++i)
                {
                    size_t bin = static_cast<size_t>(data[i] * scale);
                    if (bin >= bins) bin = bins - 1;
                    local_hist[bin]++;
                }
            }
        }
        else
        {
            // Full 16-bit histogram
            #pragma omp parallel
            {
                int thread_id = omp_get_thread_num();
                auto &local_hist = thread_histograms[thread_id];

                #pragma omp for schedule(static)
                for (size_t i = 0; i < n; ++i)
                {
                    local_hist[data[i]]++;
                }
            }
        }

        // Combine thread-local histograms
        for (int t = 0; t < num_threads; ++t)
        {
            for (size_t i = 0; i < bins; ++i)
            {
                histogram[i] += thread_histograms[t][i];
            }
        }
        return;
    }
#endif

    // Sequential version for small data or when OpenMP is not available
    // For 16-bit data, we might need to scale the values to fit in the histogram
    if (bins < 65536)
    {
        float scale = static_cast<float>(bins) / 65536.0f;

        for (size_t i = 0; i < n; ++i)
        {
            size_t bin = static_cast<size_t>(data[i] * scale);
            if (bin >= bins) bin = bins - 1;
            histogram[bin]++;
        }
    }
    else
    {
        // Full 16-bit histogram
        for (size_t i = 0; i < n; ++i)
        {
            histogram[data[i]]++;
        }
    }
}

// Threshold functions
void threshold_u8_generic(uint8_t* data, size_t n, uint8_t threshold, uint8_t low_value, uint8_t high_value)
{
#ifdef HAVE_OPENMP
    #pragma omp parallel for schedule(static) if(n > 10000)
#endif
    for (size_t i = 0; i < n; ++i)
    {
        data[i] = (data[i] > threshold) ? high_value : low_value;
    }
}

void threshold_u16_generic(uint16_t* data, size_t n, uint16_t threshold, uint16_t low_value, uint16_t high_value)
{
#ifdef HAVE_OPENMP
    #pragma omp parallel for schedule(static) if(n > 10000)
#endif
    for (size_t i = 0; i < n; ++i)
    {
        data[i] = (data[i] > threshold) ? high_value : low_value;
    }
}

// Image processing functions
void gaussian_blur_u8_generic(uint8_t* dst, const uint8_t* src, int width, int height, int channels, float sigma)
{
    // This is a more complete implementation with OpenMP support
    // First, compute the Gaussian kernel
    int kernel_radius = static_cast<int>(sigma * 3.0f);
    int kernel_size = 2 * kernel_radius + 1;

    std::vector<float> kernel(kernel_size);
    float kernel_sum = 0.0f;

    for (int i = 0; i < kernel_size; ++i)
    {
        float x = static_cast<float>(i - kernel_radius);
        kernel[i] = std::exp(-(x * x) / (2.0f * sigma * sigma));
        kernel_sum += kernel[i];
    }

    // Normalize the kernel
    for (int i = 0; i < kernel_size; ++i)
    {
        kernel[i] /= kernel_sum;
    }

    // Allocate temporary buffer for horizontal pass
    std::vector<uint8_t> temp(width * height * channels);

    // Apply horizontal pass with OpenMP
#ifdef HAVE_OPENMP
    #pragma omp parallel for collapse(2) schedule(dynamic, 16) if(width * height > 10000)
#endif
    for (int y = 0; y < height; ++y)
    {
        for (int x = 0; x < width; ++x)
        {
            for (int c = 0; c < channels; ++c)
            {
                float sum = 0.0f;
                float weight_sum = 0.0f;

                for (int kx = -kernel_radius; kx <= kernel_radius; ++kx)
                {
                    int nx = x + kx;
                    // Handle boundary conditions (clamp to edge)
                    nx = std::max(0, std::min(width - 1, nx));

                    float weight = kernel[kx + kernel_radius];
                    sum += src[(y * width + nx) * channels + c] * weight;
                    weight_sum += weight;
                }

                temp[(y * width + x) * channels + c] = static_cast<uint8_t>(sum / weight_sum);
            }
        }
    }

    // Apply vertical pass with OpenMP
#ifdef HAVE_OPENMP
    #pragma omp parallel for collapse(2) schedule(dynamic, 16) if(width * height > 10000)
#endif
    for (int y = 0; y < height; ++y)
    {
        for (int x = 0; x < width; ++x)
        {
            for (int c = 0; c < channels; ++c)
            {
                float sum = 0.0f;
                float weight_sum = 0.0f;

                for (int ky = -kernel_radius; ky <= kernel_radius; ++ky)
                {
                    int ny = y + ky;
                    // Handle boundary conditions (clamp to edge)
                    ny = std::max(0, std::min(height - 1, ny));

                    float weight = kernel[ky + kernel_radius];
                    sum += temp[(ny * width + x) * channels + c] * weight;
                    weight_sum += weight;
                }

                dst[(y * width + x) * channels + c] = static_cast<uint8_t>(sum / weight_sum);
            }
        }
    }
}

void median_filter_u8_generic(uint8_t* dst, const uint8_t* src, int width, int height, int channels, int kernel_size)
{
    // This is a more complete implementation with OpenMP support
    if (kernel_size <= 0)
    {
        // No filtering, just copy the data
        std::memcpy(dst, src, static_cast<size_t>(width * height * channels));
        return;
    }

    int kernel_radius = kernel_size / 2;
    std::vector<uint8_t> values(kernel_size * kernel_size);

    // Apply median filter with OpenMP
#ifdef HAVE_OPENMP
    #pragma omp parallel for collapse(2) schedule(dynamic, 16) if(width * height > 10000)
#endif
    for (int y = 0; y < height; ++y)
    {
        for (int x = 0; x < width; ++x)
        {
            for (int c = 0; c < channels; ++c)
            {
                int count = 0;

                // Gather values in the kernel
                for (int ky = -kernel_radius; ky <= kernel_radius; ++ky)
                {
                    for (int kx = -kernel_radius; kx <= kernel_radius; ++kx)
                    {
                        int nx = std::max(0, std::min(width - 1, x + kx));
                        int ny = std::max(0, std::min(height - 1, y + ky));

                        values[count++] = src[(ny * width + nx) * channels + c];
                    }
                }

                // Sort values and take the median
                std::sort(values.begin(), values.begin() + count);
                dst[(y * width + x) * channels + c] = values[count / 2];
            }
        }
    }
}

// SSE implementations for x86_64 architecture
#if defined(__x86_64__) || defined(_M_X64)
#ifdef HAVE_SSE

#include <immintrin.h>

// SSE implementations of the functions
float mean_u8_sse(const uint8_t* data, size_t n)
{
    return mean_u8_generic(data, n);
}

float mean_u16_sse(const uint16_t* data, size_t n)
{
    return mean_u16_generic(data, n);
}

float variance_u8_sse(const uint8_t* data, size_t n)
{
    return variance_u8_generic(data, n);
}

float variance_u16_sse(const uint16_t* data, size_t n)
{
    return variance_u16_generic(data, n);
}

uint8_t median_u8_sse(uint8_t* data, size_t n)
{
    return median_u8_generic(data, n);
}

uint16_t median_u16_sse(uint16_t* data, size_t n)
{
    return median_u16_generic(data, n);
}

void min_max_u8_sse(const uint8_t* data, size_t n, uint8_t &min_val, uint8_t &max_val)
{
    min_max_u8_generic(data, n, min_val, max_val);
}

void min_max_u16_sse(const uint16_t* data, size_t n, uint16_t &min_val, uint16_t &max_val)
{
    min_max_u16_generic(data, n, min_val, max_val);
}

void histogram_u8_sse(const uint8_t* data, size_t n, uint32_t* histogram, size_t bins)
{
    histogram_u8_generic(data, n, histogram, bins);
}

void threshold_u8_sse(uint8_t* data, size_t n, uint8_t threshold, uint8_t low_value, uint8_t high_value)
{
    threshold_u8_generic(data, n, threshold, low_value, high_value);
}

void threshold_u16_sse(uint16_t* data, size_t n, uint16_t threshold, uint16_t low_value, uint16_t high_value)
{
    threshold_u16_generic(data, n, threshold, low_value, high_value);
}

void gaussian_blur_u8_sse(uint8_t* dst, const uint8_t* src, int width, int height, int channels, float sigma)
{
    gaussian_blur_u8_generic(dst, src, width, height, channels, sigma);
}

void median_filter_u8_sse(uint8_t* dst, const uint8_t* src, int width, int height, int channels, int kernel_size)
{
    median_filter_u8_generic(dst, src, width, height, channels, kernel_size);
}

#endif // HAVE_SSE
#endif // defined(__x86_64__) || defined(_M_X64)

} // namespace simd
