#ifndef STACKER_STACK_TPP
#define STACKER_STACK_TPP

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

template <typename T>
std::shared_ptr<ImageBuffer<T>> Stacker<T>::stackMean()
{
    TIME_FUNCTION("stackMean");

    // Start benchmark timer if benchmark mode is enabled
    BenchmarkTimer benchmarkTimer;

    if (m_frames.empty())
    {
        throw std::runtime_error("No frames to stack");
    }

    if (!m_referenceFrame)
    {
        m_referenceFrame = m_frames[0];
    }

    // Create a buffer for the stacked image
    auto result = std::make_shared<ImageBuffer<T>>(
                      m_referenceFrame->getWidth(),
                      m_referenceFrame->getHeight(),
                      m_referenceFrame->getChannels()
                  );
    result->fill(0);

    // Process each frame
    int frameCount = 0;
    for (const auto &frame : m_frames)
    {
        Timer frameTimer("Frame " + std::to_string(++frameCount));

        // Apply dark subtraction if needed
        auto processedFrame = m_darkFrame ? applyDarkSubtraction(frame) : frame;

        // Align frame if needed
        if (m_settings->getAlignFrames())
        {
            Timer alignTimer("Align frame " + std::to_string(frameCount));
            processedFrame = alignFrame(processedFrame);
        }

        // Calculate frame quality
        float quality = calculateFrameQuality(processedFrame);

        // Skip frames with poor quality
        if (quality < m_settings->getQualityThreshold())
        {
            LOG_WARN("Skipping frame " + std::to_string(frameCount) + " due to poor quality: " + std::to_string(quality));
            continue;
        }

        // Add to the result
        {
            Timer addTimer("Add frame " + std::to_string(frameCount) + " to result");

            // Check if we can use direct memory access for better performance
            if (result->isContiguous() && processedFrame->isContiguous())
            {
                // Get pointers to the data
                T* resultData = result->getData();
                const T* frameData = processedFrame->getData();

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

                // Process all pixels in a single loop using SIMD
                if constexpr (std::is_same_v<T, uint8_t>)
                {
                    // For 8-bit images, use SIMD operations
                    float factor = 1.0f / m_frames.size();

#if defined(__arm__) || defined(__aarch64__)
                    // Optimized for ARM (Raspberry Pi) - process in smaller chunks to fit in cache
                    const size_t chunkSize = 4096; // 4KB chunks for better cache utilization

                    // Use static buffer to avoid repeated allocations
                    static thread_local std::vector<uint8_t> scaledFrameData;
                    if (scaledFrameData.size() < chunkSize)
                        scaledFrameData.resize(chunkSize);

                    // Process in chunks
                    for (size_t offset = 0; offset < totalPixels; offset += chunkSize)
                    {
                        // Calculate actual chunk size (might be smaller for the last chunk)
                        size_t currentChunkSize = std::min(chunkSize, totalPixels - offset);

                        // Copy chunk to temporary buffer
                        std::memcpy(scaledFrameData.data(), frameData + offset, currentChunkSize);

                        // Scale the chunk by the factor
                        simd::multiply_u8(scaledFrameData.data(), factor, currentChunkSize);

                        // Add to the result
                        simd::add_u8(resultData + offset, scaledFrameData.data(), currentChunkSize);
                    }
#else
                    // For non-ARM platforms, use the original approach
                    // Create a temporary buffer for scaled frame data
                    std::vector<uint8_t> scaledFrameData(totalPixels);
                    std::memcpy(scaledFrameData.data(), frameData, totalPixels);

                    // Scale the frame data by the factor
                    simd::multiply_u8(scaledFrameData.data(), factor, totalPixels);

                    // Add to the result
                    simd::add_u8(resultData, scaledFrameData.data(), totalPixels);
#endif
                }
                else if constexpr (std::is_same_v<T, uint16_t>)
                {
                    // For 16-bit images, use SIMD operations
                    float factor = 1.0f / m_frames.size();

#if defined(__arm__) || defined(__aarch64__)
                    // Optimized for ARM (Raspberry Pi) - process in smaller chunks to fit in cache
                    const size_t chunkSize = 2048; // 4KB chunks (2048 * 2 bytes) for better cache utilization

                    // Use static buffer to avoid repeated allocations
                    static thread_local std::vector<uint16_t> scaledFrameData;
                    if (scaledFrameData.size() < chunkSize)
                        scaledFrameData.resize(chunkSize);

                    // Process in chunks
                    for (size_t offset = 0; offset < totalPixels; offset += chunkSize)
                    {
                        // Calculate actual chunk size (might be smaller for the last chunk)
                        size_t currentChunkSize = std::min(chunkSize, totalPixels - offset);

                        // Copy chunk to temporary buffer
                        std::memcpy(scaledFrameData.data(), frameData + offset, currentChunkSize * sizeof(uint16_t));

                        // Scale the chunk by the factor
                        simd::multiply_u16(scaledFrameData.data(), factor, currentChunkSize);

                        // Add to the result
                        simd::add_u16(resultData + offset, scaledFrameData.data(), currentChunkSize);
                    }
#else
                    // For non-ARM platforms, use the original approach
                    // Create a temporary buffer for scaled frame data
                    std::vector<uint16_t> scaledFrameData(totalPixels);
                    std::memcpy(scaledFrameData.data(), frameData, totalPixels * sizeof(uint16_t));

                    // Scale the frame data by the factor
                    simd::multiply_u16(scaledFrameData.data(), factor, totalPixels);

                    // Add to the result
                    simd::add_u16(resultData, scaledFrameData.data(), totalPixels);
#endif
                }
                else
                {
                    // Fallback for other types
                    for (size_t i = 0; i < totalPixels; ++i)
                    {
                        resultData[i] += frameData[i] / m_frames.size();
                    }
                }
            }
            else
            {
                // Fallback to pixel-by-pixel processing
                for (int y = 0; y < result->getHeight(); ++y)
                {
                    for (int x = 0; x < result->getWidth(); ++x)
                    {
                        for (int c = 0; c < result->getChannels(); ++c)
                        {
                            T value = processedFrame->getPixel(x, y, c);
                            T currentValue = result->getPixel(x, y, c);
                            result->setPixel(x, y, currentValue + value / m_frames.size(), c);
                        }
                    }
                }
            }

            // Save intermediate result if diagnostic mode is enabled
            saveIntermediateResult(result, frameCount);
        }
    }

    // Apply post-processing
    {
        Timer postTimer("Post-processing");
        result = applyPostProcessing(result);
    }

    // Display total stacking time if benchmark mode is enabled
    if (m_settings->getBenchmarkMode())
    {
        long long totalTime = benchmarkTimer.getElapsedMs();
        LOG_INFO("Total stacking time: " + std::to_string(totalTime) + " ms");
    }

    return result;
}

template <typename T>
std::shared_ptr<ImageBuffer<T>> Stacker<T>::stackMedian()
{
    TIME_FUNCTION("stackMedian");

    // Start benchmark timer if benchmark mode is enabled
    BenchmarkTimer benchmarkTimer;

    if (m_frames.empty())
    {
        throw std::runtime_error("No frames to stack");
    }

    if (!m_referenceFrame)
    {
        m_referenceFrame = m_frames[0];
    }

    // Create a buffer for the stacked image
    auto result = std::make_shared<ImageBuffer<T>>(
                      m_referenceFrame->getWidth(),
                      m_referenceFrame->getHeight(),
                      m_referenceFrame->getChannels()
                  );

    // Pre-process all frames first to avoid repeated alignment
    LOG_INFO("Pre-processing frames for median stacking...");
    std::vector<std::shared_ptr<ImageBuffer<T>>> processedFrames;
    processedFrames.reserve(m_frames.size());

    {
        Timer preProcessTimer("Pre-processing all frames");
        for (const auto &frame : m_frames)
        {
            // Apply dark subtraction if needed
            auto processedFrame = m_darkFrame ? applyDarkSubtraction(frame) : frame;

            // Align frame if needed
            if (m_settings->getAlignFrames())
            {
                Timer alignTimer("Align frame");
                processedFrame = alignFrame(processedFrame);
            }

            processedFrames.push_back(processedFrame);
        }
    }

    // Process each pixel in parallel
    {
        Timer pixelProcessTimer("Processing all pixels");
        int height = result->getHeight();
        int width = result->getWidth();
        int channels = result->getChannels();

        // Use OpenMP to parallelize the outer loop
        #pragma omp parallel
        {
            // Log the number of threads being used (only once)
            #pragma omp single
            {
                LOG_INFO("Using " + std::to_string(omp_get_num_threads()) + " threads for parallel processing");
            }

            // Parallelize the row processing
            #pragma omp for schedule(dynamic, 10)
            for (int y = 0; y < height; ++y)
            {
                // Log progress occasionally
                if (y % 100 == 0)
                {
                    #pragma omp critical
                    {
                        LOG_INFO("Processing row " + std::to_string(y) + " of " + std::to_string(height));
                    }
                }

                for (int x = 0; x < width; ++x)
                {
                    for (int c = 0; c < channels; ++c)
                    {
                        // Collect values for this pixel from all frames
                        std::vector<T> values;
                        values.reserve(processedFrames.size());

                        for (const auto &processedFrame : processedFrames)
                        {
                            values.push_back(processedFrame->getPixel(x, y, c));
                        }

                        // Calculate median using SIMD if possible
                        T medianValue;
                        if constexpr (std::is_same_v<T, uint8_t>)
                        {
                            // For 8-bit images with enough frames, use SIMD median
                            if (processedFrames.size() >= 8)
                                medianValue = simd::median_u8(values.data(), values.size());
                            else
                                medianValue = simd::median_u8_generic(values.data(), values.size());
                        }
                        else if constexpr (std::is_same_v<T, uint16_t>)
                        {
                            // For 16-bit images with enough frames, use SIMD median
                            medianValue = simd::median_u16(values.data(), values.size());
                        }
                        else
                        {
                            // Fallback to standard sort and median
                            std::sort(values.begin(), values.end());
                            medianValue = values[values.size() / 2];
                        }

                        result->setPixel(x, y, medianValue, c);
                    }
                }
            }
            LOG_INFO("Finished processing all " + std::to_string(result->getHeight()) + " rows");
        }

    }

    // Apply post-processing
    {
        Timer postProcessTimer("Post-processing");
        result = applyPostProcessing(result);
    }

    // Display total stacking time if benchmark mode is enabled
    if (m_settings->getBenchmarkMode())
    {
        long long totalTime = benchmarkTimer.getElapsedMs();
        LOG_INFO("Total stacking time: " + std::to_string(totalTime) + " ms");
    }

    return result;
}

template <typename T>
std::shared_ptr<ImageBuffer<T>> Stacker<T>::stackSigmaClip()
{
    TIME_FUNCTION("stackSigmaClip");

    // Start benchmark timer if benchmark mode is enabled
    BenchmarkTimer benchmarkTimer;

    if (m_frames.empty())
    {
        throw std::runtime_error("No frames to stack");
    }

    if (!m_referenceFrame)
    {
        m_referenceFrame = m_frames[0];
    }

    // Check if we have enough frames for sigma clipping
    if (m_frames.size() < m_settings->getSigmaClippingFrames())
    {
        // Fall back to mean stacking
        LOG_WARN("Not enough frames for sigma clipping, falling back to mean stacking");
        return stackMean();
    }

    // Create a buffer for the stacked image
    auto result = std::make_shared<ImageBuffer<T>>(
                      m_referenceFrame->getWidth(),
                      m_referenceFrame->getHeight(),
                      m_referenceFrame->getChannels()
                  );

    // Pre-process all frames first to avoid repeated alignment
    LOG_INFO("Pre-processing frames for sigma clipping...");
    std::vector<std::shared_ptr<ImageBuffer<T>>> processedFrames;
    processedFrames.reserve(m_frames.size());

    {
        Timer preProcessTimer("Pre-processing all frames");
        for (const auto &frame : m_frames)
        {
            // Apply dark subtraction if needed
            auto processedFrame = m_darkFrame ? applyDarkSubtraction(frame) : frame;

            // Align frame if needed
            if (m_settings->getAlignFrames())
            {
                Timer alignTimer("Align frame");
                processedFrame = alignFrame(processedFrame);
            }

            processedFrames.push_back(processedFrame);
        }
    }

    // Process each pixel in parallel
    {
        Timer pixelProcessTimer("Processing all pixels");
        int height = result->getHeight();
        int width = result->getWidth();
        int channels = result->getChannels();

        // Use OpenMP to parallelize the outer loop
        #pragma omp parallel
        {
            // Log the number of threads being used (only once)
            #pragma omp single
            {
                LOG_INFO("Using " + std::to_string(omp_get_num_threads()) + " threads for parallel processing");
            }

            // Parallelize the row processing
            #pragma omp for schedule(dynamic, 10)
            for (int y = 0; y < height; ++y)
            {
                // Log progress occasionally
                if (y % 100 == 0)
                {
                    #pragma omp critical
                    {
                        LOG_INFO("Processing row " + std::to_string(y) + " of " + std::to_string(height));
                    }
                }

                for (int x = 0; x < width; ++x)
                {
                    for (int c = 0; c < channels; ++c)
                    {
                        // Collect values for this pixel from all frames
                        std::vector<T> values;
                        values.reserve(processedFrames.size());

                        for (const auto &processedFrame : processedFrames)
                        {
                            values.push_back(processedFrame->getPixel(x, y, c));
                        }

                        // Calculate mean and standard deviation using SIMD if possible
                        double mean, stdev;

                        if constexpr (std::is_same_v<T, uint8_t>)
                        {
                            // For 8-bit images with enough frames, use SIMD mean and variance
                            if (processedFrames.size() >= 8)
                            {
                                mean = simd::mean_u8(values.data(), values.size());
                                stdev = std::sqrt(simd::variance_u8(values.data(), values.size()));
                            }
                            else
                            {
                                mean = simd::mean_u8_generic(values.data(), values.size());
                                stdev = std::sqrt(simd::variance_u8_generic(values.data(), values.size()));
                            }
                        }
                        else if constexpr (std::is_same_v<T, uint16_t>)
                        {
                            // For 16-bit images with enough frames, use SIMD mean and variance
                            mean = simd::mean_u16(values.data(), values.size());
                            stdev = std::sqrt(simd::variance_u16(values.data(), values.size()));
                        }
                        else
                        {
                            // Fallback to standard accumulate and inner_product
                            double sum = std::accumulate(values.begin(), values.end(), 0.0);
                            mean = sum / values.size();

                            double sq_sum = std::inner_product(values.begin(), values.end(), values.begin(), 0.0);
                            stdev = std::sqrt(sq_sum / values.size() - mean * mean);
                        }

                        // Filter values outside the sigma range
                        std::vector<T> filtered_values;
                        filtered_values.reserve(values.size());

                        for (T value : values)
                        {
                            double deviation = std::abs(value - mean);
                            if (deviation <= m_settings->getSigmaClippingUpper() * stdev && deviation >= m_settings->getSigmaClippingLower() * stdev)
                            {
                                filtered_values.push_back(value);
                            }
                        }

                        // Calculate mean of filtered values
                        if (filtered_values.empty())
                        {
                            // If all values were filtered out, use the original mean
                            result->setPixel(x, y, static_cast<T>(mean), c);
                        }
                        else
                        {
                            double filtered_sum = std::accumulate(filtered_values.begin(), filtered_values.end(), 0.0);
                            double filtered_mean = filtered_sum / filtered_values.size();
                            result->setPixel(x, y, static_cast<T>(filtered_mean), c);
                        }
                    }
                }
            }
            LOG_INFO("Finished processing all " + std::to_string(result->getHeight()) + " rows");
        }

    }

    // Apply post-processing
    {
        Timer postProcessTimer("Post-processing");
        result = applyPostProcessing(result);
    }

    // Display total stacking time if benchmark mode is enabled
    if (m_settings->getBenchmarkMode())
    {
        long long totalTime = benchmarkTimer.getElapsedMs();
        LOG_INFO("Total stacking time: " + std::to_string(totalTime) + " ms");
    }

    return result;
}

#endif // STACKER_STACK_TPP
