stacker.cpp

#include "stacker.hpp"
#include <algorithm>
#include <cmath>
#include <filesystem>
#include <iostream>
#include <map>
#include <opencv2/core.hpp>
#include <opencv2/core/mat.hpp>
#include <opencv2/core/types.hpp>
#include <opencv2/imgcodecs.hpp>
#include <opencv2/opencv.hpp>
#include <string>
#include <unordered_map>
#include <vector>

// stack the given frames over each other
// cv::Mat StackFrames(vector<cv::Mat> &frames)
// {
// 	if (frames.empty())
// 		return cv::Mat();
//
// 	for (auto &frame : frames)
// 		if (frame.rows > frame.cols)
// 			cv::rotate(frame, frame, cv::ROTATE_90_CLOCKWISE);
//
// 	// reference for alignment
// 	cv::Mat reference = frames[0];
// 	cv::Mat sum = cv::Mat::zeros(reference.size(), CV_32FC3);
//
// 	cv::Mat ref_float;
// 	reference.convertTo(ref_float, CV_32F);
// 	sum += ref_float;
//
// 	for (int i = 1; i < (int)frames.size(); i++)
// 	{
// 		std::cout << "aligning frame " << i << "/" << frames.size() - 1
// 		          << std::endl;
//
// 		// align the frame
// 		cv::Mat aligned = AlignToReference(reference, frames[i]);
//
// 		cv::Mat as_float;
// 		aligned.convertTo(as_float, CV_32F);
// 		sum += as_float;
// 	}
//
// 	cv::Mat result;
// 	sum.convertTo(result, CV_8U, 1.0 / frames.size());
// 	return result;
// }

// remove the background and noise, keep only the stars
cv::Mat RemoveBackground(const cv::Mat &mat, cv::Size kSize, int sigmaX)
{
#ifdef CUDA_ENABLED
	cv::cuda::GpuMat d_gray, d_blurred, d_stars;
	d_gray.upload(mat);
	// CUDA's SeparableLinearFilter (used by createGaussianFilter) hard-limits
	// kernel width to <=32 pixels.  At background-subtraction scales (101px)
	// a box filter is perceptually equivalent and has no size restriction.
	auto filter = cv::cuda::createBoxFilter(
	    d_gray.type(), d_gray.type(), kSize);
	filter->apply(d_gray, d_blurred);
	cv::cuda::subtract(d_gray, d_blurred, d_stars);
	cv::Mat stars;
	d_stars.download(stars);
	return stars;
#else
	cv::Mat gray = mat.clone();

	// blur the current image
	cv::Mat blurred;
	cv::GaussianBlur(gray, blurred, kSize, sigmaX);

	cv::Mat stars;
	cv::subtract(gray, blurred, stars);

	return stars;
#endif
}

cv::Mat DetectStars(const cv::Mat &gray)
{
	// remove background
	cv::Mat stars = RemoveBackground(gray, cv::Size(101, 101), 50);

	// normalise
	cv::Mat norm;
	cv::normalize(stars, norm, 0, 255, cv::NORM_MINMAX);

	norm.convertTo(norm, CV_8U);

	// Adaptive threshold using percentile
	// NOTE: keeps top bright pixels, not stars
	// ISSUE: might fail with bright galaxies/halos. need to keep bright stars
	// instead of pixels
	cv::Mat flat = norm.reshape(1, 1);
	vector<uchar> pixels;

	flat.copyTo(pixels);
	std::sort(pixels.begin(), pixels.end());

	double percentile = 99.99;
	int idx = static_cast<int>((percentile / 100.0) * pixels.size());
	idx = std::min(idx, (int)pixels.size() - 1);
	uchar thresholdValue = pixels[idx];

	cv::Mat binary;
	cv::threshold(norm, binary, thresholdValue, 255, cv::THRESH_BINARY);

	// remove tiny noise
	cv::Mat kernel =
	    cv::getStructuringElement(cv::MORPH_ELLIPSE, cv::Size(3, 3));

	cv::morphologyEx(binary, binary, cv::MORPH_OPEN, kernel);

	return binary;
}

// TODO: Add isolation criteria -> stars must be spread out
vector<StarCandidate> IdentifyCandidates(const cv::Mat &stars,
                                         const cv::Mat &gray)
{
	vector<StarCandidate> candidates = {};

	cv::Mat labels, stats, centroids;
	int candidateCount =
	    cv::connectedComponentsWithStats(stars, labels, stats, centroids);

	// 0 => background. start from 1
	for (int i = 1; i < candidateCount; i++)
	{
		// ignore sensor noise
		double area = stats.at<int>(i, cv::CC_STAT_AREA);
		// TODO: use adaptive area filtering
		if (area < 5)
			continue;

		// bounding box
		int left = stats.at<int>(i, cv::CC_STAT_LEFT);
		int top = stats.at<int>(i, cv::CC_STAT_TOP);

		int width = stats.at<int>(i, cv::CC_STAT_WIDTH);
		int height = stats.at<int>(i, cv::CC_STAT_HEIGHT);

		cv::Rect bbox = {left, top, width, height};

		// get the stars original brightness value from grayscale image
		// ISSUE: bbox contains background
		// ISSUE: bbox may contain nearby stars (which pollutes
		// measurement)
		double brightness = cv::sum(gray(bbox))[0];

		// componentMask contains only the current object and nothing else
		const cv::Mat componentMask = (labels(bbox) == i);

		double circularity = CalculateCircularity(componentMask);
		double eccentricity = CalculateEccentricity(componentMask);
		double contrast = CalculateLocalContrast(gray, bbox);
		cv::Point2d centroid = CalculateCentroid(gray, bbox);

		// filter candidates further
		if (eccentricity > 0.8)
			continue;
		if (circularity < 0.4)
			continue;
		if (area < 5)
			continue;

		// create and append StarCandidate object
		StarCandidate candidate;
		candidate.bbox = bbox;
		candidate.brightness = brightness;
		candidate.circularity = circularity;
		candidate.eccentricity = eccentricity;
		candidate.contrast = contrast;
		candidate.area = area;
		candidate.centroid = centroid;

		candidates.push_back(candidate);
	}

	return candidates;
}

double CalculateEccentricity(const cv::Mat &componentMask)
{
	// find the contours
	vector<vector<cv::Point>> contours;
	cv::findContours(componentMask, contours, cv::RETR_EXTERNAL,
	                 cv::CHAIN_APPROX_SIMPLE);

	// No contour found
	if (contours.empty())
		return 1.0;

	// Need at least 5 points for ellipse fitting
	if (contours[0].size() < 5)
		return 1.0;

	// Fit ellipse
	cv::RotatedRect ellipse = cv::fitEllipse(contours[0]);

	// Extract semi-major and semi-minor axes
	double a = std::max(ellipse.size.width, ellipse.size.height) / 2.0;
	double b = std::min(ellipse.size.width, ellipse.size.height) / 2.0;

	// Prevent division by zero
	if (a <= 0.0)
		return 1.0;

	// Eccentricity formula
	double eccentricity = std::sqrt(1.0 - (b * b) / (a * a));

	return eccentricity;
}

double CalculateLocalContrast(const cv::Mat &gray, const cv::Rect &bbox)
{
	// the radius of the region to consider for measuring contrast around the
	// star
	int padding = 10;

	// region containing star + background
	cv::Rect expanded(bbox.x - padding, bbox.y - padding,
	                  bbox.width + 2 * padding, bbox.height + 2 * padding);

	// fix to prevent errors when accessing out-of-bound pixels
	// &= -> intersection while handling overflowing
	expanded &= cv::Rect(0, 0, gray.cols, gray.rows);

	cv::Mat region = gray(expanded);

	// brightness of star alone
	double peak;
	cv::minMaxLoc(gray(bbox), nullptr, &peak);

	// extract background without star
	cv::Mat bgMask(expanded.height, expanded.width, CV_8U, cv::Scalar(255));

	// color the star's bbox black within the bgmask
	cv::rectangle(bgMask,
	              cv::Rect(bbox.x - expanded.x, bbox.y - expanded.y, bbox.width,
	                       bbox.height),
	              0, cv::FILLED);

	// brightness of background
	double background = cv::mean(region, bgMask)[0];

	// contrast
	return peak - background;
}

double CalculateCircularity(const cv::Mat &componentMask)
{
	vector<vector<cv::Point>> contours;

	cv::findContours(componentMask, contours, cv::RETR_EXTERNAL,
	                 cv::CHAIN_APPROX_SIMPLE);

	if (contours.empty())
		return 0.0;

	double area = cv::contourArea(contours[0]);
	double perimeter = cv::arcLength(contours[0], true);

	if (perimeter <= 0.0)
		return 0.0;

	double circularity = 4.0 * CV_PI * area / (perimeter * perimeter);

	return circularity;
}

cv::Point2d CalculateCentroid(const cv::Mat &gray, const cv::Rect &bbox)
{
	cv::Mat patch = gray(bbox);

	// weighted sums
	double sumX = 0.0;
	double sumY = 0.0;

	// total weighted sum
	double sumI = 0.0;

	for (int y = 0; y < patch.rows; y++)
	{
		for (int x = 0; x < patch.cols; x++)
		{
			double intensity = patch.at<uchar>(y, x);

			sumX += x * intensity;
			sumY += y * intensity;
			sumI += intensity;
		}
	}

	if (sumI <= 0.0)
	{
		return cv::Point2d(bbox.x + bbox.width / 2.0,
		                   bbox.y + bbox.height / 2.0);
	}

	double cx = sumX / sumI;
	double cy = sumY / sumI;

	// convert to global coords
	cx += bbox.x;
	cy += bbox.y;

	return cv::Point2d(cx, cy);
}

vector<StarCandidate> RankCandidates(vector<StarCandidate> &candidates)
{
	vector<StarCandidate> rankedCandidates = {};

	for (auto &star : candidates)
	{

		// weighted ranking
		// TODO: fix hardcoded weights???
		double score = 3.25 * star.normalizedArea +
		               2.25 * star.normalizedBrightness +
		               3.25 * star.normalizedContrast +
		               0.25 * star.circularity + 0.25 * (1 - star.eccentricity);
		// double score = star.area * star.circularity * star.brightness *
		//                star.contrast * (1 - star.eccentricity);

		star.score = score;

		rankedCandidates.push_back(star);
	}

	std::sort(rankedCandidates.begin(), rankedCandidates.end(),
	          [](auto &a, auto &b) { return a.score > b.score; });

	return rankedCandidates;
}

vector<StarCandidate>
NormalizeCandidateFeatures(vector<StarCandidate> &candidates)
{
	if (candidates.empty())
		return candidates;

	double minBrightness = std::numeric_limits<double>::max();
	double maxBrightness = std::numeric_limits<double>::lowest();

	double minContrast = std::numeric_limits<double>::max();
	double maxContrast = std::numeric_limits<double>::lowest();

	double minArea = std::numeric_limits<double>::max();
	double maxArea = std::numeric_limits<double>::lowest();

	// find min/max values
	for (const auto &star : candidates)
	{
		minBrightness = std::min(minBrightness, star.brightness);
		maxBrightness = std::max(maxBrightness, star.brightness);
		minContrast = std::min(minContrast, star.contrast);
		maxContrast = std::max(maxContrast, star.contrast);
		minArea = std::min(minArea, star.area);
		maxArea = std::max(maxArea, star.area);
	}

	// compute ranges safely
	double brightnessRange = maxBrightness - minBrightness;
	double contrastRange = maxContrast - minContrast;
	double areaRange = maxArea - minArea;

	if (brightnessRange == 0)
		brightnessRange = 1;

	if (contrastRange == 0)
		contrastRange = 1;

	if (areaRange == 0)
		areaRange = 1;

	// normalize to 0 -> 1
	for (auto &star : candidates)
	{
		star.normalizedBrightness =
		    (star.brightness - minBrightness) / brightnessRange;
		star.normalizedContrast = (star.contrast - minContrast) / contrastRange;
		star.normalizedArea = (star.area - minArea) / areaRange;
	}

	return candidates;
}

// find k nearest neighbouring stars for each star
// PERF: use KD tree or something advanced?
// FIX: find k nearest instead of all stars
vector<vector<Neighbour>> BuildNeighbours(vector<StarCandidate> &candidates,
                                          int k)
{
	vector<vector<Neighbour>> neighbours;

	// INFO: brute forcing for now
	for (int i = 0; i < candidates.size(); i++)
	{
		vector<Neighbour> currentNeighbours;
		currentNeighbours.reserve(candidates.size());

		for (int j = 0; j < candidates.size(); j++)
		{
			if (i == j)
				continue;

			StarCandidate &star = candidates[i];
			StarCandidate &neighbour = candidates[j];

			double dx = neighbour.centroid.x - star.centroid.x;
			double dy = neighbour.centroid.y - star.centroid.y;

			double dist = dx * dx + dy * dy;

			currentNeighbours.push_back({j, dist});
		}

		// sort the current neighbours by closest distance
		std::sort(currentNeighbours.begin(), currentNeighbours.end(),
		          [](const auto &a, const auto &b) { return a.dist < b.dist; });

		// take first k neighbours alone
		if (currentNeighbours.size() > k)
			currentNeighbours.resize(k);

		neighbours.push_back(currentNeighbours);
	}

	return neighbours;
}

// build triangle description for triplets of stars
vector<TriangleDesc> BuildTriangleDesc(vector<StarCandidate> &candidates,
                                       vector<vector<Neighbour>> &allNeighbours)
{
	vector<TriangleDesc> desc;

	int n = allNeighbours.size();

	for (int i = 0; i < n; i++)
	{
		const auto &neighbours = allNeighbours[i];
		int numNeighbours = neighbours.size();

		for (int j = 0; j < numNeighbours; j++)
		{
			for (int k = j + 1; k < numNeighbours; k++)
			{
				int idxA = i;
				int idxB = neighbours[j].index;
				int idxC = neighbours[k].index;

				StarCandidate &A = candidates[idxA];
				StarCandidate &B = candidates[idxB];
				StarCandidate &C = candidates[idxC];

				// calculate side lengths
				double dxAB = A.centroid.x - B.centroid.x;
				double dyAB = A.centroid.y - B.centroid.y;

				double dxBC = B.centroid.x - C.centroid.x;
				double dyBC = B.centroid.y - C.centroid.y;

				double dxCA = C.centroid.x - A.centroid.x;
				double dyCA = C.centroid.y - A.centroid.y;

				double side1 = dxAB * dxAB + dyAB * dyAB;
				double side2 = dxBC * dxBC + dyBC * dyBC;
				double side3 = dxCA * dxCA + dyCA * dyCA;

				// sort side lengths
				std::array<double, 3> sides = {side1, side2, side3};
				std::sort(sides.begin(), sides.end());

				side1 = sides[0];
				side2 = sides[1];
				side3 = sides[2];

				// Reject degenerate triangles
				if (side3 < 1e-6)
					continue;

				// calculate side ratios
				double r1 = side1 / side3;
				double r2 = side2 / side3;

				// create TriangleDesc
				TriangleDesc t;

				t.star_i = idxA;
				t.star_j = idxB;
				t.star_k = idxC;

				t.side1 = side1;
				t.side2 = side2;
				t.side3 = side3;

				t.r1 = r1;
				t.r2 = r2;

				desc.push_back(t);
			}
		}
	}

	return desc;
}

std::vector<TriangleMatch>
MatchTriangles(const std::vector<TriangleDesc> &refTriangles,
               const std::vector<TriangleDesc> &tgtTriangles, double tolerance)
{
	std::vector<TriangleMatch> matches;

	for (const auto &tgt : tgtTriangles)
	{
		for (const auto &ref : refTriangles)
		{
			// Check if shape ratios match
			double err1 = std::abs(tgt.r1 - ref.r1);
			double err2 = std::abs(tgt.r2 - ref.r2);

			if (err1 < tolerance && err2 < tolerance)
			{
				TriangleMatch match;
				match.ref_i = ref.star_i;
				match.ref_j = ref.star_j;
				match.ref_k = ref.star_k;

				match.tgt_i = tgt.star_i;
				match.tgt_j = tgt.star_j;
				match.tgt_k = tgt.star_k;

				match.error = err1 + err2;

				matches.push_back(match);
			}
		}
	}

	// Sort by error, return best matches
	std::sort(matches.begin(), matches.end(),
	          [](const auto &a, const auto &b) { return a.error < b.error; });

	return matches;
}

// Vote-deduplicate star pair correspondences from triangle matches.
//
// Each TriangleMatch encodes 3 star-pair correspondences. If the same pair
// (ref_star, tgt_star) appears in many triangle matches it means many
// independent triangles agree — that's a strong vote for correctness.
// Only pairs with >= minVotes are returned.
vector<StarPair>
ExtractCorrespondences(const vector<TriangleMatch> &matches,
                       const vector<StarCandidate> &refCandidates,
                       const vector<StarCandidate> &tgtCandidates, int minVotes)
{
	// key = (ref_idx, tgt_idx), value = vote count
	std::map<std::pair<int, int>, int> voteMap;

	for (const auto &m : matches)
	{
		// Each triangle match gives 3 correspondences (one per vertex)
		voteMap[{m.ref_i, m.tgt_i}]++;
		voteMap[{m.ref_j, m.tgt_j}]++;
		voteMap[{m.ref_k, m.tgt_k}]++;
	}

	vector<StarPair> pairs;

	for (const auto &[key, votes] : voteMap)
	{
		if (votes < minVotes)
			continue;

		auto [refIdx, tgtIdx] = key;

		// Guard against stale indices
		if (refIdx >= (int)refCandidates.size() ||
		    tgtIdx >= (int)tgtCandidates.size())
			continue;

		StarPair pair;
		pair.ref_pt = refCandidates[refIdx].centroid;
		pair.tgt_pt = tgtCandidates[tgtIdx].centroid;
		pair.votes = votes;
		pairs.push_back(pair);
	}

	// Sort by descending votes so the caller gets the most-agreed-on pairs
	// first
	std::sort(pairs.begin(), pairs.end(),
	          [](const auto &a, const auto &b) { return a.votes > b.votes; });

	return pairs;
}

// Estimate a 4-DOF similarity transform (translation + rotation + uniform
// scale) from star pair correspondences.
//
// Uses cv::estimateAffinePartial2D with built-in RANSAC. Returns an empty M
// and 0 inliers on failure — no I/O side-effects.
TransformResult EstimateTransform(const vector<StarPair> &pairs)
{
	if ((int)pairs.size() < 3)
		return {cv::Mat(), 0};

	vector<cv::Point2f> refPts, tgtPts;
	refPts.reserve(pairs.size());
	tgtPts.reserve(pairs.size());

	for (const auto &p : pairs)
	{
		refPts.push_back(p.ref_pt);
		tgtPts.push_back(p.tgt_pt);
	}

	// estimateAffinePartial2D = similarity transform (no shear)
	// RANSAC threshold: 3px reprojection error
	cv::Mat inlierMask;
	cv::Mat M = cv::estimateAffinePartial2D(tgtPts, refPts, inlierMask,
	                                        cv::RANSAC, 3.0);

	int inlierCount = M.empty() ? 0 : cv::countNonZero(inlierMask);
	return {M, inlierCount};
}

// ---------------------------------------------------------------------------
// Pipeline subfunctions
// ---------------------------------------------------------------------------

// Draw star circles on a grayscale image and save it as a debug PNG.
void DebugWrite(const cv::Mat &gray, const vector<StarCandidate> &candidates,
                const std::string &outPath)
{
	cv::Mat dbg;
	cv::cvtColor(gray, dbg, cv::COLOR_GRAY2BGR);

	for (size_t i = 0; i < candidates.size(); i++)
	{
		const auto &star = candidates[i];
		cv::circle(dbg, star.centroid, 18, cv::Scalar(0, 255, 0), 1);
		cv::putText(dbg, std::to_string(i),
		            star.centroid + cv::Point2d(10, -10),
		            cv::FONT_HERSHEY_SIMPLEX, 0.35, cv::Scalar(0, 255, 255), 1);
	}

	cv::imwrite(outPath, dbg);
}

// Run the per-frame star detection + descriptor pipeline on one colour image.
// Returns ranked candidates and their triangle descriptors.
// Writes debug PNGs to test/<debugPrefix>_*.png when debugPrefix is non-empty.
// No console output — all reporting is the caller's responsibility.
FrameResult ProcessFrame(const cv::Mat &img, const std::string &debugPrefix)
{
	cv::Mat gray;
	cv::cvtColor(img, gray, cv::COLOR_BGR2GRAY);

	cv::Mat stars = DetectStars(gray);

	vector<StarCandidate> candidates = IdentifyCandidates(stars, gray);
	vector<StarCandidate> normCandidates =
	    NormalizeCandidateFeatures(candidates);
	vector<StarCandidate> ranked = RankCandidates(normCandidates);

	vector<vector<Neighbour>> neighbours = BuildNeighbours(ranked, 10);
	vector<TriangleDesc> triangles = BuildTriangleDesc(ranked, neighbours);

	if (!debugPrefix.empty())
	{
		DebugWrite(gray, candidates, "test/" + debugPrefix + "_detected.png");
		DebugWrite(gray, ranked, "test/" + debugPrefix + "_chosen.png");
		cv::imwrite("test/" + debugPrefix + "_starmask.png", stars);
	}

	return FrameResult{ranked, triangles};
}

// Load and validate all frames from disk.
// Returns an empty vector if any path fails to load.
// No console output — all reporting is the caller's responsibility.
vector<cv::Mat> LoadFrames(const vector<std::string> &paths)
{
	vector<cv::Mat> frames;
	frames.reserve(paths.size());

	for (const auto &p : paths)
	{
		cv::Mat f = cv::imread(p);
		if (f.empty())
		{
			std::cerr << "  [warn] LoadFrames: could not read \"" << p
			          << "\" — skipping\n";
			continue;
		}
		frames.push_back(std::move(f));
	}

	return frames;
}

// Warp a target frame onto the reference coordinate system using transform M.
// Writes a debug aligned PNG to test/<debugPrefix>_aligned.png.
// When compiled with CUDA_ENABLED the warp runs on the GPU via cv::cuda::warpAffine.
cv::Mat AlignFrame(const cv::Mat &frame, const cv::Mat &M,
                   const cv::Size &refSize, const std::string &debugPrefix)
{
#ifdef CUDA_ENABLED
	cv::cuda::GpuMat d_frame, d_aligned;
	d_frame.upload(frame);
	cv::cuda::warpAffine(d_frame, d_aligned, M, refSize);
	cv::Mat aligned;
	d_aligned.download(aligned);
	if (!debugPrefix.empty())
		cv::imwrite("test/" + debugPrefix + "_aligned.png", aligned);
	return aligned;
#else
	cv::Mat aligned;
	cv::warpAffine(frame, aligned, M, refSize);
	if (!debugPrefix.empty())
		cv::imwrite("test/" + debugPrefix + "_aligned.png", aligned);
	return aligned;
#endif
}

// Divide the float32 accumulator by count, then apply a percentile-based
// histogram stretch before converting to 8-bit.
// Without stretching, a raw stacked astrophoto looks identical to a single
// frame because all signal lives in a tiny low-value band of the 0-255 range.
cv::Mat MeanStack(const cv::Mat &sum, int count)
{
	cv::Mat meanF = sum / static_cast<double>(count);

	// --- Percentile stretch ---
	// Flatten all BGR float values into one sorted list and find the
	// 0.1th and 99.9th percentiles.  This maps faint galaxy detail into
	// the displayable range while cleanly clipping hot pixels at the top.
	cv::Mat flat = meanF.reshape(1, 1); // 1-channel row vector
	std::vector<float> pixels(flat.begin<float>(), flat.end<float>());
	std::sort(pixels.begin(), pixels.end());

	size_t n = pixels.size();
	double lo = pixels[static_cast<size_t>(0.001 * n)];
	double hi = pixels[static_cast<size_t>(0.999 * n)];
	if (hi <= lo)
		hi = lo + 1.0; // guard against flat frames

	cv::Mat stretched = (meanF - lo) * (255.0 / (hi - lo));
	// clamp negatives to zero (pixels darker than the 0.1th percentile)
	cv::threshold(stretched, stretched, 0.0, 0.0, cv::THRESH_TOZERO);

	std::cout << "[stretch] lo=" << lo << "  hi=" << hi << "\n";

	cv::Mat result;
	stretched.convertTo(result, CV_8U);
	cv::imwrite("test/stacked.png", result);

	return result;
}

// ---------------------------------------------------------------------------
// Calibration helpers
// ---------------------------------------------------------------------------

// Build a master calibration frame by mean-stacking all images at `paths`.
// Each frame is loaded as BGR and converted to float32 before accumulation.
// Returns an empty Mat on failure (load error or empty paths).
cv::Mat BuildMaster(const vector<std::string> &paths)
{
	if (paths.empty())
		return cv::Mat();

	vector<cv::Mat> frames = LoadFrames(paths);
	if (frames.empty())
		return cv::Mat();

	cv::Mat sum = cv::Mat::zeros(frames[0].size(), CV_32FC3);
	for (auto &f : frames)
	{
		cv::Mat fF;
		f.convertTo(fF, CV_32F);
		sum += fF;
	}

	return sum / static_cast<double>(frames.size());
}

// Build the three master calibration frames from raw frame paths.
//   master dark  : mean dark, with master bias subtracted if provided
//   master flat  : mean flat, normalised so its mean pixel value == 1.0
// Any type whose path vector is empty is left as an empty Mat (step skipped).
CalibrationFrames BuildCalibrationFrames(const vector<std::string> &biasPaths,
                                         const vector<std::string> &darkPaths,
                                         const vector<std::string> &flatPaths)
{
	CalibrationFrames cal;

	cal.masterBias = BuildMaster(biasPaths);

	// --- Master dark: subtract bias to isolate thermal noise ---
	cal.masterDark = BuildMaster(darkPaths);
	if (!cal.masterDark.empty() && !cal.masterBias.empty())
	{
		cal.masterDark -= cal.masterBias;
		// cv::max(mat, scalar) has undefined multi-channel semantics in
		// OpenCV 5. Use the two-array form: create a zero-filled mat of the
		// same type.
		cv::max(cal.masterDark,
		        cv::Mat::zeros(cal.masterDark.size(), cal.masterDark.type()),
		        cal.masterDark);
	}

	// --- Master flat: normalise so dividing by it is brightness-neutral ---
	cal.masterFlat = BuildMaster(flatPaths);
	if (!cal.masterFlat.empty())
	{
		// Subtract bias from flat before normalising (if available)
		if (!cal.masterBias.empty())
		{
			cal.masterFlat -= cal.masterBias;
			// Clamp flat to >= 1.0 so flat-dividing never amplifies noise to
			// inf.
			cv::Mat ones(cal.masterFlat.size(), cal.masterFlat.type(),
			             cv::Scalar::all(1.0));
			cv::max(cal.masterFlat, ones, cal.masterFlat);
		}

		// Normalise: divide by per-channel mean so mean pixel == 1.0
		cv::Scalar meanVal = cv::mean(cal.masterFlat);
		if (meanVal[0] > 0.0)
			cal.masterFlat /= meanVal[0];

		// Guard against zero/near-zero pixels that would blow up division
		cv::Mat eps(cal.masterFlat.size(), cal.masterFlat.type(),
		            cv::Scalar::all(1e-6));
		cv::max(cal.masterFlat, eps, cal.masterFlat);
	}

	if (!cal.masterBias.empty())
		std::cout << "[cal] master bias built\n";
	if (!cal.masterDark.empty())
		std::cout << "[cal] master dark built (bias-subtracted: "
		          << (!cal.masterBias.empty() ? "yes" : "no") << ")\n";
	if (!cal.masterFlat.empty())
		std::cout << "[cal] master flat built (normalised)\n";

	return cal;
}

// Apply calibration to a single raw light frame.
// Input must be float32 (CV_32FC3). Returns a calibrated float32 frame.
//   1. Subtract masterDark  (removes thermal noise + bias signal)
//   2. Divide by masterFlat (corrects vignetting and dust)
// Each step is skipped when the corresponding master is empty.
cv::Mat CalibrateFrame(const cv::Mat &rawLightF, const CalibrationFrames &cal)
{
	cv::Mat calibrated = rawLightF.clone();

	if (!cal.masterDark.empty())
	{
		calibrated -= cal.masterDark;
		// Clamp negatives: use two-array form (scalar overload undefined for
		// multi-channel arrays in OpenCV 5).
		cv::max(calibrated,
		        cv::Mat::zeros(calibrated.size(), calibrated.type()),
		        calibrated);
	}

	if (!cal.masterFlat.empty())
		cv::divide(calibrated, cal.masterFlat, calibrated);

	return calibrated;
}

// ---------------------------------------------------------------------------
// Stack — the main pipeline
//
// Takes a list of image paths, aligns every frame to the first (reference)
// using star-pattern matching + RANSAC, then mean-stacks all aligned frames.
//
// Returns the stacked cv::Mat (BGR, 8-bit), or an empty Mat on failure.
// ---------------------------------------------------------------------------

// Log the reference frame (frame 0) header and star/triangle counts.
void LogReference(const FrameResult &ref)
{
	std::cout << "\n=== Frame 0 (reference) ===\n";
	std::cout << "  candidates: " << ref.candidates.size()
	          << "  triangles: " << ref.triangles.size() << "\n";
}

// Log all results produced during one loop iteration of Stack().
// Called once per frame, after all computation for that frame is done.
void Log(int frameIndex, int totalFrames, const FrameResult &ref,
         const FrameResult &tgt, const vector<TriangleMatch> &matches,
         const vector<StarPair> &pairs, const TransformResult &tr)
{
	(void)ref; // available for future cross-frame diagnostics

	std::cout << "\n=== Frame " << frameIndex << " / " << totalFrames - 1
	          << " ===\n";
	std::cout << "  candidates: " << tgt.candidates.size()
	          << "  triangles: " << tgt.triangles.size() << "\n";
	std::cout << "  triangle matches: " << matches.size() << "\n";

	if (matches.empty())
	{
		std::cout << "  [skip] no triangle matches\n";
		return;
	}

	std::cout << "  correspondences: " << pairs.size() << "\n";
	std::cout << "  transform: " << pairs.size() << " pairs -> " << tr.inliers
	          << " inliers\n";

	if (tr.M.empty())
		std::cout << "  [skip] transform estimation failed\n";
}

cv::Mat Stack(const vector<std::string> &paths, const CalibrationFrames &cal)
{
	if (paths.empty())
	{
		std::cerr << "Stack: no input paths provided.\n";
		return cv::Mat();
	}

	std::filesystem::create_directories("test");

	// --- Stream frames one at a time to avoid loading the entire dataset
	// into RAM simultaneously.  Previously LoadFrames() decoded all N light
	// frames upfront, keeping gigabytes of decoded pixel data alive for the
	// full duration of Stack().  Now only the reference frame is loaded
	// eagerly; every subsequent frame is loaded, processed, and released
	// inside the loop so peak memory stays roughly constant regardless of N.

	// Load and validate the reference frame (frame 0)
	cv::Mat reference = cv::imread(paths[0]);
	if (reference.empty())
	{
		std::cerr << "Stack: could not load reference frame: " << paths[0] << "\n";
		return cv::Mat();
	}

	// Convert reference to float32, calibrate, accumulate
	cv::Mat refF;
	reference.convertTo(refF, CV_32F);
	cv::Mat refCalibrated = CalibrateFrame(refF, cal);
	refF.release(); // float32 reference no longer needed

	// Accumulator: float32 so we can sum without overflow
	cv::Mat sum = cv::Mat::zeros(reference.size(), CV_32FC3);
	sum += refCalibrated;
	refCalibrated.release(); // accumulated; free the copy

	// Star detection runs on the 8-bit reference (cleaner stars)
	FrameResult ref = ProcessFrame(reference, "frame0");
	LogReference(ref);

	// Keep reference size for warpAffine; release pixel data once star
	// detection is done — reference.size() is stored in the FrameResult.
	cv::Size refSize = reference.size();
	reference.release();

	int alignedCount = 1; // reference always counts
	int totalFrames   = static_cast<int>(paths.size());

	for (int i = 1; i < totalFrames; i++)
	{
		std::string prefix = "frame" + std::to_string(i);

		try
		{
			// Load this frame from disk (not kept in a pre-loaded vector)
			cv::Mat raw = cv::imread(paths[i]);
			if (raw.empty())
			{
				std::cerr << "  [warn] could not load frame " << i
				          << " (" << paths[i] << ") — skipping\n";
				continue;
			}

			// Convert to float32 and calibrate; release the 8-bit source
			// as soon as the conversion is done.
			cv::Mat frameF;
			raw.convertTo(frameF, CV_32F);
			raw.release(); // 8-bit no longer needed

			cv::Mat calibrated = CalibrateFrame(frameF, cal);
			frameF.release(); // pre-calibration float32 no longer needed

			// Convert calibrated float32 back to 8-bit for ProcessFrame
			// (DetectStars and IdentifyCandidates expect CV_8U internally)
			cv::Mat calibrated8u;
			calibrated.convertTo(calibrated8u, CV_8U);

			// --- Computation (no I/O) ---
			FrameResult tgt = ProcessFrame(calibrated8u, prefix);
			calibrated8u.release(); // star detection done; 8-bit copy no longer needed

			vector<TriangleMatch> matches =
			    MatchTriangles(ref.triangles, tgt.triangles);

			vector<StarPair> pairs;
			TransformResult tr;

			if (!matches.empty())
			{
				pairs = ExtractCorrespondences(matches, ref.candidates,
				                               tgt.candidates);
				tr = EstimateTransform(pairs);

				if (!tr.M.empty())
				{
					cv::Mat alignedF;
					try
					{
						alignedF = AlignFrame(calibrated, tr.M, refSize, prefix);
					}
					catch (const cv::Exception &e)
					{
						std::cout << "  [error] frame " << i << " AlignFrame threw: " << e.what() << "\n";
						calibrated.release();
						continue;
					}

					calibrated.release(); // warp done; float32 source no longer needed

					// Reject warps that mapped content entirely off-canvas.
					cv::Scalar meanVal = cv::mean(alignedF);
					bool offCanvas = (meanVal[0] < 0.5 && meanVal[1] < 0.5 &&
					                  meanVal[2] < 0.5);

					if (offCanvas)
					{
						std::cout
						    << "  [skip] warp mapped content off-canvas\n";
					}
					else
					{
						try
						{
							sum += alignedF;
							alignedCount++;
						}
						catch (const cv::Exception &e)
						{
							std::cout << "  [error] frame " << i << " sum += alignedF threw: " << e.what() << "\n";
							continue;
						}
					}
					alignedF.release(); // accumulated (or skipped); free warped frame
				}
				else
				{
					calibrated.release(); // no valid transform; discard
					std::cout << "  [skip] no valid transform\n";
				}
			}
			else
			{
				calibrated.release(); // no matches; discard
			}

			// --- Single log call for this iteration ---
			Log(i, totalFrames, ref, tgt, matches, pairs, tr);
		}
		catch (const cv::Exception &e)
		{
			std::cerr << "\n  [error] frame " << i
			          << " threw OpenCV exception — skipping.\n"
			          << "  " << e.what() << "\n";
		}
	}

	cv::Mat result = MeanStack(sum, alignedCount);

	std::cout << "\nDone. Stacked " << alignedCount << " / " << totalFrames
	          << " frames. Output: test/stacked.png\n";

	return result;
}