#ifndef STACKER_OPTIMIZED_QUALITY_TPP
#define STACKER_OPTIMIZED_QUALITY_TPP

#include "../simd/simd_ops.h"
#include <algorithm>
#include <cmath>

// Check if OpenMP is available
#ifdef _OPENMP
#include <omp.h>
#else
// Provide a fallback for systems without OpenMP
#define omp_get_num_threads() 1
#define omp_get_thread_num() 0
#endif

template <typename T>
float Stacker<T>::calculateFrameQuality(const std::shared_ptr<ImageBuffer<T>> &frame)
{
    TIME_FUNCTION("calculateFrameQuality");

    if (!m_referenceFrame)
    {
        return 1.0f; // If no reference frame, assume good quality
    }

    // Check if dimensions match
    if (frame->getWidth() != m_referenceFrame->getWidth() ||
            frame->getHeight() != m_referenceFrame->getHeight() ||
            frame->getChannels() != m_referenceFrame->getChannels())
    {
        LOG_WARN("Frame dimensions do not match reference frame dimensions");
        return 0.0f; // Return poor quality for mismatched dimensions
    }

    // Use SIMD operations for calculating statistics if data is contiguous
    if (frame->isContiguous() && m_referenceFrame->isContiguous())
    {
        // Get pointers to the data
        const T* frameData = frame->getData();
        const T* refData = m_referenceFrame->getData();

        // Calculate the total number of pixels
        size_t totalPixels = frame->getWidth() * frame->getHeight() * frame->getChannels();

        // Use sampling for large images to improve performance
        // Only process a subset of pixels for very large images
        size_t pixelsToProcess = totalPixels;
        size_t stride = 1;

        // For large images (>2MP), use sampling to reduce computation
        if (totalPixels > 2000000)
        {
#if defined(__arm__) || defined(__aarch64__)
            // More aggressive sampling on ARM platforms (like Raspberry Pi)
            stride = 8; // Process every 8th pixel for better performance on ARM
            LOG_INFO("Using aggressive sampling (stride=" + std::to_string(stride) + ") for ARM platform");
#else
            stride = 4; // Process every 4th pixel on other platforms
            LOG_INFO("Using sampling (stride=" + std::to_string(stride) + ") for large image");
#endif
            pixelsToProcess = totalPixels / stride;
        }

        // Calculate statistics using SIMD operations
        float frameMean, frameStdDev;
        float refMean, refStdDev;

        if constexpr (std::is_same_v<T, uint8_t>)
        {
            // For 8-bit images, use SIMD operations
            if (simd::is_simd_enabled())
            {
                // Create sampled data arrays if using sampling
                std::vector<uint8_t> sampledFrameData;
                std::vector<uint8_t> sampledRefData;

                if (stride > 1)
                {
                    sampledFrameData.reserve(pixelsToProcess);
                    sampledRefData.reserve(pixelsToProcess);

                    for (size_t i = 0; i < totalPixels; i += stride)
                    {
                        sampledFrameData.push_back(frameData[i]);
                        sampledRefData.push_back(refData[i]);
                    }

                    // Calculate mean and variance using SIMD
                    frameMean = simd::mean_u8(sampledFrameData.data(), pixelsToProcess);
                    frameStdDev = std::sqrt(simd::variance_u8(sampledFrameData.data(), pixelsToProcess));

                    refMean = simd::mean_u8(sampledRefData.data(), pixelsToProcess);
                    refStdDev = std::sqrt(simd::variance_u8(sampledRefData.data(), pixelsToProcess));
                }
                else
                {
                    // Calculate mean and variance using SIMD on full data
                    frameMean = simd::mean_u8(frameData, totalPixels);
                    frameStdDev = std::sqrt(simd::variance_u8(frameData, totalPixels));

                    refMean = simd::mean_u8(refData, totalPixels);
                    refStdDev = std::sqrt(simd::variance_u8(refData, totalPixels));
                }
            }
            else
            {
                // Fallback to generic implementations if SIMD is disabled
                if (stride > 1)
                {
                    std::vector<uint8_t> sampledFrameData;
                    std::vector<uint8_t> sampledRefData;

                    for (size_t i = 0; i < totalPixels; i += stride)
                    {
                        sampledFrameData.push_back(frameData[i]);
                        sampledRefData.push_back(refData[i]);
                    }

                    frameMean = simd::mean_u8_generic(sampledFrameData.data(), pixelsToProcess);
                    frameStdDev = std::sqrt(simd::variance_u8_generic(sampledFrameData.data(), pixelsToProcess));

                    refMean = simd::mean_u8_generic(sampledRefData.data(), pixelsToProcess);
                    refStdDev = std::sqrt(simd::variance_u8_generic(sampledRefData.data(), pixelsToProcess));
                }
                else
                {
                    frameMean = simd::mean_u8_generic(frameData, totalPixels);
                    frameStdDev = std::sqrt(simd::variance_u8_generic(frameData, totalPixels));

                    refMean = simd::mean_u8_generic(refData, totalPixels);
                    refStdDev = std::sqrt(simd::variance_u8_generic(refData, totalPixels));
                }
            }
        }
        else if constexpr (std::is_same_v<T, uint16_t>)
        {
            // For 16-bit images, use SIMD operations
            if (simd::is_simd_enabled())
            {
                // Create sampled data arrays if using sampling
                std::vector<uint16_t> sampledFrameData;
                std::vector<uint16_t> sampledRefData;

                if (stride > 1)
                {
                    sampledFrameData.reserve(pixelsToProcess);
                    sampledRefData.reserve(pixelsToProcess);

                    for (size_t i = 0; i < totalPixels; i += stride)
                    {
                        sampledFrameData.push_back(frameData[i]);
                        sampledRefData.push_back(refData[i]);
                    }

                    // Calculate mean and variance using SIMD
                    frameMean = simd::mean_u16(sampledFrameData.data(), pixelsToProcess);
                    frameStdDev = std::sqrt(simd::variance_u16(sampledFrameData.data(), pixelsToProcess));

                    refMean = simd::mean_u16(sampledRefData.data(), pixelsToProcess);
                    refStdDev = std::sqrt(simd::variance_u16(sampledRefData.data(), pixelsToProcess));
                }
                else
                {
                    // Calculate mean and variance using SIMD on full data
                    frameMean = simd::mean_u16(frameData, totalPixels);
                    frameStdDev = std::sqrt(simd::variance_u16(frameData, totalPixels));

                    refMean = simd::mean_u16(refData, totalPixels);
                    refStdDev = std::sqrt(simd::variance_u16(refData, totalPixels));
                }
            }
            else
            {
                // Fallback to generic implementations if SIMD is disabled
                if (stride > 1)
                {
                    std::vector<uint16_t> sampledFrameData;
                    std::vector<uint16_t> sampledRefData;

                    for (size_t i = 0; i < totalPixels; i += stride)
                    {
                        sampledFrameData.push_back(frameData[i]);
                        sampledRefData.push_back(refData[i]);
                    }

                    frameMean = simd::mean_u16_generic(sampledFrameData.data(), pixelsToProcess);
                    frameStdDev = std::sqrt(simd::variance_u16_generic(sampledFrameData.data(), pixelsToProcess));

                    refMean = simd::mean_u16_generic(sampledRefData.data(), pixelsToProcess);
                    refStdDev = std::sqrt(simd::variance_u16_generic(sampledRefData.data(), pixelsToProcess));
                }
                else
                {
                    frameMean = simd::mean_u16_generic(frameData, totalPixels);
                    frameStdDev = std::sqrt(simd::variance_u16_generic(frameData, totalPixels));

                    refMean = simd::mean_u16_generic(refData, totalPixels);
                    refStdDev = std::sqrt(simd::variance_u16_generic(refData, totalPixels));
                }
            }
        }

        // Compare standard deviations
        // If they are similar, the frame is likely good quality
        float ratio = std::min(frameStdDev, refStdDev) / std::max(frameStdDev, refStdDev);

        LOG_INFO("Frame quality: " + std::to_string(ratio));
        return ratio;
    }
    else
    {
        // Fallback to a more efficient implementation for non-contiguous data
        // Use OpenMP to parallelize the calculation

        // Calculate the total number of pixels
        int width = frame->getWidth();
        int height = frame->getHeight();
        int channels = frame->getChannels();
        size_t totalPixels = width * height * channels;

        // Use sampling for large images
#if defined(__arm__) || defined(__aarch64__)
        // More aggressive sampling on ARM platforms (like Raspberry Pi)
        int stride = (totalPixels > 2000000) ? 8 : 1;
#else
        int stride = (totalPixels > 2000000) ? 4 : 1;
#endif

        // Calculate frame statistics
        double frameSum = 0.0;
        double frameSumSq = 0.0;
        double refSum = 0.0;
        double refSumSq = 0.0;
        int count = 0;

        #pragma omp parallel for reduction(+:frameSum,frameSumSq,refSum,refSumSq,count) schedule(static)
        for (int y = 0; y < height; y += stride)
        {
            for (int x = 0; x < width; x += stride)
            {
                for (int c = 0; c < channels; c += stride)
                {
                    T frameValue = frame->getPixel(x, y, c);
                    T refValue = m_referenceFrame->getPixel(x, y, c);

                    frameSum += frameValue;
                    frameSumSq += frameValue * frameValue;

                    refSum += refValue;
                    refSumSq += refValue * refValue;

                    count++;
                }
            }
        }

        // Calculate mean and standard deviation
        double frameMean = frameSum / count;
        double frameVariance = (frameSumSq / count) - (frameMean * frameMean);
        double frameStdDev = std::sqrt(frameVariance);

        double refMean = refSum / count;
        double refVariance = (refSumSq / count) - (refMean * refMean);
        double refStdDev = std::sqrt(refVariance);

        // Compare standard deviations
        float ratio = std::min(frameStdDev, refStdDev) / std::max(frameStdDev, refStdDev);

        LOG_INFO("Frame quality: " + std::to_string(ratio));
        return ratio;
    }
}

#endif // STACKER_OPTIMIZED_QUALITY_TPP
