#include "backgroundextractor.h"
#include "star_detector.h"
#include <algorithm>
#include <cmath>
#include <iostream>
#include "../io/logger.h"
#include <memory>
#include <atomic>

#ifdef HAVE_OPENMP
#include <omp.h>
#endif

BackgroundExtractor::BackgroundExtractor()
    : m_polynomialDegree(2)
    , m_sampleGridSize(16)
    , m_protectionThreshold(0.5f)
    , m_smoothnessFactor(0.5f)
{
}

BackgroundExtractor::~BackgroundExtractor() = default;

void BackgroundExtractor::setPolynomialDegree(int degree)
{
    m_polynomialDegree = std::max(1, std::min(4, degree));
}

void BackgroundExtractor::setSampleGridSize(int gridSize)
{
    m_sampleGridSize = std::max(8, std::min(64, gridSize));
}

void BackgroundExtractor::setProtectionThreshold(float threshold)
{
    m_protectionThreshold = std::max(0.0f, std::min(1.0f, threshold));
}

void BackgroundExtractor::setSmoothnessFactor(float smoothness)
{
    m_smoothnessFactor = std::max(0.0f, std::min(1.0f, smoothness));
}

cv::Mat BackgroundExtractor::detectObjects(const cv::Mat &image)
{
    // Try StarDetector first, fall back to OpenCV if it fails
    try
    {
        StarDetector detector;
        return detector.detectStars(image);
    }
    catch (const std::exception &e)
    {
        LOG_ERROR("StarDetector failed: " + std::string(e.what()));
        LOG_ERROR("Falling back to OpenCV detection");
        return detectStarsWithOpenCV(image);
    }
}

cv::Mat BackgroundExtractor::detectStarsWithOpenCV(const cv::Mat &image)
{
    // Convert to grayscale and 8-bit if needed
    cv::Mat gray;
    if (image.channels() > 1)
    {
        // Convert to grayscale
        cv::cvtColor(image, gray, cv::COLOR_BGR2GRAY);
    }
    else
    {
        gray = image.clone();
    }

    // Convert to 8-bit if needed
    if (gray.type() != CV_8U)
    {
        cv::Mat temp;
        cv::normalize(gray, temp, 0, 255, cv::NORM_MINMAX, CV_8U);
        gray = temp;
    }

    // Set up blob detector parameters
    cv::SimpleBlobDetector::Params params;
    params.filterByArea = true;
    params.minArea = 3;
    params.maxArea = 100;
    params.filterByCircularity = true;
    params.minCircularity = 0.5f;
    params.filterByConvexity = true;
    params.minConvexity = 0.8f;
    params.filterByInertia = true;
    params.minInertiaRatio = 0.5f;

    // Create detector
    cv::Ptr<cv::SimpleBlobDetector> detector = cv::SimpleBlobDetector::create(params);

    // Detect blobs
    std::vector<cv::KeyPoint> keypoints;
    detector->detect(gray, keypoints);

    // Create mask
    cv::Mat mask = cv::Mat::zeros(image.size(), CV_8U);

    // Draw detected stars on mask
    for (const auto &kp : keypoints)
    {
        int radius = static_cast<int>(kp.size * 1.5f);
        cv::circle(mask, kp.pt, radius, cv::Scalar(255), -1);
    }

    return mask;
}

cv::Mat BackgroundExtractor::createProtectionMask(const cv::Mat &image)
{
    // Convert to grayscale and 8-bit if needed
    cv::Mat gray;
    if (image.channels() > 1)
    {
        // Convert to grayscale
        cv::cvtColor(image, gray, cv::COLOR_BGR2GRAY);
    }
    else
    {
        gray = image.clone();
    }

    // Convert to 8-bit if needed
    if (gray.type() != CV_8U)
    {
        cv::Mat temp;
        cv::normalize(gray, temp, 0, 255, cv::NORM_MINMAX, CV_8U);
        gray = temp;
    }

    // Create a copy of the original image for color-based processing
    cv::Mat colorImage;
    if (image.channels() > 1)
    {
        colorImage = image.clone();
    }
    else
    {
        cv::cvtColor(gray, colorImage, cv::COLOR_GRAY2BGR);
    }

    // Apply CLAHE to enhance local contrast
    cv::Ptr<cv::CLAHE> clahe = cv::createCLAHE(2.0, cv::Size(8, 8));
    cv::Mat enhanced;
    clahe->apply(gray, enhanced);

    // Detect stars first using our specialized detector
    cv::Mat starMask = detectObjects(image);

    // ---- Method 1: Edge-based detection ----
    // Apply Canny edge detection
    cv::Mat edges;
    cv::Canny(enhanced, edges, 50, 150);

    // Dilate edges to connect nearby edges
    cv::Mat dilatedEdges;
    cv::Mat edgeElement = cv::getStructuringElement(cv::MORPH_ELLIPSE, cv::Size(3, 3));
    cv::dilate(edges, dilatedEdges, edgeElement);

    // Find contours in the edge image
    std::vector<std::vector<cv::Point>> contours;
    cv::findContours(dilatedEdges, contours, cv::RETR_EXTERNAL, cv::CHAIN_APPROX_SIMPLE);

    // Create a mask from the contours
    cv::Mat contourMask = cv::Mat::zeros(gray.size(), CV_8U);
    for (const auto &contour : contours)
    {
        // Filter contours by area to remove noise
        double area = cv::contourArea(contour);
        if (area > 10.0)  // Minimum area threshold
        {
            cv::drawContours(contourMask, std::vector<std::vector<cv::Point>> {contour}, 0, cv::Scalar(255), -1);
        }
    }

    // ---- Method 2: Adaptive thresholding ----
    // Apply Gaussian blur to smooth the image before thresholding
    cv::Mat blurred;
    cv::GaussianBlur(enhanced, blurred, cv::Size(5, 5), 0);

    // Use adaptive thresholding to better handle varying brightness
    cv::Mat adaptiveBinary;
    cv::adaptiveThreshold(blurred, adaptiveBinary, 255, cv::ADAPTIVE_THRESH_GAUSSIAN_C,
                          cv::THRESH_BINARY, 11, -2);

    // ---- Method 3: Otsu's thresholding ----
    // Use Otsu's method for global thresholding
    cv::Mat otsuBinary;
    double otsuThresh = cv::threshold(enhanced, otsuBinary, 0, 255, cv::THRESH_BINARY | cv::THRESH_OTSU);

    // Adjust threshold based on protection threshold parameter
    double adjustedThresh = otsuThresh * m_protectionThreshold;
    cv::Mat adjustedBinary;
    cv::threshold(enhanced, adjustedBinary, adjustedThresh, 255, cv::THRESH_BINARY);

    // ---- Method 4: Color-based segmentation using K-means ----
    cv::Mat kMeansMask = cv::Mat::zeros(gray.size(), CV_8U);

    // Reshape the image to a single column of pixels (for k-means)
    cv::Mat reshapedImage = colorImage.reshape(3, colorImage.rows * colorImage.cols);
    reshapedImage.convertTo(reshapedImage, CV_32F);

    // Define criteria for k-means
    cv::TermCriteria criteria(cv::TermCriteria::EPS + cv::TermCriteria::MAX_ITER, 10, 1.0);

    // Number of clusters (background + foreground)
    int k = 2;
    cv::Mat labels, centers;

    // Apply k-means clustering
    cv::kmeans(reshapedImage, k, labels, criteria, 3, cv::KMEANS_PP_CENTERS, centers);

    // Reshape labels back to image dimensions
    labels = labels.reshape(1, colorImage.rows);

    // Determine which cluster is foreground (usually the one with higher intensity)
    float bgIntensity = 0, fgIntensity = 0;
    for (int i = 0; i < k; i++)
    {
        cv::Vec3f center = centers.at<cv::Vec3f>(i, 0);
        float intensity = (center[0] + center[1] + center[2]) / 3.0f;
        if (i == 0)
        {
            bgIntensity = intensity;
            fgIntensity = intensity;
        }
        else
        {
            if (intensity > fgIntensity)
                fgIntensity = intensity;
            if (intensity < bgIntensity)
                bgIntensity = intensity;
        }
    }

    // Create mask based on cluster labels
    for (int y = 0; y < labels.rows; y++)
    {
        for (int x = 0; x < labels.cols; x++)
        {
            int cluster = labels.at<int>(y, x);
            cv::Vec3f center = centers.at<cv::Vec3f>(cluster, 0);
            float intensity = (center[0] + center[1] + center[2]) / 3.0f;

            // If pixel belongs to foreground cluster
            if (std::abs(intensity - fgIntensity) < std::abs(intensity - bgIntensity))
            {
                kMeansMask.at<uchar>(y, x) = 255;
            }
        }
    }

    // ---- Combine all masks ----
    // Combine all binary masks
    cv::Mat combinedBinary;
    cv::bitwise_or(adaptiveBinary, adjustedBinary, combinedBinary);
    cv::bitwise_or(combinedBinary, contourMask, combinedBinary);
    cv::bitwise_or(combinedBinary, kMeansMask, combinedBinary);

    // Apply morphological operations to connect nearby regions
    cv::Mat element = cv::getStructuringElement(cv::MORPH_ELLIPSE, cv::Size(7, 7));
    cv::Mat mask;
    cv::morphologyEx(combinedBinary, mask, cv::MORPH_CLOSE, element);

    // Dilate the mask slightly to ensure we protect the edges of objects
    cv::dilate(mask, mask, element, cv::Point(-1, -1), 1);

    // Combine with star mask
    cv::bitwise_or(mask, starMask, mask);

    // Calculate the percentage of masked pixels
    double maskedPercentage = (cv::countNonZero(mask) * 100.0) / (mask.rows * mask.cols);
    LOG_INFO("Protection mask covers " + std::to_string(maskedPercentage) + "% of the image");

    // If the mask covers too much of the image (>80%), it's likely incorrect
    // In this case, revert to a more conservative approach
    if (maskedPercentage > 80.0)
    {
        LOG_WARN("Warning: Protection mask covers too much of the image. Adjusting...");

        // Use a higher threshold to be more selective
        cv::threshold(enhanced, mask, otsuThresh * 1.2, 255, cv::THRESH_BINARY);

        // Only keep the star mask and the brightest parts
        cv::bitwise_or(mask, starMask, mask);

        // Recalculate the percentage
        maskedPercentage = (cv::countNonZero(mask) * 100.0) / (mask.rows * mask.cols);
        LOG_INFO("Adjusted protection mask covers " + std::to_string(maskedPercentage) + "% of the image");
    }

    return mask;
}

std::vector<cv::Point3f> BackgroundExtractor::sampleBackgroundPoints(const cv::Mat &image, const cv::Mat &mask)
{
    std::vector<cv::Point3f> points;

    // Convert to grayscale if needed
    cv::Mat grayImage;
    if (image.channels() > 1)
    {
        cv::cvtColor(image, grayImage, cv::COLOR_BGR2GRAY);
    }
    else
    {
        grayImage = image.clone();
    }

    // Calculate grid step
    int stepX = std::max(1, grayImage.cols / m_sampleGridSize);
    int stepY = std::max(1, grayImage.rows / m_sampleGridSize);

    // Count of masked points
    int maskedPoints = 0;
    int totalPoints = 0;

    // Calculate the minimum number of points needed based on polynomial degree
    int minPointsNeeded;
    switch (m_polynomialDegree)
    {
        case 1:
            minPointsNeeded = 4;
            break;  // 1 + x + y (3 coeffs)
        case 2:
            minPointsNeeded = 7;
            break;  // 1 + x + y + x^2 + xy + y^2 (6 coeffs)
        case 3:
            minPointsNeeded = 11;
            break; // 1 + x + y + x^2 + xy + y^2 + x^3 + x^2y + xy^2 + y^3 (10 coeffs)
        case 4:
            minPointsNeeded = 16;
            break; // 1 + x + y + x^2 + xy + y^2 + x^3 + x^2y + xy^2 + y^3 + x^4 + ... (15 coeffs)
        default:
            minPointsNeeded = 7;
            break; // Default to degree 2
    }

    // Add 50% more points to ensure a good fit
    minPointsNeeded = static_cast<int>(minPointsNeeded * 1.5);

    // First pass: sample points on a grid, only adding unmasked points
    // Use thread-local vectors to avoid contention
    std::vector<cv::Point3f> threadPoints;
    int localTotalPoints = 0;
    int localMaskedPoints = 0;

    #pragma omp parallel private(threadPoints) reduction(+:localTotalPoints, localMaskedPoints)
    {
#ifdef HAVE_OPENMP
        threadPoints.reserve(grayImage.rows * grayImage.cols / (stepX * stepY) / omp_get_num_threads());
#else
        threadPoints.reserve(grayImage.rows * grayImage.cols / (stepX * stepY));
#endif

        #pragma omp for schedule(dynamic, 16)
        for (int y = stepY / 2; y < grayImage.rows; y += stepY)
        {
            for (int x = stepX / 2; x < grayImage.cols; x += stepX)
            {
                localTotalPoints++;

                // Check if point is masked
                bool isMasked = false;
                if (mask.at<uchar>(y, x) > 0)
                {
                    localMaskedPoints++;
                    isMasked = true;
                }

                // Get pixel value
                float value;
                if (grayImage.type() == CV_8U)
                {
                    value = static_cast<float>(grayImage.at<uchar>(y, x));
                }
                else if (grayImage.type() == CV_16U)
                {
                    value = static_cast<float>(grayImage.at<ushort>(y, x));
                }
                else if (grayImage.type() == CV_32F)
                {
                    value = grayImage.at<float>(y, x);
                }
                else
                {
                    continue; // Unsupported type
                }

                // Add point if not masked
                if (!isMasked)
                {
                    threadPoints.push_back(cv::Point3f(x, y, value));
                }
            }
        }

        // Combine thread-local points
        #pragma omp critical
        {
            points.insert(points.end(), threadPoints.begin(), threadPoints.end());
        }
    }

    // Update counters
    totalPoints = localTotalPoints;
    maskedPoints = localMaskedPoints;

    // Print statistics from first pass
    LOG_INFO("First pass: Sampled " + std::to_string(points.size()) + " background points out of " +
             std::to_string(totalPoints) + " total points");
    LOG_INFO("Masked points: " + std::to_string(maskedPoints) + " (" +
             std::to_string(static_cast<double>(maskedPoints) / totalPoints * 100.0) + "%)");

    // If we have too few points, try a second pass with a finer grid
    if (points.size() < minPointsNeeded)
    {
        LOG_INFO("Not enough background points. Trying with a finer grid...");

        // Use a finer grid
        stepX = std::max(1, stepX / 2);
        stepY = std::max(1, stepY / 2);

        // Second pass: sample points on a finer grid
        for (int y = stepY / 2; y < grayImage.rows; y += stepY)
        {
            for (int x = stepX / 2; x < grayImage.cols; x += stepX)
            {
                // Skip points that are too close to existing points
                bool tooClose = false;
                for (const auto &p : points)
                {
                    if (std::abs(p.x - x) < stepX / 2 && std::abs(p.y - y) < stepY / 2)
                    {
                        tooClose = true;
                        break;
                    }
                }

                if (tooClose)
                    continue;

                // Check if point is masked
                bool isMasked = false;
                if (mask.at<uchar>(y, x) > 0)
                {
                    isMasked = true;
                }

                // Get pixel value
                float value;
                if (grayImage.type() == CV_8U)
                {
                    value = static_cast<float>(grayImage.at<uchar>(y, x));
                }
                else if (grayImage.type() == CV_16U)
                {
                    value = static_cast<float>(grayImage.at<ushort>(y, x));
                }
                else if (grayImage.type() == CV_32F)
                {
                    value = grayImage.at<float>(y, x);
                }
                else
                {
                    continue; // Unsupported type
                }

                // Add point if not masked
                if (!isMasked)
                {
                    points.push_back(cv::Point3f(x, y, value));
                }
            }
        }

        LOG_INFO("Second pass: Now have " + std::to_string(points.size()) + " background points");
    }

    // If we still have too few points, add some points from the edges and corners
    if (points.size() < minPointsNeeded)
    {
        LOG_INFO("Still not enough background points. Adding edge and corner points...");

        // Add corners and edge points
        int margin = 10;
        int numEdgePoints = 10; // Number of points to sample along each edge

        // Add corners
        std::vector<cv::Point> corners =
        {
            cv::Point(margin, margin),
            cv::Point(margin, grayImage.rows - margin),
            cv::Point(grayImage.cols - margin, margin),
            cv::Point(grayImage.cols - margin, grayImage.rows - margin)
        };

        for (const auto &corner : corners)
        {
            // Skip if masked
            if (mask.at<uchar>(corner.y, corner.x) > 0)
                continue;

            float value;
            if (grayImage.type() == CV_8U)
            {
                value = static_cast<float>(grayImage.at<uchar>(corner.y, corner.x));
            }
            else if (grayImage.type() == CV_16U)
            {
                value = static_cast<float>(grayImage.at<ushort>(corner.y, corner.x));
            }
            else if (grayImage.type() == CV_32F)
            {
                value = grayImage.at<float>(corner.y, corner.x);
            }
            else
            {
                continue; // Unsupported type
            }

            points.push_back(cv::Point3f(corner.x, corner.y, value));
        }

        // Add edge points
        // Top edge
        for (int i = 1; i < numEdgePoints - 1; i++)
        {
            int x = margin + i * (grayImage.cols - 2 * margin) / (numEdgePoints - 1);
            int y = margin;

            // Skip if masked
            if (mask.at<uchar>(y, x) > 0)
                continue;

            float value;
            if (grayImage.type() == CV_8U)
            {
                value = static_cast<float>(grayImage.at<uchar>(y, x));
            }
            else if (grayImage.type() == CV_16U)
            {
                value = static_cast<float>(grayImage.at<ushort>(y, x));
            }
            else if (grayImage.type() == CV_32F)
            {
                value = grayImage.at<float>(y, x);
            }
            else
            {
                continue; // Unsupported type
            }

            points.push_back(cv::Point3f(x, y, value));
        }

        // Bottom edge
        for (int i = 1; i < numEdgePoints - 1; i++)
        {
            int x = margin + i * (grayImage.cols - 2 * margin) / (numEdgePoints - 1);
            int y = grayImage.rows - margin;

            // Skip if masked
            if (mask.at<uchar>(y, x) > 0)
                continue;

            float value;
            if (grayImage.type() == CV_8U)
            {
                value = static_cast<float>(grayImage.at<uchar>(y, x));
            }
            else if (grayImage.type() == CV_16U)
            {
                value = static_cast<float>(grayImage.at<ushort>(y, x));
            }
            else if (grayImage.type() == CV_32F)
            {
                value = grayImage.at<float>(y, x);
            }
            else
            {
                continue; // Unsupported type
            }

            points.push_back(cv::Point3f(x, y, value));
        }

        // Left edge
        for (int i = 1; i < numEdgePoints - 1; i++)
        {
            int x = margin;
            int y = margin + i * (grayImage.rows - 2 * margin) / (numEdgePoints - 1);

            // Skip if masked
            if (mask.at<uchar>(y, x) > 0)
                continue;

            float value;
            if (grayImage.type() == CV_8U)
            {
                value = static_cast<float>(grayImage.at<uchar>(y, x));
            }
            else if (grayImage.type() == CV_16U)
            {
                value = static_cast<float>(grayImage.at<ushort>(y, x));
            }
            else if (grayImage.type() == CV_32F)
            {
                value = grayImage.at<float>(y, x);
            }
            else
            {
                continue; // Unsupported type
            }

            points.push_back(cv::Point3f(x, y, value));
        }

        // Right edge
        for (int i = 1; i < numEdgePoints - 1; i++)
        {
            int x = grayImage.cols - margin;
            int y = margin + i * (grayImage.rows - 2 * margin) / (numEdgePoints - 1);

            // Skip if masked
            if (mask.at<uchar>(y, x) > 0)
                continue;

            float value;
            if (grayImage.type() == CV_8U)
            {
                value = static_cast<float>(grayImage.at<uchar>(y, x));
            }
            else if (grayImage.type() == CV_16U)
            {
                value = static_cast<float>(grayImage.at<ushort>(y, x));
            }
            else if (grayImage.type() == CV_32F)
            {
                value = grayImage.at<float>(y, x);
            }
            else
            {
                continue; // Unsupported type
            }

            points.push_back(cv::Point3f(x, y, value));
        }

        LOG_INFO("After adding edge points: Now have " + std::to_string(points.size()) + " background points");
    }

    // If we still don't have enough points, we'll have to use some masked points
    if (points.size() < minPointsNeeded)
    {
        LOG_WARN("Warning: Still not enough background points. Using some masked points...");

        // Create a distance transform of the mask to find points that are far from objects
        cv::Mat distTransform;
        cv::distanceTransform(255 - mask, distTransform, cv::DIST_L2, 3);

        // Normalize the distance transform to 0-255
        cv::Mat normDist;
        cv::normalize(distTransform, normDist, 0, 255, cv::NORM_MINMAX);
        normDist.convertTo(normDist, CV_8U);

        // Find points with the largest distance from masked objects
        std::vector<std::pair<float, cv::Point>> distPoints;
        for (int y = stepY / 2; y < grayImage.rows; y += stepY)
        {
            for (int x = stepX / 2; x < grayImage.cols; x += stepX)
            {
                // Skip if already added
                bool alreadyAdded = false;
                for (const auto &p : points)
                {
                    if (std::abs(p.x - x) < 1 && std::abs(p.y - y) < 1)
                    {
                        alreadyAdded = true;
                        break;
                    }
                }

                if (alreadyAdded)
                    continue;

                // Get the distance value
                float dist = normDist.at<uchar>(y, x);
                distPoints.push_back(std::make_pair(dist, cv::Point(x, y)));
            }
        }

        // Sort by distance (descending)
        std::sort(distPoints.begin(), distPoints.end(),
                  [](const std::pair<float, cv::Point> &a, const std::pair<float, cv::Point> &b)
        {
            return a.first > b.first;
        });

        // Add points until we have enough
        for (const auto &dp : distPoints)
        {
            if (points.size() >= minPointsNeeded)
                break;

            int x = dp.second.x;
            int y = dp.second.y;

            float value;
            if (grayImage.type() == CV_8U)
            {
                value = static_cast<float>(grayImage.at<uchar>(y, x));
            }
            else if (grayImage.type() == CV_16U)
            {
                value = static_cast<float>(grayImage.at<ushort>(y, x));
            }
            else if (grayImage.type() == CV_32F)
            {
                value = grayImage.at<float>(y, x);
            }
            else
            {
                continue; // Unsupported type
            }

            points.push_back(cv::Point3f(x, y, value));
        }

        LOG_INFO("Final count: " + std::to_string(points.size()) + " background points");
    }

    return points;
}

std::vector<double> BackgroundExtractor::fitPolynomialSurface(const std::vector<cv::Point3f> &points, int width, int height)
{
    if (points.empty())
    {
        throw std::runtime_error("No background points to fit");
    }

    // Number of coefficients depends on polynomial degree
    int ncoeffs;
    switch (m_polynomialDegree)
    {
        case 1:
            ncoeffs = 3; // 1 + x + y
            break;
        case 2:
            ncoeffs = 6; // 1 + x + y + x^2 + xy + y^2
            break;
        case 3:
            ncoeffs = 10; // 1 + x + y + x^2 + xy + y^2 + x^3 + x^2y + xy^2 + y^3
            break;
        case 4:
            ncoeffs = 15; // 1 + x + y + x^2 + xy + y^2 + x^3 + x^2y + xy^2 + y^3 + x^4 + x^3y + x^2y^2 + xy^3 + y^4
            break;
        default:
            ncoeffs = 6; // Default to degree 2
            break;
    }

    // Normalize coordinates to [0, 1] range
    float normX = static_cast<float>(width - 1);
    float normY = static_cast<float>(height - 1);

    // Set up GSL workspace
    gsl_matrix *X = gsl_matrix_alloc(points.size(), ncoeffs);
    gsl_vector *Y = gsl_vector_alloc(points.size());
    gsl_vector *c = gsl_vector_alloc(ncoeffs);
    gsl_matrix *cov = gsl_matrix_alloc(ncoeffs, ncoeffs);

    // Fill design matrix and data vector
    for (size_t i = 0; i < points.size(); i++)
    {
        double x = points[i].x / normX;
        double y = points[i].y / normY;
        double z = points[i].z;

        // Fill row of design matrix based on polynomial degree
        int col = 0;
        gsl_matrix_set(X, i, col++, 1.0); // Constant term

        // Linear terms
        gsl_matrix_set(X, i, col++, x);
        gsl_matrix_set(X, i, col++, y);

        if (m_polynomialDegree >= 2)
        {
            // Quadratic terms
            gsl_matrix_set(X, i, col++, x * x);
            gsl_matrix_set(X, i, col++, x * y);
            gsl_matrix_set(X, i, col++, y * y);
        }

        if (m_polynomialDegree >= 3)
        {
            // Cubic terms
            gsl_matrix_set(X, i, col++, x * x * x);
            gsl_matrix_set(X, i, col++, x * x * y);
            gsl_matrix_set(X, i, col++, x * y * y);
            gsl_matrix_set(X, i, col++, y * y * y);
        }

        if (m_polynomialDegree >= 4)
        {
            // Quartic terms
            gsl_matrix_set(X, i, col++, x * x * x * x);
            gsl_matrix_set(X, i, col++, x * x * x * y);
            gsl_matrix_set(X, i, col++, x * x * y * y);
            gsl_matrix_set(X, i, col++, x * y * y * y);
            gsl_matrix_set(X, i, col++, y * y * y * y);
        }

        // Set data value
        gsl_vector_set(Y, i, z);
    }

    // Perform the fit
    double chisq;
    gsl_multifit_linear_workspace *work = gsl_multifit_linear_alloc(points.size(), ncoeffs);
    gsl_multifit_linear(X, Y, c, cov, &chisq, work);
    gsl_multifit_linear_free(work);

    // Extract coefficients
    std::vector<double> coeffs(ncoeffs);
    for (int i = 0; i < ncoeffs; i++)
    {
        coeffs[i] = gsl_vector_get(c, i);
    }

    // Clean up
    gsl_matrix_free(X);
    gsl_vector_free(Y);
    gsl_vector_free(c);
    gsl_matrix_free(cov);

    return coeffs;
}

cv::Mat BackgroundExtractor::generateBackgroundModelFromCoefficients(const std::vector<double> &coeffs, int width,
        int height)
{
    cv::Mat model(height, width, CV_32F);

    // Normalize coordinates to [0, 1] range
    float normX = static_cast<float>(width - 1);
    float normY = static_cast<float>(height - 1);

    // Generate background model with OpenMP parallelization
    #pragma omp parallel for collapse(2) schedule(dynamic, 16)
    for (int y = 0; y < height; y++)
    {
        for (int x = 0; x < width; x++)
        {
            double nx = x / normX;
            double ny = y / normY;

            // Evaluate polynomial
            double value = coeffs[0]; // Constant term

            int idx = 1;
            // Linear terms
            value += coeffs[idx++] * nx;
            value += coeffs[idx++] * ny;

            if (m_polynomialDegree >= 2 && idx + 2 < coeffs.size())
            {
                // Quadratic terms
                value += coeffs[idx++] * nx * nx;
                value += coeffs[idx++] * nx * ny;
                value += coeffs[idx++] * ny * ny;
            }

            if (m_polynomialDegree >= 3 && idx + 3 < coeffs.size())
            {
                // Cubic terms
                value += coeffs[idx++] * nx * nx * nx;
                value += coeffs[idx++] * nx * nx * ny;
                value += coeffs[idx++] * nx * ny * ny;
                value += coeffs[idx++] * ny * ny * ny;
            }

            if (m_polynomialDegree >= 4 && idx + 4 < coeffs.size())
            {
                // Quartic terms
                value += coeffs[idx++] * nx * nx * nx * nx;
                value += coeffs[idx++] * nx * nx * nx * ny;
                value += coeffs[idx++] * nx * nx * ny * ny;
                value += coeffs[idx++] * nx * ny * ny * ny;
                value += coeffs[idx++] * ny * ny * ny * ny;
            }

            model.at<float>(y, x) = static_cast<float>(value);
        }
    }

    return model;
}

cv::Mat BackgroundExtractor::applyBackgroundModel(const cv::Mat &image, const cv::Mat &backgroundModel)
{
    // Handle multi-channel images
    if (image.channels() > 1)
    {
        // Split the image into channels
        std::vector<cv::Mat> channels;
        cv::split(image, channels);

        // Process each channel separately
        std::vector<cv::Mat> processedChannels(channels.size());

        // Use OpenMP to parallelize channel processing
        #pragma omp parallel for schedule(dynamic, 1)
        for (int c = 0; c < static_cast<int>(channels.size()); ++c)
        {
            // Convert to float for processing
            cv::Mat floatChannel;
            channels[c].convertTo(floatChannel, CV_32F);

            // Calculate original mean and standard deviation
            cv::Scalar originalMean, originalStdDev;
            cv::meanStdDev(floatChannel, originalMean, originalStdDev);

            // Subtract background model with a moderate scaling factor
            // The scaling factor is based on the smoothness factor
            float subtractionFactor = 1.0f + (1.0f - m_smoothnessFactor) * 0.3f;
            cv::Mat result = floatChannel - (backgroundModel * subtractionFactor);

            // Add back a portion of the mean background level to preserve foreground details
            cv::Scalar meanBg = cv::mean(backgroundModel);
            float addBackFactor = m_smoothnessFactor * 0.7f;
            result += meanBg[0] * addBackFactor;

            // Ensure no negative values
            cv::threshold(result, result, 0, 0, cv::THRESH_TOZERO);

            // Calculate new mean and standard deviation
            cv::Scalar newMean, newStdDev;
            cv::meanStdDev(result, newMean, newStdDev);

            // Scale the result to match the original standard deviation
            if (newStdDev[0] > 0 && originalStdDev[0] > 0)
            {
                double scaleFactor = originalStdDev[0] / newStdDev[0];
                // Scale by a factor that's a blend between 1.0 and the calculated scale factor
                // This prevents extreme scaling while still improving the quality assessment
                double blendedScaleFactor = 1.0 + (scaleFactor - 1.0) * m_smoothnessFactor;
                // Use a higher mean factor to preserve foreground details
                double meanFactor = 0.9;
                result = (result - newMean[0]) * blendedScaleFactor + originalMean[0] * meanFactor;
            }

            // Ensure no negative values after scaling
            cv::threshold(result, result, 0, 0, cv::THRESH_TOZERO);

            // Apply mild contrast enhancement to make stars and galaxies stand out more
            double min, max;
            cv::minMaxLoc(result, &min, &max);
            if (max > 0)
            {
                // Apply a gentle power function to enhance contrast
                float contrastFactor = 1.0f - m_smoothnessFactor * 0.3f;
                // Ensure we're working with CV_32F for the pow operation
                cv::Mat normalizedResult = result / max;
                cv::Mat powResult;
                cv::pow(normalizedResult, contrastFactor, powResult);
                result = powResult * max;
            }

            // Convert back to original type
            cv::Mat outputChannel;
            if (channels[c].type() == CV_8U)
            {
                result.convertTo(outputChannel, CV_8U);
            }
            else if (channels[c].type() == CV_16U)
            {
                result.convertTo(outputChannel, CV_16U);
            }
            else
            {
                outputChannel = result;
            }

            processedChannels[c] = outputChannel;
        }

        // Merge the processed channels back into an image
        cv::Mat output;
        cv::merge(processedChannels, output);
        return output;
    }
    else
    {
        // Single channel image
        // Convert to float for processing
        cv::Mat floatImage;
        image.convertTo(floatImage, CV_32F);

        // Calculate original mean and standard deviation
        cv::Scalar originalMean, originalStdDev;
        cv::meanStdDev(floatImage, originalMean, originalStdDev);

        // Subtract background model with a scaling factor to make it more aggressive
        // The scaling factor is based on the smoothness factor - lower smoothness means more aggressive subtraction
        float subtractionFactor = 1.0f + (1.0f - m_smoothnessFactor) * 0.5f;
        cv::Mat result = floatImage - (backgroundModel * subtractionFactor);

        // Add back a small portion of the mean background level
        cv::Scalar meanBg = cv::mean(backgroundModel);
        // Use a smaller factor for adding back the mean to make background removal more aggressive
        float addBackFactor = m_smoothnessFactor * 0.5f;
        result += meanBg[0] * addBackFactor;

        // Ensure no negative values
        cv::threshold(result, result, 0, 0, cv::THRESH_TOZERO);

        // Calculate new mean and standard deviation
        cv::Scalar newMean, newStdDev;
        cv::meanStdDev(result, newMean, newStdDev);

        // Scale the result to match the original standard deviation
        if (newStdDev[0] > 0 && originalStdDev[0] > 0)
        {
            double scaleFactor = originalStdDev[0] / newStdDev[0];
            // Scale by a factor that's a blend between 1.0 and the calculated scale factor
            // This prevents extreme scaling while still improving the quality assessment
            double blendedScaleFactor = 1.0 + (scaleFactor - 1.0) * m_smoothnessFactor;
            result = (result - newMean[0]) * blendedScaleFactor + originalMean[0] * 0.8f; // Reduce the mean to darken the background
        }

        // Ensure no negative values after scaling
        cv::threshold(result, result, 0, 0, cv::THRESH_TOZERO);

        // Apply additional contrast enhancement to make stars and galaxies stand out more
        double min, max;
        cv::minMaxLoc(result, &min, &max);
        if (max > 0)
        {
            // Apply a power function to enhance contrast
            // Lower smoothness factor means more contrast enhancement
            float contrastFactor = 1.0f - m_smoothnessFactor * 0.5f;
            // Ensure we're working with CV_32F for the pow operation
            cv::Mat normalizedResult = result / max;
            cv::Mat powResult;
            cv::pow(normalizedResult, contrastFactor, powResult);
            result = powResult * max;
        }

        // Convert back to original type
        cv::Mat output;
        if (image.type() == CV_8U)
        {
            result.convertTo(output, CV_8U);
        }
        else if (image.type() == CV_16U)
        {
            result.convertTo(output, CV_16U);
        }
        else
        {
            output = result;
        }

        return output;
    }
}
