#pragma once

#include <opencv2/opencv.hpp>
#include <opencv2/features2d.hpp>
#include <opencv2/calib3d.hpp>
#include <iostream>
#include <cmath>
#include <algorithm>

// Include the star detector
#include "star_detector.h"

template <typename T>
std::shared_ptr<ImageBuffer<T>> Stacker<T>::alignFrame(const std::shared_ptr<ImageBuffer<T>> &frame)
{
    TIME_FUNCTION("alignFrame");

    // If alignment is disabled, return the original frame
    if (!m_settings->getAlignFrames())
    {
        return frame;
    }

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

    try
    {
        // Convert ImageBuffer to OpenCV Mat
        cv::Mat refMat, frameMat;
        {
            Timer conversionTimer("Convert to OpenCV Mat");

            // For 8-bit images
            if (std::is_same<T, uint8_t>::value)
            {
                int refType = m_referenceFrame->getChannels() == 1 ? CV_8UC1 : CV_8UC3;
                int frameType = frame->getChannels() == 1 ? CV_8UC1 : CV_8UC3;

                refMat = cv::Mat(m_referenceFrame->getHeight(), m_referenceFrame->getWidth(), refType, m_referenceFrame->getData());
                frameMat = cv::Mat(frame->getHeight(), frame->getWidth(), frameType, frame->getData());
            }
            // For 16-bit images
            else
            {
                int refType = m_referenceFrame->getChannels() == 1 ? CV_16UC1 : CV_16UC3;
                int frameType = frame->getChannels() == 1 ? CV_16UC1 : CV_16UC3;

                refMat = cv::Mat(m_referenceFrame->getHeight(), m_referenceFrame->getWidth(), refType, m_referenceFrame->getData());
                frameMat = cv::Mat(frame->getHeight(), frame->getWidth(), frameType, frame->getData());
            }
        }

        // Convert to grayscale if needed for feature detection
        cv::Mat refGray, frameGray;
        {
            Timer grayscaleTimer("Convert to grayscale");

            if (m_referenceFrame->getChannels() == 1)
            {
                refGray = refMat;
            }
            else
            {
                cv::cvtColor(refMat, refGray, cv::COLOR_BGR2GRAY);
            }

            if (frame->getChannels() == 1)
            {
                frameGray = frameMat;
            }
            else
            {
                cv::cvtColor(frameMat, frameGray, cv::COLOR_BGR2GRAY);
            }
        }

        // Use StarDetector for star detection (best for astronomical images)
        std::vector<cv::Point2f> refPoints, framePoints;
        bool starDetectionSuccess = false;
        StarDetector starDetector; // Create a star detector at a higher scope

        {
            Timer starDetectorTimer("Star detection");
            try
            {
                // Set the maximum number of stars to use for matching
                // Limit to a reasonable number for Raspberry Pi
                int maxStars = std::min(100, m_settings->getMaxAlignmentStars());
                starDetector.setMaxStars(maxStars);

                LOG_INFO("Using maximum of " + std::to_string(maxStars) + " stars for alignment");

                // Detect and match stars
                starDetectionSuccess = starDetector.detectAndMatchStars(refGray, frameGray, refPoints, framePoints);

                if (starDetectionSuccess)
                {
                    LOG_INFO("StarDetector detected " + std::to_string(refPoints.size()) + " stars in reference frame");
                    LOG_INFO("StarDetector detected " + std::to_string(framePoints.size()) + " stars in current frame");
                    LOG_INFO("StarDetector matched " + std::to_string(std::min(refPoints.size(),
                             framePoints.size())) + " stars between frames");
                }
            }
            catch (const std::exception &e)
            {
                LOG_ERROR("Star detection failed: " + std::string(e.what()));
                LOG_ERROR("Falling back to feature-based detection");
                starDetectionSuccess = false;
            }
        }

        // If StellarSolver failed or didn't find enough stars, fall back to feature-based detection
        if (!starDetectionSuccess || refPoints.size() < 3 || framePoints.size() < 3)
        {
            LOG_INFO("Trying feature-based detection methods...");

            // Clear points from failed SEP attempt
            refPoints.clear();
            framePoints.clear();

            // First try: ORB with more features and better parameters
            cv::Ptr<cv::FeatureDetector> detector = cv::ORB::create(m_settings->getMaxAlignmentStars() * 2, // Detect more features
                                                    1.2f,  // Scale factor
                                                    8,     // Number of levels
                                                    31,    // Edge threshold
                                                    0,     // First level
                                                    2,     // WTA_K
                                                    cv::ORB::HARRIS_SCORE, // Score type
                                                    31,    // Patch size
                                                    20);   // Fast threshold
            std::vector<cv::KeyPoint> refKeypoints, frameKeypoints;
            cv::Mat refDescriptors, frameDescriptors;

            {
                Timer featureTimer("Feature detection");
                detector->detectAndCompute(refGray, cv::noArray(), refKeypoints, refDescriptors);
                detector->detectAndCompute(frameGray, cv::noArray(), frameKeypoints, frameDescriptors);

                LOG_INFO("Detected " + std::to_string(refKeypoints.size()) + " features in reference frame");
                LOG_INFO("Detected " + std::to_string(frameKeypoints.size()) + " features in current frame");
            }

            // If ORB didn't find enough features, try AKAZE which is better for low-contrast astronomical images
            if (refKeypoints.size() < 10 || frameKeypoints.size() < 10)
            {
                LOG_INFO("ORB found too few features, trying AKAZE detector...");
                cv::Ptr<cv::AKAZE> akazeDetector = cv::AKAZE::create();
                akazeDetector->detectAndCompute(refGray, cv::noArray(), refKeypoints, refDescriptors);
                akazeDetector->detectAndCompute(frameGray, cv::noArray(), frameKeypoints, frameDescriptors);

                LOG_INFO("AKAZE detected " + std::to_string(refKeypoints.size()) + " features in reference frame");
                LOG_INFO("AKAZE detected " + std::to_string(frameKeypoints.size()) + " features in current frame");
            }

            // Check if enough features were detected
            if (refKeypoints.size() < 3 || frameKeypoints.size() < 3)
            {
                LOG_INFO("Not enough features for alignment, returning original frame");
                return frame->clone();
            }

            // Match features with better parameters
            std::vector<cv::DMatch> matches;
            std::vector<cv::DMatch> goodMatches;
            {
                Timer matchTimer("Feature matching");

                // Use FLANN matcher for better performance with many features
                cv::Ptr<cv::DescriptorMatcher> matcher;

                // Choose matcher based on descriptor type
                if (refDescriptors.type() == CV_8U)
                {
                    // For binary descriptors like ORB
                    matcher = cv::DescriptorMatcher::create(cv::DescriptorMatcher::BRUTEFORCE_HAMMING);
                }
                else
                {
                    // For floating-point descriptors like AKAZE
                    matcher = cv::DescriptorMatcher::create(cv::DescriptorMatcher::BRUTEFORCE);
                }

                matcher->match(frameDescriptors, refDescriptors, matches);
                LOG_INFO("Found " + std::to_string(matches.size()) + " initial matches");

                // Filter matches based on distance
                double maxDist = 0;
                double minDist = std::numeric_limits<double>::max();

                for (const auto &match : matches)
                {
                    double dist = match.distance;
                    if (dist < minDist) minDist = dist;
                    if (dist > maxDist) maxDist = dist;
                }

                LOG_INFO("Min distance: " + std::to_string(minDist) + ", Max distance: " + std::to_string(maxDist));

                // Use a more adaptive threshold for good matches
                double threshold = std::max(1.5 * minDist, 30.0);
                for (const auto &match : matches)
                {
                    if (match.distance < threshold)
                    {
                        goodMatches.push_back(match);
                    }
                }

                LOG_INFO("Found " + std::to_string(goodMatches.size()) + " good matches after filtering");
            }

            // Check if enough good matches were found
            if (goodMatches.size() < 3)
            {
                LOG_INFO("Not enough good matches for alignment, returning original frame");
                return frame->clone();
            }

            // Extract location of good matches
            for (const auto &match : goodMatches)
            {
                refPoints.push_back(refKeypoints[match.trainIdx].pt);
                framePoints.push_back(frameKeypoints[match.queryIdx].pt);
            }
        }

        // Try different alignment methods in order of complexity
        cv::Mat H;
        bool alignmentFound = false;

        // Use the StarDetector's analyzeTransformation method to calculate and analyze the transformation
        cv::Mat homography;
        bool analysisSuccess = starDetector.analyzeTransformation(refPoints, framePoints, homography);

        if (analysisSuccess && !homography.empty())
        {
            H = homography;
            alignmentFound = true;
            LOG_INFO("Transformation analysis successful");
        }
        else
        {
            // Calculate translation for reference (fallback)
            Timer translationTimer("Find translation");
            if (refPoints.size() >= 1)
            {
                // Calculate mean translation
                cv::Point2f sumDiff(0, 0);
                for (size_t i = 0; i < refPoints.size(); i++)
                {
                    sumDiff += (refPoints[i] - framePoints[i]);
                }
                cv::Point2f meanTranslation = sumDiff * (1.0f / refPoints.size());
                LOG_INFO("Translation alignment: (" + std::to_string(meanTranslation.x) + ", " + std::to_string(
                             meanTranslation.y) + ") pixels");
            }

            // Always try full homography when we have enough points
            if (refPoints.size() >= 3)
            {
                Timer homographyTimer("Find homography");
                // Use the same parameters as test_rotation.cpp
                cv::Mat homography = cv::findHomography(framePoints, refPoints, cv::RANSAC, 3.0);

                // Check if homography is valid
                if (!homography.empty())
                {
                    H = homography;
                    alignmentFound = true;
                    LOG_INFO("Full homography alignment found");
                }
            }
        }

        // If homography failed, fall back to translation-only
        if (!alignmentFound && refPoints.size() >= 1)
        {
            // Calculate mean translation
            cv::Point2f sumDiff(0, 0);
            for (size_t i = 0; i < refPoints.size(); i++)
            {
                sumDiff += (refPoints[i] - framePoints[i]);
            }
            cv::Point2f meanTranslation = sumDiff * (1.0f / refPoints.size());

            // Create translation matrix
            H = cv::Mat::eye(3, 3, CV_64F);
            H.at<double>(0, 2) = meanTranslation.x;
            H.at<double>(1, 2) = meanTranslation.y;

            alignmentFound = true;
            LOG_INFO("Using translation-only alignment as fallback");
        }

        // Check if homography is valid
        if (H.empty())
        {
            // Invalid homography, return original frame
            return frame->clone();
        }

        // Extract translation, rotation, and scaling from homography matrix
        // Decompose the homography matrix to get the transformation parameters
        if (!H.empty() && H.rows == 3 && H.cols == 3)
        {
            // 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;
            if (det < 0)
            {
                LOG_INFO("Detected a reflection in the transformation (determinant: " + std::to_string(det) + ")");
            }

            // Check if transformation is physically reasonable
            bool unreasonableTranslation = (std::abs(tx) > refMat.cols * 2 || std::abs(ty) > refMat.rows * 2);
            bool extremeScaling = (sx < 0.25 || sx > 4.0 || sy < 0.25 || sy > 4.0);

            // Only apply sanity checks for truly extreme transformations
            if (unreasonableTranslation || extremeScaling)
            {
                LOG_INFO("Detected extreme transformation parameters, trying simpler model");

                // Calculate centroid of matched points for a fallback translation
                cv::Point2f refCentroid(0, 0);
                cv::Point2f frameCentroid(0, 0);

                for (size_t i = 0; i < refPoints.size(); i++)
                {
                    refCentroid += refPoints[i];
                    frameCentroid += framePoints[i];
                }

                refCentroid *= (1.0f / refPoints.size());
                frameCentroid *= (1.0f / framePoints.size());

                // Use estimateAffinePartial2D which finds rotation + translation + uniform scale
                // This is more robust than full homography for extreme cases
                cv::Mat rigidTransform = cv::estimateAffinePartial2D(framePoints, refPoints);

                if (!rigidTransform.empty())
                {
                    // Convert to 3x3 homography matrix
                    H = cv::Mat::eye(3, 3, CV_64F);
                    rigidTransform.copyTo(H.rowRange(0, 2));

                    // Recalculate parameters
                    tx = H.at<double>(0, 2);
                    ty = H.at<double>(1, 2);
                    a = H.at<double>(0, 0);
                    b = H.at<double>(0, 1);
                    c = H.at<double>(1, 0);
                    d = H.at<double>(1, 1);
                    sx = std::sqrt(a * a + b * b);
                    sy = std::sqrt(c * c + d * d);
                    angle = std::atan2(b, a) * 180.0 / M_PI;
                    det = a * d - b * c;

                    LOG_INFO("Using rigid transformation (rotation + translation + uniform scale)");
                }
                else
                {
                    // If rigid transform fails, fall back to translation only
                    H = cv::Mat::eye(3, 3, CV_64F);
                    H.at<double>(0, 2) = refCentroid.x - frameCentroid.x;
                    H.at<double>(1, 2) = refCentroid.y - frameCentroid.y;

                    // Recalculate parameters
                    tx = H.at<double>(0, 2);
                    ty = H.at<double>(1, 2);
                    a = H.at<double>(0, 0);
                    b = H.at<double>(0, 1);
                    c = H.at<double>(1, 0);
                    d = H.at<double>(1, 1);
                    sx = std::sqrt(a * a + b * b);
                    sy = std::sqrt(c * c + d * d);
                    angle = std::atan2(b, a) * 180.0 / M_PI;
                    det = a * d - b * c;

                    LOG_INFO("Using translation-only transformation as fallback");
                }
            }

            // Log the transformation parameters
            LOG_INFO("Frame alignment parameters:");
            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) + ")");

            // 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);
            }
        }

        // Create a new buffer for the transformed image
        auto transformed = std::make_shared<ImageBuffer<T>>(
                               m_referenceFrame->getWidth(),
                               m_referenceFrame->getHeight(),
                               m_referenceFrame->getChannels()
                           );

        // Apply homography transformation
        cv::Mat transformedMat;
        {
            Timer warpTimer("Warp perspective");
            cv::warpPerspective(frameMat, transformedMat, H, cv::Size(m_referenceFrame->getWidth(), m_referenceFrame->getHeight()));
        }

        // Copy transformed image back to ImageBuffer
        {
            Timer copyTimer("Copy back to ImageBuffer");

            // Use memcpy for faster data transfer if the data is contiguous
            if (transformedMat.isContinuous() && transformed->isContiguous())
            {
                std::memcpy(transformed->getData(), transformedMat.data,
                            transformed->getWidth() * transformed->getHeight() *
                            transformed->getChannels() * sizeof(T));
            }
            else
            {
                // Fallback to pixel-by-pixel copy if data is not contiguous
                for (int y = 0; y < transformed->getHeight(); ++y)
                {
                    for (int x = 0; x < transformed->getWidth(); ++x)
                    {
                        for (int c = 0; c < transformed->getChannels(); ++c)
                        {
                            if (std::is_same<T, uint8_t>::value)
                            {
                                transformed->setPixel(x, y, transformedMat.at<uint8_t>(y, x * transformed->getChannels() + c), c);
                            }
                            else
                            {
                                transformed->setPixel(x, y, transformedMat.at<uint16_t>(y, x * transformed->getChannels() + c), c);
                            }
                        }
                    }
                }
            }
        }

        return transformed;
    }
    catch (const std::exception &e)
    {
        // If an error occurs, return the original frame
        LOG_ERROR("Error in alignFrame: " + std::string(e.what()));
        return frame->clone();
    }
}
