#ifndef STACKER_MEDIAN_FILTER_TPP
#define STACKER_MEDIAN_FILTER_TPP

#include "../io/timer.h"
#include <opencv2/opencv.hpp>

template <typename T>
std::shared_ptr<ImageBuffer<T>> Stacker<T>::applyMedianFilter(const std::shared_ptr<ImageBuffer<T>> &frame, int size)
{
    TIME_FUNCTION("applyMedianFilter");

    // Create a copy of the frame
    auto result = frame->clone();

    // Ensure size is odd (required by OpenCV)
    if (size % 2 == 0)
    {
        size++;
    }

    LOG_INFO("Applying median filter with size " + std::to_string(size));

    // Process each channel separately using OpenCV's medianBlur
    for (int c = 0; c < result->getChannels(); ++c)
    {
        Timer channelTimer("medianFilter_channel" + std::to_string(c));

        // Create OpenCV Mat for the current channel
        cv::Mat channelMat(result->getHeight(), result->getWidth(),
                           std::is_same<T, uint8_t>::value ? CV_8UC1 : CV_16UC1);

        // Copy data from ImageBuffer to OpenCV Mat
        Timer copyToMatTimer("medianFilter_copyToMat");
        for (int y = 0; y < result->getHeight(); ++y)
        {
            for (int x = 0; x < result->getWidth(); ++x)
            {
                if (std::is_same<T, uint8_t>::value)
                {
                    channelMat.at<uint8_t>(y, x) = static_cast<uint8_t>(frame->getPixel(x, y, c));
                }
                else
                {
                    channelMat.at<uint16_t>(y, x) = static_cast<uint16_t>(frame->getPixel(x, y, c));
                }
            }
        }
        copyToMatTimer.getElapsedMs(); // End timer

        // Apply median filter using OpenCV (much faster than manual implementation)
        Timer blurTimer("medianFilter_blur");
        cv::Mat filteredMat;
        cv::medianBlur(channelMat, filteredMat, size);
        blurTimer.getElapsedMs(); // End timer

        // Copy data back from OpenCV Mat to ImageBuffer
        Timer copyBackTimer("medianFilter_copyBack");
        for (int y = 0; y < result->getHeight(); ++y)
        {
            for (int x = 0; x < result->getWidth(); ++x)
            {
                if (std::is_same<T, uint8_t>::value)
                {
                    result->setPixel(x, y, static_cast<T>(filteredMat.at<uint8_t>(y, x)), c);
                }
                else
                {
                    result->setPixel(x, y, static_cast<T>(filteredMat.at<uint16_t>(y, x)), c);
                }
            }
        }
        copyBackTimer.getElapsedMs(); // End timer
    }

    return result;
}

#endif // STACKER_MEDIAN_FILTER_TPP
