#ifndef STACKER_LOCAL_HISTOGRAM_EQUALIZATION_TPP
#define STACKER_LOCAL_HISTOGRAM_EQUALIZATION_TPP

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

template <typename T>
std::shared_ptr<ImageBuffer<T>> Stacker<T>::applyLocalHistogramEqualization(
                                 const std::shared_ptr<ImageBuffer<T>> &frame,
                                 int windowSize,
                                 float clipLimit)
{
    TIME_FUNCTION("applyLocalHistogramEqualization");

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

    // Ensure window size is odd
    if (windowSize % 2 == 0)
    {
        windowSize++;
    }

    // Calculate tile size for CLAHE (OpenCV expects tile grid size, not window size)
    // For astronomical images, we want larger tiles to avoid halos around stars
    // A good rule of thumb is to have 8-12 tiles across the image
    int tileGridSize = std::max(2, std::min(8, frame->getWidth() / windowSize));

    // Process each channel separately using OpenCV's CLAHE
    for (int c = 0; c < result->getChannels(); ++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);
        cv::Mat originalMat(result->getHeight(), result->getWidth(),
                            std::is_same<T, uint8_t>::value ? CV_8UC1 : CV_16UC1);

        // Copy data from ImageBuffer to OpenCV Mat
        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));
                    originalMat.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));
                    originalMat.at<uint16_t>(y, x) = static_cast<uint16_t>(frame->getPixel(x, y, c));
                }
            }
        }

        // Create CLAHE object with reduced clip limit for astronomical images
        cv::Ptr<cv::CLAHE> clahe = cv::createCLAHE();
        clahe->setClipLimit(clipLimit);
        clahe->setTilesGridSize(cv::Size(tileGridSize, tileGridSize));

        // Apply CLAHE
        cv::Mat equalizedMat;
        clahe->apply(channelMat, equalizedMat);

        // Apply a small amount of Gaussian blur to reduce halos (sigma = 0.5)
        cv::GaussianBlur(equalizedMat, equalizedMat, cv::Size(3, 3), 0.5);

        // Blend the original and equalized images to reduce artifacts
        // Use 70% of the equalized image and 30% of the original
        cv::addWeighted(equalizedMat, 0.7, originalMat, 0.3, 0, equalizedMat);

        // Copy data back from OpenCV Mat to ImageBuffer
        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>(equalizedMat.at<uint8_t>(y, x)), c);
                }
                else
                {
                    result->setPixel(x, y, static_cast<T>(equalizedMat.at<uint16_t>(y, x)), c);
                }
            }
        }
    }

    return result;
}

#endif // STACKER_LOCAL_HISTOGRAM_EQUALIZATION_TPP
