#include "star_detector.h"
#include <stellarsolver.h>
#include <QSharedPointer>
#include "../fits/fitsdata.h"
#include "../io/logger.h"
#include <algorithm>
#include <cmath>
#include <map>
#include <set>

// Helper struct for triangle matching
struct Triangle
{
    int vertex1, vertex2, vertex3;  // Indices of stars forming the triangle
    double side1, side2, side3;     // Lengths of the sides
    double ratio1, ratio2, ratio3;  // Ratios of sides (for scale invariance)
    double area;                    // Area of the triangle
    double perimeter;               // Perimeter of the triangle
    double angle1, angle2, angle3;  // Angles of the triangle (in radians)
    double orientation;             // Orientation of the triangle (positive for counterclockwise, negative for clockwise)
    double centroidX, centroidY;    // Centroid of the triangle

    // Constructor with star coordinates
    Triangle(int v1, int v2, int v3, float x1, float y1, float x2, float y2, float x3, float y3)
        : vertex1(v1), vertex2(v2), vertex3(v3)
    {
        // Calculate side lengths
        side1 = std::sqrt(std::pow(x2 - x1, 2) + std::pow(y2 - y1, 2));
        side2 = std::sqrt(std::pow(x3 - x2, 2) + std::pow(y3 - y2, 2));
        side3 = std::sqrt(std::pow(x1 - x3, 2) + std::pow(y1 - y3, 2));

        // Sort sides to ensure consistent ordering
        if (side1 > side2) std::swap(side1, side2);
        if (side2 > side3) std::swap(side2, side3);
        if (side1 > side2) std::swap(side1, side2);

        // Calculate ratios for scale invariance
        ratio1 = side1 / side3;
        ratio2 = side2 / side3;
        ratio3 = 1.0; // side3/side3

        // Calculate perimeter
        perimeter = side1 + side2 + side3;

        // Calculate area using Heron's formula
        double s = perimeter / 2.0;
        area = std::sqrt(s * (s - side1) * (s - side2) * (s - side3));

        // Calculate angles using the law of cosines
        angle1 = std::acos((side2 * side2 + side3 * side3 - side1 * side1) / (2 * side2 * side3));
        angle2 = std::acos((side1 * side1 + side3 * side3 - side2 * side2) / (2 * side1 * side3));
        angle3 = std::acos((side1 * side1 + side2 * side2 - side3 * side3) / (2 * side1 * side2));

        // Sort angles to ensure consistent ordering
        if (angle1 > angle2) std::swap(angle1, angle2);
        if (angle2 > angle3) std::swap(angle2, angle3);
        if (angle1 > angle2) std::swap(angle1, angle2);

        // Calculate orientation using the cross product
        // For triangle ABC, calculate (B-A) × (C-A)
        // This will be positive for counterclockwise orientation, negative for clockwise
        orientation = (x2 - x1) * (y3 - y1) - (y2 - y1) * (x3 - x1);

        // Calculate centroid
        centroidX = (x1 + x2 + x3) / 3.0;
        centroidY = (y1 + y2 + y3) / 3.0;
    }

    // Check if two triangles are similar based on multiple criteria
    bool isSimilar(const Triangle &other, double tolerance = 0.1) const
    {
        // Check side ratios (scale invariant)
        bool sideRatiosSimilar = std::abs(ratio1 - other.ratio1) < tolerance &&
                                 std::abs(ratio2 - other.ratio2) < tolerance;

        // Check angles (shape invariant)
        bool anglesSimilar = std::abs(angle1 - other.angle1) < 0.1 &&
                             std::abs(angle2 - other.angle2) < 0.1 &&
                             std::abs(angle3 - other.angle3) < 0.1;

        // Check area ratio relative to perimeter (shape invariant)
        double areaPerimeterRatio = area / (perimeter * perimeter);
        double otherAreaPerimeterRatio = other.area / (other.perimeter * other.perimeter);
        bool areaRatioSimilar = std::abs(areaPerimeterRatio - otherAreaPerimeterRatio) < 0.05;

        // A triangle is similar if it passes all tests
        return sideRatiosSimilar && anglesSimilar && areaRatioSimilar;
    }

    // Calculate a similarity score between two triangles (0.0 to 1.0)
    double similarityScore(const Triangle &other) const
    {
        // Calculate similarity based on side ratios
        double sideRatioScore = 1.0 - (std::abs(ratio1 - other.ratio1) + std::abs(ratio2 - other.ratio2)) / 2.0;

        // Calculate similarity based on angles
        double angleScore = 1.0 - (std::abs(angle1 - other.angle1) +
                                   std::abs(angle2 - other.angle2) +
                                   std::abs(angle3 - other.angle3)) / (3.0 * M_PI);

        // Calculate similarity based on area/perimeter ratio
        double areaPerimeterRatio = area / (perimeter * perimeter);
        double otherAreaPerimeterRatio = other.area / (other.perimeter * other.perimeter);
        double areaRatioScore = 1.0 - std::abs(areaPerimeterRatio - otherAreaPerimeterRatio) / 0.1;

        // Combine scores (weighted average)
        return 0.4 * sideRatioScore + 0.4 * angleScore + 0.2 * areaRatioScore;
    }
};

StarDetector::StarDetector()
    : m_maxStars(20)
{
}

StarDetector::~StarDetector() = default;

void StarDetector::setMaxStars(int maxStars)
{
    m_maxStars = maxStars;
}

std::string StarDetector::getLastError() const
{
    return m_lastError;
}

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

    // Create StellarSolver instance
    std::unique_ptr<StellarSolver> solver(new StellarSolver());

    // Configure StellarSolver parameters for star extraction only
    solver->setProperty("ProcessType", SSolver::EXTRACT);

    // Set up parameters for star extraction
    SSolver::Parameters params;
    params.minarea = 5;
    params.deblend_thresh = 32;
    params.deblend_contrast = 0.005;
    params.kron_fact = 2.5;
    params.subpix = 5;

    solver->setParameters(params);

    // Create FITSData object for the image
    QSharedPointer<FITSData> fits(new FITSData());

    // Convert OpenCV image to FITSData
    if (!fits->loadFromOpenCV(grayImage, "image.fits"))
    {
        m_lastError = "Failed to load image into FITSData";
        LOG_ERROR(m_lastError);
        throw std::runtime_error(m_lastError);
    }

    // Load the image buffer into StellarSolver
    solver->loadNewImageBuffer(fits->getStatistics(), fits->getImageBuffer());

    // Extract stars
    solver->extract();

    // Wait for extraction to complete
    if (!solver->extractionDone() || solver->failed())
    {
        m_lastError = "StellarSolver failed to extract stars";
        LOG_ERROR(m_lastError);
        throw std::runtime_error(m_lastError);
    }

    // Get the extracted stars
    const QList<FITSImage::Star> &stars = solver->getStarList();

    // Check if we have enough stars
    if (stars.size() < 3)
    {
        m_lastError = "Not enough stars detected";
        LOG_ERROR(m_lastError);
        throw std::runtime_error(m_lastError);
    }

    LOG_INFO("StellarSolver detected " + std::to_string(stars.size()) + " stars");

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

    // Draw detected stars on mask
    for (const auto &star : stars)
    {
        // Convert to pixels
        int px = static_cast<int>(star.x + 0.5f);
        int py = static_cast<int>(star.y + 0.5f);

        // Calculate radius based on flux (brighter stars get larger radius)
        float radius = 5.0f + 5.0f * std::log10(1.0f + star.flux / 1000.0f);

        // Draw circle for each star
        cv::circle(mask, cv::Point(px, py), static_cast<int>(radius), cv::Scalar(255), -1);
    }

    return mask;
}

bool StarDetector::detectAndMatchStars(const cv::Mat &refImage, const cv::Mat &frameImage,
                                       std::vector<cv::Point2f> &refPoints,
                                       std::vector<cv::Point2f> &framePoints)
{
    // Clear output vectors
    refPoints.clear();
    framePoints.clear();

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

        if (frameImage.channels() > 1)
        {
            cv::cvtColor(frameImage, frameGray, cv::COLOR_BGR2GRAY);
        }
        else
        {
            frameGray = frameImage.clone();
        }

        // Create StellarSolver instances
        std::unique_ptr<StellarSolver> refSolver(new StellarSolver());
        std::unique_ptr<StellarSolver> frameSolver(new StellarSolver());

        // Configure StellarSolver parameters for star extraction only
        refSolver->setProperty("ProcessType", SSolver::EXTRACT);
        frameSolver->setProperty("ProcessType", SSolver::EXTRACT);

        // Set up parameters for star extraction
        SSolver::Parameters params;
        params.minarea = 5;
        params.deblend_thresh = 32;
        params.deblend_contrast = 0.005;
        params.kron_fact = 2.5;
        params.subpix = 5;

        refSolver->setParameters(params);
        frameSolver->setParameters(params);

        // Create FITSData objects for the images
        QSharedPointer<FITSData> refFITS(new FITSData());
        QSharedPointer<FITSData> frameFITS(new FITSData());

        // Convert OpenCV images to FITSData
        if (!refFITS->loadFromOpenCV(refGray, "reference.fits"))
        {
            m_lastError = "Failed to load reference image into FITSData";
            LOG_ERROR(m_lastError);
            return false;
        }

        if (!frameFITS->loadFromOpenCV(frameGray, "frame.fits"))
        {
            m_lastError = "Failed to load frame image into FITSData";
            LOG_ERROR(m_lastError);
            return false;
        }

        // Load the image buffers into StellarSolver
        refSolver->loadNewImageBuffer(refFITS->getStatistics(), refFITS->getImageBuffer());
        frameSolver->loadNewImageBuffer(frameFITS->getStatistics(), frameFITS->getImageBuffer());

        // Extract stars from reference image
        refSolver->extract();

        // Wait for extraction to complete
        if (!refSolver->extractionDone() || refSolver->failed())
        {
            m_lastError = "StellarSolver failed to extract stars from reference image";
            LOG_ERROR(m_lastError);
            return false;
        }

        // Extract stars from frame image
        frameSolver->extract();

        // Wait for extraction to complete
        if (!frameSolver->extractionDone() || frameSolver->failed())
        {
            m_lastError = "StellarSolver failed to extract stars from frame image";
            LOG_ERROR(m_lastError);
            return false;
        }

        // Get the extracted stars
        const QList<FITSImage::Star> &refStars = refSolver->getStarList();
        const QList<FITSImage::Star> &frameStars = frameSolver->getStarList();

        // Check if we have enough stars
        if (refStars.size() < 3 || frameStars.size() < 3)
        {
            m_lastError = "Not enough stars detected for alignment";
            LOG_ERROR(m_lastError);
            return false;
        }

        LOG_INFO("StellarSolver detected " + std::to_string(refStars.size()) + " stars in reference image");
        LOG_INFO("StellarSolver detected " + std::to_string(frameStars.size()) + " stars in frame image");

        // Convert to vectors for processing
        std::vector<std::tuple<float, float, float>> refStarsVec;
        std::vector<std::tuple<float, float, float>> frameStarsVec;

        for (const auto &star : refStars)
        {
            refStarsVec.push_back(std::make_tuple(star.x, star.y, star.flux));
        }

        for (const auto &star : frameStars)
        {
            frameStarsVec.push_back(std::make_tuple(star.x, star.y, star.flux));
        }

        // Sort stars by flux (brightest first)
        std::sort(refStarsVec.begin(), refStarsVec.end(),
                  [](const auto & a, const auto & b)
        {
            return std::get<2>(a) > std::get<2>(b);
        });

        std::sort(frameStarsVec.begin(), frameStarsVec.end(),
                  [](const auto & a, const auto & b)
        {
            return std::get<2>(a) > std::get<2>(b);
        });

        // Use a more robust approach for matching stars
        // Extract the top N brightest stars from each image
        int maxStars = std::min(static_cast<int>(refStarsVec.size()), static_cast<int>(frameStarsVec.size()));
        maxStars = std::min(maxStars, m_maxStars);

        if (maxStars < 3)
        {
            m_lastError = "Not enough stars for matching";
            LOG_ERROR(m_lastError);
            return false;
        }

        // Create vectors of points for the brightest stars
        std::vector<cv::Point2f> refCandidates, frameCandidates;
        for (int i = 0; i < maxStars; i++)
        {
            refCandidates.push_back(cv::Point2f(std::get<0>(refStarsVec[i]), std::get<1>(refStarsVec[i])));
            frameCandidates.push_back(cv::Point2f(std::get<0>(frameStarsVec[i]), std::get<1>(frameStarsVec[i])));
        }

        // Try to find the best transformation between the two sets of points
        // First, try to find a translation-only transformation
        cv::Mat translationModel = cv::estimateAffinePartial2D(
                                       frameCandidates, refCandidates, cv::noArray(), cv::RANSAC, 3.0,
                                       2000, 0.99, maxStars * 2); // Use RANSAC with a high confidence

        if (!translationModel.empty())
        {
            // Apply the transformation to the frame points
            std::vector<cv::Point2f> transformedFramePoints;
            cv::transform(frameCandidates, transformedFramePoints, translationModel);

            // Find the closest pairs of points
            std::vector<std::pair<int, int>> matches; // (refIdx, frameIdx)
            std::vector<bool> refMatched(refCandidates.size(), false);
            std::vector<bool> frameMatched(frameCandidates.size(), false);

            // For each reference point, find the closest transformed frame point
            for (size_t i = 0; i < refCandidates.size(); i++)
            {
                double minDist = std::numeric_limits<double>::max();
                int bestMatch = -1;

                for (size_t j = 0; j < transformedFramePoints.size(); j++)
                {
                    if (frameMatched[j]) continue; // Skip already matched frame points

                    double dist = cv::norm(refCandidates[i] - transformedFramePoints[j]);
                    if (dist < minDist)
                    {
                        minDist = dist;
                        bestMatch = j;
                    }
                }

                // If we found a good match and it's within a reasonable distance
                if (bestMatch >= 0 && minDist < 50.0) // 50 pixel threshold
                {
                    matches.push_back(std::make_pair(i, bestMatch));
                    refMatched[i] = true;
                    frameMatched[bestMatch] = true;
                }
            }

            // If we have enough matches, use them
            if (matches.size() >= 3)
            {
                for (const auto &match : matches)
                {
                    refPoints.push_back(refCandidates[match.first]);
                    framePoints.push_back(frameCandidates[match.second]);
                }
                LOG_INFO("Found " + std::to_string(matches.size()) + " star matches using transformation model");
                return true;
            }
        }

        // If the transformation-based matching failed, fall back to the simple approach
        refPoints.clear();
        framePoints.clear();

        // Assume the brightest stars correspond to each other
        for (int i = 0; i < maxStars; i++)
        {
            refPoints.push_back(cv::Point2f(std::get<0>(refStarsVec[i]), std::get<1>(refStarsVec[i])));
            framePoints.push_back(cv::Point2f(std::get<0>(frameStarsVec[i]), std::get<1>(frameStarsVec[i])));
        }

        LOG_INFO("Using " + std::to_string(refPoints.size()) + " brightest stars for matching (fallback method)");

        // Calculate the transformation matrix to check for special cases
        if (refPoints.size() >= 3)
        {
            cv::Mat H = cv::findHomography(framePoints, refPoints, cv::RANSAC, 3.0);

            if (!H.empty() && H.rows == 3 && H.cols == 3)
            {
                // For a homography matrix H = [h11 h12 h13; h21 h22 h23; h31 h32 h33]
                // Calculate the determinant of the upper-left 2x2 submatrix
                double a = H.at<double>(0, 0);
                double b = H.at<double>(0, 1);
                double c = H.at<double>(1, 0);
                double d = H.at<double>(1, 1);

                double det = a * d - b * c;

                // If determinant is negative, the transformation includes a reflection (flip)
                if (det < 0)
                {
                    LOG_INFO("Detected a meridian flip (180-degree rotation with reflection)");
                }

                // Calculate rotation angle in degrees
                double angle = std::atan2(b, a) * 180.0 / M_PI;
                LOG_INFO("Rotation angle: " + std::to_string(angle) + " degrees");

                // Check if we have a rotation close to 180 degrees
                if (std::abs(std::abs(angle) - 180.0) < 20.0)
                {
                    LOG_INFO("Detected a 180-degree rotation");
                    LOG_WARN("WARNING: 180-degree rotation detected. This may affect alignment quality.");
                }
            }
        }

        return true;
    }
    catch (const std::exception &e)
    {
        m_lastError = "Error in star detection: " + std::string(e.what());
        LOG_ERROR(m_lastError);
        return false;
    }
}

bool StarDetector::analyzeTransformation(const std::vector<cv::Point2f> &refPoints,
        const std::vector<cv::Point2f> &framePoints,
        cv::Mat &H)
{
    // Check if we have enough points
    if (refPoints.size() < 3 || framePoints.size() < 3)
    {
        m_lastError = "Not enough points for transformation analysis";
        LOG_ERROR(m_lastError);
        return false;
    }

    // Calculate the transformation matrix (homography)
    if (refPoints.size() >= 4)
    {
        // Use full homography if we have enough points
        H = cv::findHomography(framePoints, refPoints, cv::RANSAC, 3.0);
    }
    else if (refPoints.size() >= 2)
    {
        // Use affine transformation if we have at least 2 points
        H = cv::estimateAffinePartial2D(framePoints, refPoints);

        // Convert to 3x3 homography matrix
        if (!H.empty())
        {
            cv::Mat fullH = cv::Mat::eye(3, 3, CV_64F);
            H.copyTo(fullH.rowRange(0, 2));
            H = fullH;
        }
    }
    else
    {
        // Use translation only if we have only 1 point
        H = cv::Mat::eye(3, 3, CV_64F);
        H.at<double>(0, 2) = refPoints[0].x - framePoints[0].x;
        H.at<double>(1, 2) = refPoints[0].y - framePoints[0].y;
    }

    if (H.empty())
    {
        m_lastError = "Failed to calculate transformation matrix";
        LOG_ERROR(m_lastError);
        return false;
    }

    // Extract transformation parameters
    // For a homography matrix H = [h11 h12 h13; h21 h22 h23; h31 h32 h33]

    // Translation (tx, ty) is directly from h13, h23
    double tx = H.at<double>(0, 2);
    double ty = H.at<double>(1, 2);

    // For rotation and scaling, we need to analyze the upper-left 2x2 submatrix
    // [h11 h12; h21 h22]
    double a = H.at<double>(0, 0);
    double b = H.at<double>(0, 1);
    double c = H.at<double>(1, 0);
    double d = H.at<double>(1, 1);

    // Calculate scaling factors
    double sx = std::sqrt(a * a + b * b);
    double sy = std::sqrt(c * c + d * d);

    // Calculate rotation angle in degrees
    double angle = std::atan2(b, a) * 180.0 / M_PI;

    // Calculate determinant to check if there's a reflection
    double det = a * d - b * c;

    // Calculate RMSE (Root Mean Square Error) of the transformation
    double totalError = 0.0;
    std::vector<cv::Point2f> transformedPoints;
    cv::perspectiveTransform(framePoints, transformedPoints, H);

    for (size_t i = 0; i < refPoints.size(); i++)
    {
        double dx = refPoints[i].x - transformedPoints[i].x;
        double dy = refPoints[i].y - transformedPoints[i].y;
        totalError += dx * dx + dy * dy;
    }

    double rmse = std::sqrt(totalError / refPoints.size());

    // Log the transformation parameters
    LOG_INFO("Transformation Analysis:");
    LOG_INFO("  Translation: (" + std::to_string(tx) + ", " + std::to_string(ty) + ") pixels");
    LOG_INFO("  Rotation: " + std::to_string(angle) + " degrees");
    LOG_INFO("  Scaling: (" + std::to_string(sx) + ", " + std::to_string(sy) + ")");
    LOG_INFO("  Determinant: " + std::to_string(det) + (det < 0 ? " (includes reflection)" : ""));
    LOG_INFO("  RMSE: " + std::to_string(rmse) + " pixels");

    // Check for special transformations
    if (std::abs(std::abs(angle) - 180.0) < 20.0)
    {
        LOG_INFO("Detected a 180-degree rotation");
    }

    if (det < 0)
    {
        LOG_INFO("Detected a meridian flip (reflection)");
    }

    // Print the homography matrix for debugging
    LOG_INFO("Homography matrix:");
    for (int i = 0; i < 3; i++)
    {
        std::string row;
        for (int j = 0; j < 3; j++)
        {
            row += std::to_string(H.at<double>(i, j)) + " ";
        }
        LOG_INFO(row);
    }

    return true;
}
