#ifndef BACKGROUNDEXTRACTOR_TPP
#define BACKGROUNDEXTRACTOR_TPP

#include "backgroundextractor.h"

template <typename T>
cv::Mat BackgroundExtractor::imageBufferToMat(const std::shared_ptr<ImageBuffer<T>> &image)
{
    cv::Mat mat;

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

    // Create a deep copy to avoid memory issues
    return mat.clone();
}

template <typename T>
std::shared_ptr<ImageBuffer<T>> BackgroundExtractor::matToImageBuffer(const cv::Mat &mat)
{
    int width = mat.cols;
    int height = mat.rows;
    int channels = mat.channels();

    // Create a new ImageBuffer
    auto buffer = std::make_shared<ImageBuffer<T>>(width, height, channels);

    // Convert Mat to the appropriate type if needed
    cv::Mat convertedMat;
    if (std::is_same<T, uint8_t>::value && mat.type() != CV_8UC1 && mat.type() != CV_8UC3)
    {
        mat.convertTo(convertedMat, channels == 1 ? CV_8UC1 : CV_8UC3);
    }
    else if (std::is_same<T, uint16_t>::value && mat.type() != CV_16UC1 && mat.type() != CV_16UC3)
    {
        mat.convertTo(convertedMat, channels == 1 ? CV_16UC1 : CV_16UC3);
    }
    else if (std::is_same<T, float>::value && mat.type() != CV_32FC1 && mat.type() != CV_32FC3)
    {
        mat.convertTo(convertedMat, channels == 1 ? CV_32FC1 : CV_32FC3);
    }
    else
    {
        convertedMat = mat;
    }

    // Copy data from Mat to ImageBuffer
    std::memcpy(buffer->getData(), convertedMat.data, width * height * channels * sizeof(T));

    return buffer;
}

template <typename T>
std::shared_ptr<ImageBuffer<T>> BackgroundExtractor::process(const std::shared_ptr<ImageBuffer<T>> &image)
{
    // Convert ImageBuffer to OpenCV Mat
    cv::Mat mat = imageBufferToMat(image);

    // Create protection mask for nebulae and faint objects
    cv::Mat mask = createProtectionMask(mat);

    // Sample background points
    std::vector<cv::Point3f> points = sampleBackgroundPoints(mat, mask);

    // Fit polynomial surface
    std::vector<double> coeffs = fitPolynomialSurface(points, mat.cols, mat.rows);

    // Generate background model
    cv::Mat backgroundModel = generateBackgroundModelFromCoefficients(coeffs, mat.cols, mat.rows);

    // Apply background model
    cv::Mat result = applyBackgroundModel(mat, backgroundModel);

    // Convert back to ImageBuffer
    return matToImageBuffer<T>(result);
}

template <typename T>
std::shared_ptr<ImageBuffer<T>> BackgroundExtractor::generateBackgroundModel(const std::shared_ptr<ImageBuffer<T>> &image)
{
    // Convert ImageBuffer to OpenCV Mat
    cv::Mat mat = imageBufferToMat(image);

    // Create protection mask for nebulae and faint objects
    cv::Mat mask = createProtectionMask(mat);

    // Sample background points
    std::vector<cv::Point3f> points = sampleBackgroundPoints(mat, mask);

    // Fit polynomial surface
    std::vector<double> coeffs = fitPolynomialSurface(points, mat.cols, mat.rows);

    // Generate background model
    cv::Mat backgroundModel = generateBackgroundModelFromCoefficients(coeffs, mat.cols, mat.rows);

    // Convert to the same type as the input image
    cv::Mat convertedModel;
    if (mat.type() == CV_8U || mat.type() == CV_8UC3)
    {
        backgroundModel.convertTo(convertedModel, CV_8U);
    }
    else if (mat.type() == CV_16U || mat.type() == CV_16UC3)
    {
        backgroundModel.convertTo(convertedModel, CV_16U);
    }
    else
    {
        convertedModel = backgroundModel;
    }

    // Convert back to ImageBuffer
    return matToImageBuffer<T>(convertedModel);
}

#endif // BACKGROUNDEXTRACTOR_TPP
