#ifndef STACKER_DENOISING_TPP
#define STACKER_DENOISING_TPP

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

template <typename T>
std::shared_ptr<ImageBuffer<T>> Stacker<T>::applyDenoising(
                                 const std::shared_ptr<ImageBuffer<T>> &frame,
                                 StackerSettings::DenoisingMethod method,
                                 int strength)
{
    TIME_FUNCTION("applyDenoising");

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

    // Convert to OpenCV Mat for processing
    cv::Mat mat;
    if (result->getChannels() == 1)
    {
        // Single channel image
        if (std::is_same<T, uint8_t>::value)
        {
            mat = cv::Mat(result->getHeight(), result->getWidth(), CV_8UC1, result->getData());
        }
        else if (std::is_same<T, uint16_t>::value)
        {
            mat = cv::Mat(result->getHeight(), result->getWidth(), CV_16UC1, result->getData());
        }
    }
    else if (result->getChannels() == 3)
    {
        // RGB image
        if (std::is_same<T, uint8_t>::value)
        {
            mat = cv::Mat(result->getHeight(), result->getWidth(), CV_8UC3, result->getData());
        }
        else if (std::is_same<T, uint16_t>::value)
        {
            mat = cv::Mat(result->getHeight(), result->getWidth(), CV_16UC3, result->getData());
        }
    }

    // Create a deep copy to avoid memory issues
    cv::Mat processedMat = mat.clone();

    // Apply denoising based on the selected method
    if (method == StackerSettings::DenoisingMethod::NonLocalMeans)
    {
        LOG_INFO("Applying Non-Local Means denoising with strength " + std::to_string(strength));

        // Convert 16-bit to 8-bit if necessary for processing
        cv::Mat tempMat;
        bool is16bit = std::is_same<T, uint16_t>::value;
        if (is16bit)
        {
            processedMat.convertTo(tempMat, processedMat.channels() == 1 ? CV_8UC1 : CV_8UC3, 1.0 / 256.0);
        }
        else
        {
            tempMat = processedMat;
        }

        // Calculate parameters based on strength
        // h parameter controls filtering strength (higher = stronger filtering)
        // For NLM, h is typically in range 3-20
        float h = 3.0f + (strength - 1) * 0.6f;

        // Template window size and search window size
        // For astronomical images, we want a larger search window to find similar patches
        int templateWindowSize = 7;  // Size of compared patches (7x7)
        int searchWindowSize = 21;   // Size of search region (21x21)

        // Apply Non-Local Means denoising
        if (tempMat.channels() == 1)
        {
            // Grayscale image
            cv::fastNlMeansDenoising(tempMat, tempMat, h, templateWindowSize, searchWindowSize);
        }
        else
        {
            // Color image
            cv::fastNlMeansDenoisingColored(tempMat, tempMat, h, h, templateWindowSize, searchWindowSize);
        }

        // Convert back to 16-bit if necessary
        if (is16bit)
        {
            tempMat.convertTo(processedMat, processedMat.channels() == 1 ? CV_16UC1 : CV_16UC3, 256.0);
        }
        else
        {
            processedMat = tempMat;
        }
    }
    else if (method == StackerSettings::DenoisingMethod::Bilateral)
    {
        LOG_INFO("Applying Bilateral Filter denoising with strength " + std::to_string(strength));

        // Convert 16-bit to 8-bit if necessary for processing
        cv::Mat tempMat;
        bool is16bit = std::is_same<T, uint16_t>::value;
        if (is16bit)
        {
            processedMat.convertTo(tempMat, processedMat.channels() == 1 ? CV_8UC1 : CV_8UC3, 1.0 / 256.0);
        }
        else
        {
            tempMat = processedMat;
        }

        // Calculate parameters based on strength
        // d: Diameter of each pixel neighborhood (larger = slower)
        // sigmaColor: Filter sigma in the color space (larger = more aggressive filtering)
        // sigmaSpace: Filter sigma in the coordinate space (larger = larger features preserved)
        int d = 9;  // Fixed diameter for performance
        double sigmaColor = 5.0 + strength * 1.0;  // Scale with strength
        double sigmaSpace = 5.0 + strength * 0.5;  // Scale with strength but less aggressively

        // Process each channel separately for color images
        if (tempMat.channels() == 3)
        {
            cv::Mat channels[3];
            cv::Mat filteredChannels[3];
            cv::split(tempMat, channels);

            for (int i = 0; i < 3; i++)
            {
                // Create a separate destination matrix for each channel
                filteredChannels[i] = cv::Mat(channels[i].size(), channels[i].type());
                cv::bilateralFilter(channels[i], filteredChannels[i], d, sigmaColor, sigmaSpace);
            }

            cv::merge(filteredChannels, 3, tempMat);
        }
        else
        {
            // For grayscale images
            cv::Mat filteredMat = cv::Mat(tempMat.size(), tempMat.type());
            cv::bilateralFilter(tempMat, filteredMat, d, sigmaColor, sigmaSpace);
            tempMat = filteredMat;
        }

        // Convert back to 16-bit if necessary
        if (is16bit)
        {
            tempMat.convertTo(processedMat, processedMat.channels() == 1 ? CV_16UC1 : CV_16UC3, 256.0);
        }
        else
        {
            processedMat = tempMat;
        }
    }

    // Save the denoised image for diagnostic purposes if needed
    if (m_settings->getDiagnosticMode() && !m_settings->getOutputPrefix().empty())
    {
        std::string denoisedFilename = m_settings->getOutputPrefix() + "denoised.jpg";
        cv::imwrite(denoisedFilename, processedMat);
        LOG_INFO("Saved denoised image to " + denoisedFilename);
    }

    // Copy the processed Mat back to the ImageBuffer
    std::memcpy(result->getData(), processedMat.data,
                result->getWidth() * result->getHeight() * result->getChannels() * sizeof(T));

    return result;
}

#endif // STACKER_DENOISING_TPP
