#include "nebulosity_detector.h"

NebulosityDetector::NebulosityDetector()
    : m_detectionThreshold(0.4f)  // Lower default threshold to better detect nebulosity with the new algorithm
{
}

void NebulosityDetector::setDetectionThreshold(float threshold)
{
    m_detectionThreshold = std::max(0.0f, std::min(1.0f, threshold));
}

float NebulosityDetector::getDetectionThreshold() const
{
    return m_detectionThreshold;
}

const NebulosityDetector::NebulosityMetrics &NebulosityDetector::getLastMetrics() const
{
    return m_lastMetrics;
}

bool NebulosityDetector::detectExtensiveNebulosity(const cv::Mat &image, float &nebulosityScore)
{
    // Call the detailed version and just return the score
    NebulosityMetrics metrics;
    bool result = detectExtensiveNebulosity(image, metrics);
    nebulosityScore = metrics.finalScore;
    return result;
}

bool NebulosityDetector::detectExtensiveNebulosity(const cv::Mat &image, NebulosityMetrics &metrics)
{
    // Convert to grayscale if needed
    cv::Mat gray;
    if (image.channels() > 1)
        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;
    }

    // 1. Calculate image statistics
    cv::Scalar mean, stddev;
    cv::meanStdDev(gray, mean, stddev);

    // 2. Apply CLAHE to enhance local contrast
    // Use stronger clip limit and larger tile size for better nebulosity enhancement
    cv::Ptr<cv::CLAHE> clahe = cv::createCLAHE(3.0, cv::Size(16, 16));
    cv::Mat enhanced;
    clahe->apply(gray, enhanced);

    // 3. Calculate local standard deviation (texture measure)
    cv::Mat localStdDev;
    cv::Scalar localMean, localStdDevMean;

    // Create a kernel for local statistics - use larger kernel for nebulosity
    cv::Mat kernel = cv::getStructuringElement(cv::MORPH_RECT, cv::Size(25, 25));

    // Calculate local standard deviation with larger window to better capture nebulosity texture
    cv::Mat localVar;
    // Convert enhanced to float for multiplication
    cv::Mat enhancedFloat;
    enhanced.convertTo(enhancedFloat, CV_32F);
    cv::blur(enhancedFloat.mul(enhancedFloat), localVar, cv::Size(25, 25));
    cv::Mat localMeanMat;
    cv::blur(enhancedFloat, localMeanMat, cv::Size(25, 25));
    localVar -= localMeanMat.mul(localMeanMat);
    cv::sqrt(localVar, localStdDev);

    // Calculate mean of local standard deviation
    cv::meanStdDev(localStdDev, localStdDevMean, cv::Scalar());

    // 4. Apply adaptive thresholding to find regions of interest
    // Use a larger block size and more aggressive constant subtraction for nebulosity
    cv::Mat binary;
    cv::adaptiveThreshold(enhanced, binary, 255, cv::ADAPTIVE_THRESH_GAUSSIAN_C,
                          cv::THRESH_BINARY, 75, -5);  // Increased block size, less negative constant

    // 5. Calculate percentage of pixels above threshold
    double pixelPercentage = (cv::countNonZero(binary) * 100.0) / (binary.rows * binary.cols);

    // 6. Calculate texture uniformity (lower means more texture variation, typical of nebulosity)
    double textureUniformity = localStdDevMean[0] / mean[0];

    // 7. Calculate gradient magnitude to analyze edge characteristics
    // Nebulosity typically has smoother gradients compared to distinct objects
    cv::Mat gradX, gradY, gradMag;
    cv::Sobel(enhanced, gradX, CV_32F, 1, 0, 3);
    cv::Sobel(enhanced, gradY, CV_32F, 0, 1, 3);
    cv::magnitude(gradX, gradY, gradMag);

    // Normalize gradient magnitude for analysis
    cv::Mat normGradMag;
    cv::normalize(gradMag, normGradMag, 0, 1, cv::NORM_MINMAX);

    // Calculate gradient statistics
    cv::Scalar gradMean, gradStdDev;
    cv::meanStdDev(normGradMag, gradMean, gradStdDev);

    // 8. Calculate diffuse area percentage
    // Nebulosity typically has large areas with low gradients
    cv::Mat lowGradientMask;
    double gradientThreshold = 0.2; // Threshold for low gradient areas
    cv::threshold(normGradMag, lowGradientMask, gradientThreshold, 1, cv::THRESH_BINARY_INV);
    lowGradientMask.convertTo(lowGradientMask, CV_8U);

    // Calculate percentage of low gradient areas
    double lowGradientPercentage = (cv::countNonZero(lowGradientMask) * 100.0) /
                                   (lowGradientMask.rows * lowGradientMask.cols);

    // 9. Analyze spatial distribution using image quadrants
    // Divide the image into a grid and analyze each cell
    const int gridSize = 4; // 4x4 grid
    std::vector<double> cellBrightness(gridSize * gridSize, 0.0);
    std::vector<double> cellGradients(gridSize * gridSize, 0.0);

    int cellWidth = gray.cols / gridSize;
    int cellHeight = gray.rows / gridSize;

    // Calculate brightness and gradient statistics for each cell
    for (int y = 0; y < gridSize; y++)
    {
        for (int x = 0; x < gridSize; x++)
        {
            cv::Rect cellRect(x * cellWidth, y * cellHeight, cellWidth, cellHeight);

            // Extract cell regions
            cv::Mat cellGray = gray(cellRect);
            cv::Mat cellGradient = normGradMag(cellRect);
            cv::Mat cellLowGradient = lowGradientMask(cellRect);

            // Calculate mean brightness for this cell
            cv::Scalar cellMean = cv::mean(cellGray);
            cellBrightness[y * gridSize + x] = cellMean[0];

            // Calculate percentage of low gradient pixels in this cell
            double cellLowGradientPct = (cv::countNonZero(cellLowGradient) * 100.0) /
                                        (cellLowGradient.rows * cellLowGradient.cols);
            cellGradients[y * gridSize + x] = cellLowGradientPct;
        }
    }

    // Calculate statistics of cell measurements
    double sumBrightness = 0.0, sumGradients = 0.0;
    double sumSqBrightness = 0.0, sumSqGradients = 0.0;

    for (int i = 0; i < gridSize * gridSize; i++)
    {
        sumBrightness += cellBrightness[i];
        sumSqBrightness += cellBrightness[i] * cellBrightness[i];

        sumGradients += cellGradients[i];
        sumSqGradients += cellGradients[i] * cellGradients[i];
    }

    double meanBrightness = sumBrightness / (gridSize * gridSize);
    double meanGradients = sumGradients / (gridSize * gridSize);

    double varBrightness = (sumSqBrightness / (gridSize * gridSize)) - (meanBrightness * meanBrightness);
    double varGradients = (sumSqGradients / (gridSize * gridSize)) - (meanGradients * meanGradients);

    double stdDevBrightness = std::sqrt(std::max(0.0, varBrightness));
    double stdDevGradients = std::sqrt(std::max(0.0, varGradients));

    // Calculate coefficient of variation for brightness and gradients
    double cvBrightness = stdDevBrightness / (meanBrightness + 1e-5);
    double cvGradients = stdDevGradients / (meanGradients + 1e-5);

    // 10. Calculate spatial uniformity factor
    // Lower values indicate more uniform distribution across the image (typical of nebulosity)
    // Higher values indicate concentrated brightness (typical of galaxies, stars, etc.)
    double spatialUniformityFactor = (cvBrightness + cvGradients) / 2.0;

    // 11. Calculate nebulosity score based on multiple factors specific to nebulosity
    double stddevFactor = stddev[0] / mean[0];

    // Calculate enhanced factors for nebulosity detection
    double enhancedPixelPercentage = std::pow(pixelPercentage / 100.0, 0.7);
    double enhancedTextureUniformity = std::pow(1.0 - std::min(1.0, textureUniformity), 1.5);
    double enhancedStddevFactor = std::pow(std::min(1.0, stddevFactor * 4.0), 0.3);

    // Calculate diffuse structure factor (higher for nebulosity)
    double diffuseFactor = lowGradientPercentage / 100.0;

    // Calculate spatial distribution factor (higher for nebulosity)
    // Invert the uniformity factor so higher values indicate more uniform distribution
    double spatialFactor = 1.0 / (1.0 + spatialUniformityFactor);

    // Calculate object concentration factor
    // Lower values for concentrated objects (galaxies, stars), higher for diffuse nebulosity
    double concentrationFactor = 0.0;

    // Count bright cells in the center vs periphery
    int centerBrightCells = 0;
    int peripheryBrightCells = 0;
    double brightnessCutoff = meanBrightness * 1.2; // 20% above mean

    for (int y = 0; y < gridSize; y++)
    {
        for (int x = 0; x < gridSize; x++)
        {
            bool isCenter = (x >= gridSize / 4 && x < 3 * gridSize / 4 &&
                             y >= gridSize / 4 && y < 3 * gridSize / 4);

            if (cellBrightness[y * gridSize + x] > brightnessCutoff)
            {
                if (isCenter)
                    centerBrightCells++;
                else
                    peripheryBrightCells++;
            }
        }
    }

    // Calculate ratio of center to periphery brightness
    double centerPeripheryRatio = 0.0;
    if (peripheryBrightCells > 0)
        centerPeripheryRatio = static_cast<double>(centerBrightCells) / peripheryBrightCells;
    else if (centerBrightCells > 0)
        centerPeripheryRatio = 10.0; // High value for center-only brightness

    // Convert to a factor where lower values indicate concentrated objects
    concentrationFactor = 1.0 / (1.0 + centerPeripheryRatio);

    // Give more weight to concentration factor for distinguishing galaxies from nebulae
    // Square the concentration factor to make it more significant
    double weightedConcentrationFactor = std::pow(concentrationFactor, 2.0);

    // Calculate base nebulosity score with focus on diffuse, extended structures
    double baseScore = enhancedPixelPercentage * enhancedTextureUniformity *
                       enhancedStddevFactor * diffuseFactor * spatialFactor * weightedConcentrationFactor;

    // Apply additional scoring factors for different types of nebulosity

    // Factor 1: Boost for images with high pixel percentage regardless of standard deviation
    // This helps detect extensive nebulosity that might have varying brightness levels
    double pixelBoost = 0.0;
    if (pixelPercentage > 15.0)
    {
        // Linear boost based on pixel percentage above 15%
        pixelBoost = (pixelPercentage - 15.0) / 85.0 * 0.3;

        // Reduce boost for concentrated objects (likely galaxies)
        pixelBoost *= concentrationFactor;
    }

    // Factor 2: Boost for images with moderate texture uniformity
    // This helps detect nebulosity with complex texture patterns
    double textureBoost = 0.0;
    if (textureUniformity < 0.8) // Increased threshold to handle more uniform textures
    {
        // Inverse linear boost based on texture uniformity below 0.8
        textureBoost = (0.8 - textureUniformity) / 0.8 * 0.3;

        // Reduce boost for concentrated objects (likely galaxies)
        textureBoost *= concentrationFactor;
    }

    // Factor 3: Special case for HA (hydrogen alpha) images which often have
    // diffuse nebulosity with low contrast but distinctive texture patterns
    double haBoost = 0.0;
    if (pixelPercentage > 20.0 && stddevFactor < 0.3) // Removed texture uniformity constraint
    {
        // Apply HA boost but scale it by concentration factor
        // This will reduce the boost for galaxies which have low concentration factor
        haBoost = 0.3 * concentrationFactor;
    }

    // Factor 4: Special boost for images with high pixel percentage and diffuse structure
    // This helps detect extensive nebulosity that might have more uniform texture
    double extensiveBoost = 0.0;
    if (pixelPercentage > 25.0) // Removed texture uniformity constraint
    {
        // Apply extensive boost but scale it by concentration factor
        // This will reduce the boost for galaxies which have low concentration factor
        extensiveBoost = 0.25 * concentrationFactor;
    }

    // Apply all boosts to the base score
    float nebulosityScore = baseScore + pixelBoost + textureBoost + haBoost + extensiveBoost;

    // Ensure score is in valid range (0.0-1.0)
    double tempScore = static_cast<double>(nebulosityScore);
    tempScore = std::min(1.0, std::max(0.0, tempScore));
    nebulosityScore = static_cast<float>(tempScore);

    // Store all metrics in the output parameter
    metrics.meanPixelValue = mean[0];
    metrics.standardDeviation = stddev[0];
    metrics.pixelPercentage = pixelPercentage;
    metrics.textureUniformity = textureUniformity;
    metrics.lowGradientPercentage = lowGradientPercentage;
    metrics.diffuseFactor = diffuseFactor;
    metrics.spatialUniformityFactor = spatialUniformityFactor;
    metrics.spatialFactor = spatialFactor;
    metrics.centerPeripheryRatio = centerPeripheryRatio;
    metrics.concentrationFactor = concentrationFactor;
    metrics.enhancedPixelPercentage = enhancedPixelPercentage;
    metrics.enhancedTextureUniformity = enhancedTextureUniformity;
    metrics.enhancedStddevFactor = enhancedStddevFactor;
    metrics.baseScore = baseScore;
    metrics.pixelBoost = pixelBoost;
    metrics.textureBoost = textureBoost;
    metrics.haBoost = haBoost;
    metrics.extensiveBoost = extensiveBoost;
    metrics.finalScore = nebulosityScore;

    // Also store in the last metrics
    m_lastMetrics = metrics;

    // Log the detection results
    LOG_INFO("Nebulosity detection results:");
    LOG_INFO("  Mean pixel value: " + std::to_string(mean[0]));
    LOG_INFO("  Standard deviation: " + std::to_string(stddev[0]));
    LOG_INFO("  Pixel percentage above threshold: " + std::to_string(pixelPercentage) + "%");
    LOG_INFO("  Texture uniformity: " + std::to_string(textureUniformity));
    LOG_INFO("  Nebulosity score: " + std::to_string(nebulosityScore));

    // Add additional logging for the enhanced factors and nebulosity metrics
    LOG_INFO("  Enhanced pixel percentage factor: " + std::to_string(enhancedPixelPercentage));
    LOG_INFO("  Enhanced texture uniformity factor: " + std::to_string(enhancedTextureUniformity));
    LOG_INFO("  Enhanced standard deviation factor: " + std::to_string(enhancedStddevFactor));
    LOG_INFO("  Low gradient percentage: " + std::to_string(lowGradientPercentage) + "%");
    LOG_INFO("  Diffuse factor: " + std::to_string(diffuseFactor));
    LOG_INFO("  Spatial uniformity factor: " + std::to_string(spatialUniformityFactor));
    LOG_INFO("  Spatial distribution factor: " + std::to_string(spatialFactor));
    LOG_INFO("  Center/periphery ratio: " + std::to_string(centerPeripheryRatio));
    LOG_INFO("  Concentration factor: " + std::to_string(concentrationFactor));
    LOG_INFO("  Base score: " + std::to_string(baseScore));
    LOG_INFO("  Pixel boost: " + std::to_string(pixelBoost));
    LOG_INFO("  Texture boost: " + std::to_string(textureBoost));
    LOG_INFO("  HA boost: " + std::to_string(haBoost));
    LOG_INFO("  Extensive boost: " + std::to_string(extensiveBoost));

    // Consider extensive nebulosity if score is above threshold
    bool result = nebulosityScore > m_detectionThreshold;
    LOG_INFO("Extensive nebulosity detected: " + std::string(result ? "YES" : "NO"));

    return result;
}

cv::Mat NebulosityDetector::generateVisualization(const cv::Mat &image, float nebulosityScore)
{
    // Process the image to get the binary mask and other intermediate results
    cv::Mat gray, enhanced, binary, localStdDev;

    // Convert to grayscale if needed
    if (image.channels() > 1)
        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;
    }

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

    // Apply adaptive thresholding to find regions of interest
    cv::adaptiveThreshold(enhanced, binary, 255, cv::ADAPTIVE_THRESH_GAUSSIAN_C,
                          cv::THRESH_BINARY, 75, -5);

    // Calculate local standard deviation
    cv::Mat enhancedFloat, localVar, localMeanMat;
    enhanced.convertTo(enhancedFloat, CV_32F);
    cv::blur(enhancedFloat.mul(enhancedFloat), localVar, cv::Size(25, 25));
    cv::blur(enhancedFloat, localMeanMat, cv::Size(25, 25));
    localVar -= localMeanMat.mul(localMeanMat);
    cv::sqrt(localVar, localStdDev);

    // Normalize localStdDev for visualization
    cv::normalize(localStdDev, localStdDev, 0, 255, cv::NORM_MINMAX, CV_8U);

    // Create a copy of the image for visualization
    cv::Mat visualization;

    // If the image is grayscale, convert to color for better visualization
    if (image.channels() == 1)
    {
        cv::cvtColor(image, visualization, cv::COLOR_GRAY2BGR);
    }
    else
    {
        visualization = image.clone();
    }

    // Calculate gradient magnitude for visualization
    cv::Mat gradX, gradY, gradMag;
    cv::Sobel(enhanced, gradX, CV_32F, 1, 0, 3);
    cv::Sobel(enhanced, gradY, CV_32F, 0, 1, 3);
    cv::magnitude(gradX, gradY, gradMag);

    // Normalize gradient magnitude for visualization
    cv::Mat normGradMag;
    cv::normalize(gradMag, normGradMag, 0, 255, cv::NORM_MINMAX, CV_8U);

    // Calculate low gradient mask for visualization
    cv::Mat lowGradientMask;
    double gradientThreshold = 0.2;
    cv::Mat normGradMagFloat;
    cv::normalize(gradMag, normGradMagFloat, 0, 1, cv::NORM_MINMAX, CV_32F);
    cv::threshold(normGradMagFloat, lowGradientMask, gradientThreshold, 1, cv::THRESH_BINARY_INV);
    lowGradientMask.convertTo(lowGradientMask, CV_8U, 255);

    // Create a composite visualization with multiple views
    int width = visualization.cols;
    int height = visualization.rows;
    cv::Mat composite(height * 2, width * 2, CV_8UC3);

    // Original image (top-left)
    cv::Mat roi = composite(cv::Rect(0, 0, width, height));
    if (visualization.channels() == 3)
        visualization.copyTo(roi);
    else
        cv::cvtColor(visualization, roi, cv::COLOR_GRAY2BGR);

    // Enhanced image (top-right)
    roi = composite(cv::Rect(width, 0, width, height));
    cv::cvtColor(enhanced, roi, cv::COLOR_GRAY2BGR);

    // Gradient magnitude (bottom-left)
    roi = composite(cv::Rect(0, height, width, height));
    cv::cvtColor(normGradMag, roi, cv::COLOR_GRAY2BGR);

    // Low gradient mask (bottom-right)
    roi = composite(cv::Rect(width, height, width, height));
    cv::cvtColor(lowGradientMask, roi, cv::COLOR_GRAY2BGR);

    // Add labels to each quadrant
    cv::putText(composite, "Original Image", cv::Point(20, 30),
                cv::FONT_HERSHEY_SIMPLEX, 1.0, cv::Scalar(255, 255, 255), 2);
    cv::putText(composite, "Enhanced Image", cv::Point(width + 20, 30),
                cv::FONT_HERSHEY_SIMPLEX, 1.0, cv::Scalar(255, 255, 255), 2);
    cv::putText(composite, "Gradient Magnitude", cv::Point(20, height + 30),
                cv::FONT_HERSHEY_SIMPLEX, 1.0, cv::Scalar(255, 255, 255), 2);
    cv::putText(composite, "Low Gradient Mask", cv::Point(width + 20, height + 30),
                cv::FONT_HERSHEY_SIMPLEX, 1.0, cv::Scalar(255, 255, 255), 2);

    // Add text with the nebulosity score
    std::string scoreText = "Nebulosity Score: " + std::to_string(nebulosityScore);
    cv::putText(composite, scoreText, cv::Point(20, height - 50),
                cv::FONT_HERSHEY_SIMPLEX, 1.0, cv::Scalar(255, 255, 255), 2);

    // Add text indicating if extensive nebulosity was detected
    std::string detectionText = "Extensive Nebulosity: " +
                                std::string(nebulosityScore > m_detectionThreshold ? "YES" : "NO");
    cv::putText(composite, detectionText, cv::Point(20, height - 20),
                cv::FONT_HERSHEY_SIMPLEX, 1.0,
                nebulosityScore > m_detectionThreshold ? cv::Scalar(0, 255, 0) : cv::Scalar(0, 0, 255), 2);

    // Create a colored border based on detection result
    int borderSize = 10;
    cv::Scalar borderColor = nebulosityScore > m_detectionThreshold ?
                             cv::Scalar(0, 255, 0) : cv::Scalar(0, 0, 255);

    cv::Mat bordered;
    cv::copyMakeBorder(composite, bordered, borderSize, borderSize, borderSize, borderSize,
                       cv::BORDER_CONSTANT, borderColor);

    return bordered;
}
