eclipse_align.py

#!/usr/bin/env python
"""
Align a time-ordered FITS eclipse sequence using a linear drift model.

The solar disk is detected only on the first and last reference frames. All
intermediate frames are shifted according to their real FITS timestamps.
"""

from __future__ import annotations

import argparse
import csv
import math
import random
import sys
from dataclasses import dataclass
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Iterable

import numpy as np
from PIL import Image, ImageDraw
from scipy import ndimage, optimize

try:
    import cv2
except ImportError:  # pragma: no cover - depends on the local environment.
    cv2 = None

try:
    from astropy.io import fits as astropy_fits
except ImportError:  # pragma: no cover - depends on the local environment.
    astropy_fits = None


FITS_SUFFIXES = {".fit", ".fits", ".fts"}
BAYER_PATTERNS = {"RGGB", "BGGR", "GRBG", "GBRG"}
BAYER_TO_CV2 = {
    "RGGB": "COLOR_BayerRGGB2RGB",
    "BGGR": "COLOR_BayerBGGR2RGB",
    "GRBG": "COLOR_BayerGRBG2RGB",
    "GBRG": "COLOR_BayerGBRG2RGB",
}
TIMESTAMP_KEYS = (
    "DATE-OBS",
    "DATEOBS",
    "DATE_OBS",
    "OBS-DATE",
    "DATE",
)
TIME_KEYS = (
    "TIME-OBS",
    "TIMEOBS",
    "TIME_OBS",
    "UT",
    "UTC",
    "TIME",
)


@dataclass(frozen=True)
class FrameInfo:
    path: Path
    timestamp: datetime
    shape: tuple[int, int]


@dataclass(frozen=True)
class SolarCircle:
    x: float
    y: float
    radius: float
    mask: np.ndarray
    boundary: np.ndarray
    inlier_points: np.ndarray


@dataclass(frozen=True)
class AlignmentRow:
    filename: str
    timestamp: datetime
    alpha: float
    predicted_x: float
    predicted_y: float
    shift_x: float
    shift_y: float


class AlignmentError(RuntimeError):
    """Raised for user-facing alignment errors."""


def split_fits_comment(value: str) -> str:
    in_quote = False
    result: list[str] = []
    i = 0
    while i < len(value):
        char = value[i]
        if char == "'":
            in_quote = not in_quote
        if char == "/" and not in_quote:
            break
        result.append(char)
        i += 1
    return "".join(result).strip()


def parse_fits_value(raw_value: str) -> Any:
    value = split_fits_comment(raw_value).strip()
    if not value:
        return ""
    if value.startswith("'"):
        chars: list[str] = []
        i = 1
        while i < len(value):
            if value[i] == "'":
                if i + 1 < len(value) and value[i + 1] == "'":
                    chars.append("'")
                    i += 2
                    continue
                break
            chars.append(value[i])
            i += 1
        return "".join(chars).strip()
    if value == "T":
        return True
    if value == "F":
        return False
    try:
        if any(marker in value.upper() for marker in (".", "E", "D")):
            return float(value.replace("D", "E").replace("d", "e"))
        return int(value)
    except ValueError:
        return value


def read_fits_header(path: Path) -> tuple[dict[str, Any], int]:
    header: dict[str, Any] = {}
    cards_read = 0
    with path.open("rb") as handle:
        while True:
            block = handle.read(2880)
            if not block:
                raise AlignmentError(f"{path.name}: FITS header has no END card")
            if len(block) != 2880:
                raise AlignmentError(f"{path.name}: truncated FITS header")
            for offset in range(0, 2880, 80):
                cards_read += 1
                card = block[offset : offset + 80].decode("ascii", errors="replace")
                key = card[:8].strip()
                if key == "END":
                    header_size = int(math.ceil(cards_read * 80 / 2880.0) * 2880)
                    return header, header_size
                if key and card[8:10] == "= ":
                    header[key.upper()] = parse_fits_value(card[10:])


def fits_data_dtype(bitpix: int) -> np.dtype:
    mapping = {
        8: np.dtype("u1"),
        16: np.dtype(">i2"),
        32: np.dtype(">i4"),
        64: np.dtype(">i8"),
        -32: np.dtype(">f4"),
        -64: np.dtype(">f8"),
    }
    try:
        return mapping[bitpix]
    except KeyError as exc:
        raise AlignmentError(f"Unsupported FITS BITPIX={bitpix}") from exc


def read_fits_minimal(path: Path) -> tuple[np.ndarray, dict[str, Any]]:
    header, header_size = read_fits_header(path)
    naxis = int(header.get("NAXIS", 0))
    if naxis != 2:
        raise AlignmentError(f"{path.name}: expected a 2D FITS image, got NAXIS={naxis}")

    width = int(header["NAXIS1"])
    height = int(header["NAXIS2"])
    bitpix = int(header["BITPIX"])
    dtype = fits_data_dtype(bitpix)
    count = width * height

    with path.open("rb") as handle:
        handle.seek(header_size)
        raw = np.fromfile(handle, dtype=dtype, count=count)

    if raw.size != count:
        raise AlignmentError(f"{path.name}: truncated FITS image data")

    data = raw.reshape((height, width))
    bscale = float(header.get("BSCALE", 1.0))
    bzero = float(header.get("BZERO", 0.0))
    if bscale != 1.0 or bzero != 0.0 or not np.issubdtype(data.dtype, np.floating):
        data = data.astype(np.float32) * bscale + bzero
    else:
        data = data.astype(np.float32, copy=False)
    return data, header


def read_fits(path: Path) -> tuple[np.ndarray, dict[str, Any]]:
    if astropy_fits is not None:
        data = astropy_fits.getdata(path)
        header = dict(astropy_fits.getheader(path))
        if data.ndim != 2:
            raise AlignmentError(f"{path.name}: expected a 2D FITS image, got {data.ndim}D")
        return np.asarray(data, dtype=np.float32), {str(k).upper(): v for k, v in header.items()}
    return read_fits_minimal(path)


def read_timestamp(header: dict[str, Any], filename: str) -> datetime:
    date_value = None
    for key in TIMESTAMP_KEYS:
        if key in header and str(header[key]).strip():
            date_value = str(header[key]).strip()
            break

    if date_value is None:
        raise AlignmentError(f"{filename}: no usable FITS timestamp field found")

    stamp = date_value
    if "T" not in stamp and " " not in stamp:
        for key in TIME_KEYS:
            if key in header and str(header[key]).strip():
                stamp = f"{stamp}T{str(header[key]).strip()}"
                break

    stamp = stamp.strip().replace("Z", "+00:00")
    try:
        parsed = datetime.fromisoformat(stamp)
    except ValueError as exc:
        raise AlignmentError(f"{filename}: cannot parse timestamp {stamp!r}") from exc

    if parsed.tzinfo is not None:
        parsed = parsed.astimezone(timezone.utc).replace(tzinfo=None)
    return parsed


def fits_shape_from_header(header: dict[str, Any], filename: str) -> tuple[int, int]:
    try:
        return int(header["NAXIS2"]), int(header["NAXIS1"])
    except KeyError as exc:
        raise AlignmentError(f"{filename}: missing NAXIS1/NAXIS2 in FITS header") from exc


def collect_frames(input_dir: Path) -> list[FrameInfo]:
    paths = sorted(
        path for path in input_dir.iterdir()
        if path.is_file() and path.suffix.lower() in FITS_SUFFIXES
    )
    if not paths:
        raise AlignmentError(f"No FITS files found in {input_dir}")

    frames: list[FrameInfo] = []
    for path in paths:
        header, _ = read_fits_header(path)
        frames.append(FrameInfo(path, read_timestamp(header, path.name), fits_shape_from_header(header, path.name)))

    frames.sort(key=lambda frame: (frame.timestamp, frame.path.name))
    first_shape = frames[0].shape
    mismatches = [frame.path.name for frame in frames if frame.shape != first_shape]
    if mismatches:
        joined = "\n  ".join(mismatches[:10])
        more = "" if len(mismatches) <= 10 else f"\n  ... and {len(mismatches) - 10} more"
        raise AlignmentError(
            "All FITS images must have the same dimensions. Mismatches:\n"
            f"  {joined}{more}"
        )
    return frames


def resolve_reference(input_dir: Path, name: str) -> Path:
    path = Path(name)
    if not path.is_absolute():
        path = input_dir / path
    if not path.exists():
        raise AlignmentError(f"Reference file does not exist: {path}")
    if path.suffix.lower() not in FITS_SUFFIXES:
        raise AlignmentError(f"Reference file is not a FITS file: {path}")
    return path.resolve()


def normalize_for_detection(data: np.ndarray) -> np.ndarray:
    finite = data[np.isfinite(data)]
    if finite.size == 0:
        raise AlignmentError("Image contains no finite pixel values")
    low, high = np.percentile(finite, (0.5, 99.8))
    if high <= low:
        low, high = float(np.nanmin(finite)), float(np.nanmax(finite))
    if high <= low:
        raise AlignmentError("Image has no usable brightness range")
    normalized = (np.asarray(data, dtype=np.float32) - low) / (high - low)
    return np.clip(normalized, 0.0, 1.0)


def otsu_threshold(image: np.ndarray) -> float:
    hist, edges = np.histogram(image[np.isfinite(image)], bins=256, range=(0.0, 1.0))
    total = hist.sum()
    if total == 0:
        return 0.5
    centers = (edges[:-1] + edges[1:]) * 0.5
    weight_background = np.cumsum(hist)
    weight_foreground = total - weight_background
    mean_background = np.cumsum(hist * centers)
    mean_total = mean_background[-1]
    valid = (weight_background > 0) & (weight_foreground > 0)
    score = np.zeros_like(centers)
    mb = mean_background[valid] / weight_background[valid]
    mf = (mean_total - mean_background[valid]) / weight_foreground[valid]
    score[valid] = weight_background[valid] * weight_foreground[valid] * (mb - mf) ** 2
    return float(centers[int(np.argmax(score))])


def largest_component(mask: np.ndarray) -> np.ndarray:
    labeled, count = ndimage.label(mask)
    if count == 0:
        raise AlignmentError("Solar segmentation failed: no bright component found")
    sizes = np.bincount(labeled.ravel())
    sizes[0] = 0
    label = int(np.argmax(sizes))
    component = labeled == label
    if sizes[label] < max(100, mask.size * 0.001):
        raise AlignmentError("Solar segmentation failed: bright component is too small")
    return component


def segment_solar_region(data: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
    normalized = normalize_for_detection(data)
    blurred = ndimage.gaussian_filter(normalized, sigma=2.0)
    threshold = max(otsu_threshold(blurred), 0.08)
    mask = largest_component(blurred > threshold)
    mask = ndimage.binary_fill_holes(mask)
    mask = ndimage.binary_opening(mask, iterations=1)
    mask = largest_component(mask)
    boundary = mask ^ ndimage.binary_erosion(mask, iterations=1)
    if int(boundary.sum()) < 50:
        raise AlignmentError("Solar segmentation failed: contour is too small")
    return mask, boundary


def fit_circle_least_squares(points_xy: np.ndarray) -> tuple[float, float, float]:
    x = points_xy[:, 0].astype(np.float64)
    y = points_xy[:, 1].astype(np.float64)
    matrix = np.column_stack((x, y, np.ones_like(x)))
    rhs = -(x * x + y * y)
    solution, _, _, _ = np.linalg.lstsq(matrix, rhs, rcond=None)
    a, b, c = solution
    cx = -a / 2.0
    cy = -b / 2.0
    radius_sq = cx * cx + cy * cy - c
    if radius_sq <= 0:
        raise ValueError("invalid circle")
    return float(cx), float(cy), float(math.sqrt(radius_sq))


def circle_from_three_points(sample: np.ndarray) -> tuple[float, float, float]:
    p1, p2, p3 = sample.astype(np.float64)
    temp = p2[0] * p2[0] + p2[1] * p2[1]
    bc = (p1[0] * p1[0] + p1[1] * p1[1] - temp) / 2.0
    cd = (temp - p3[0] * p3[0] - p3[1] * p3[1]) / 2.0
    det = (p1[0] - p2[0]) * (p2[1] - p3[1]) - (p2[0] - p3[0]) * (p1[1] - p2[1])
    if abs(det) < 1e-6:
        raise ValueError("collinear points")
    cx = (bc * (p2[1] - p3[1]) - cd * (p1[1] - p2[1])) / det
    cy = ((p1[0] - p2[0]) * cd - (p2[0] - p3[0]) * bc) / det
    radius = math.hypot(cx - p1[0], cy - p1[1])
    return float(cx), float(cy), float(radius)


def fit_solar_circle(
    boundary: np.ndarray,
    image_shape: tuple[int, int],
    *,
    max_points: int = 12000,
    iterations: int = 900,
) -> tuple[float, float, float, np.ndarray]:
    yx = np.column_stack(np.nonzero(boundary))
    points = np.column_stack((yx[:, 1], yx[:, 0])).astype(np.float64)
    if points.shape[0] < 50:
        raise AlignmentError("Not enough contour points to fit a solar circle")

    rng = random.Random(42)
    if points.shape[0] > max_points:
        indices = rng.sample(range(points.shape[0]), max_points)
        sample_points = points[indices]
    else:
        sample_points = points

    height, width = image_shape
    min_radius = min(width, height) * 0.12
    max_radius = max(width, height) * 0.75
    tolerance = max(2.5, min(width, height) * 0.004)

    best_score = -1
    best_circle: tuple[float, float, float] | None = None
    best_inliers = np.empty((0, 2), dtype=np.float64)
    indices = range(sample_points.shape[0])

    for _ in range(iterations):
        sample = sample_points[rng.sample(indices, 3)]
        try:
            cx, cy, radius = circle_from_three_points(sample)
        except ValueError:
            continue
        if not (min_radius <= radius <= max_radius):
            continue
        if not (-radius <= cx <= width + radius and -radius <= cy <= height + radius):
            continue
        residual = np.abs(np.hypot(sample_points[:, 0] - cx, sample_points[:, 1] - cy) - radius)
        inliers = sample_points[residual <= tolerance]
        score = int(inliers.shape[0])
        if score > best_score:
            best_score = score
            best_circle = (cx, cy, radius)
            best_inliers = inliers

    if best_circle is None or best_inliers.shape[0] < 50:
        try:
            cx, cy, radius = fit_circle_least_squares(sample_points)
        except ValueError as exc:
            raise AlignmentError("Solar circle fitting failed") from exc
        residual = np.abs(np.hypot(points[:, 0] - cx, points[:, 1] - cy) - radius)
        inliers = points[residual <= tolerance * 2.0]
        return cx, cy, radius, inliers

    cx, cy, radius = best_circle
    inliers = best_inliers
    for _ in range(3):
        try:
            cx, cy, radius = fit_circle_least_squares(inliers)
        except ValueError:
            break
        residual = np.abs(np.hypot(points[:, 0] - cx, points[:, 1] - cy) - radius)
        inliers = points[residual <= tolerance * 1.5]
        if inliers.shape[0] < 50:
            break

    return cx, cy, radius, inliers


def fit_center_with_known_radius(
    data: np.ndarray,
    boundary: np.ndarray,
    expected_radius: float,
) -> tuple[float, float, float, np.ndarray]:
    normalized = normalize_for_detection(data)
    blurred = ndimage.gaussian_filter(normalized, sigma=2.0)
    gy, gx = np.gradient(blurred)

    yx = np.column_stack(np.nonzero(boundary))
    points = np.column_stack((yx[:, 1], yx[:, 0])).astype(np.float64)
    if points.shape[0] < 50:
        raise AlignmentError("Not enough contour points to fit a solar circle")

    grad_x = gx[yx[:, 0], yx[:, 1]]
    grad_y = gy[yx[:, 0], yx[:, 1]]
    grad_mag = np.hypot(grad_x, grad_y)
    keep = grad_mag > max(0.002, float(np.percentile(grad_mag, 35)))
    if int(np.count_nonzero(keep)) < 20:
        raise AlignmentError("Known-radius solar fitting failed: weak contour gradient")

    center_votes = np.column_stack(
        (
            points[keep, 0] + grad_x[keep] / grad_mag[keep] * expected_radius,
            points[keep, 1] + grad_y[keep] / grad_mag[keep] * expected_radius,
        )
    )

    height, width = data.shape
    valid = (
        (center_votes[:, 0] >= -expected_radius)
        & (center_votes[:, 0] <= width + expected_radius)
        & (center_votes[:, 1] >= -expected_radius)
        & (center_votes[:, 1] <= height + expected_radius)
    )
    center_votes = center_votes[valid]
    if center_votes.shape[0] < 20:
        raise AlignmentError("Known-radius solar fitting failed: no plausible center votes")

    bin_size = max(1.5, expected_radius * 0.02)
    bins = np.floor(center_votes / bin_size).astype(np.int64)
    unique_bins, counts = np.unique(bins, axis=0, return_counts=True)
    peak_bin = unique_bins[int(np.argmax(counts))]
    near_peak = np.all(np.abs(bins - peak_bin) <= 2, axis=1)
    if int(np.count_nonzero(near_peak)) < 10:
        raise AlignmentError("Known-radius solar fitting failed: center votes are not clustered")

    center = center_votes[near_peak].mean(axis=0)
    tolerance = max(2.5, min(width, height) * 0.004)

    for _ in range(3):
        residual = np.abs(np.hypot(points[:, 0] - center[0], points[:, 1] - center[1]) - expected_radius)
        inliers = points[residual <= tolerance * 1.8]
        if inliers.shape[0] < 20:
            break

        def objective(candidate: np.ndarray) -> np.ndarray:
            return np.hypot(inliers[:, 0] - candidate[0], inliers[:, 1] - candidate[1]) - expected_radius

        result = optimize.least_squares(
            objective,
            center,
            loss="soft_l1",
            f_scale=tolerance,
            max_nfev=100,
        )
        center = result.x

    residual = np.abs(np.hypot(points[:, 0] - center[0], points[:, 1] - center[1]) - expected_radius)
    inliers = points[residual <= tolerance * 1.8]
    if inliers.shape[0] < 20:
        raise AlignmentError("Known-radius solar fitting failed: too few final inliers")
    return float(center[0]), float(center[1]), float(expected_radius), inliers


def detect_solar_disk(data: np.ndarray, expected_radius: float | None = None) -> SolarCircle:
    mask, boundary = segment_solar_region(data)
    if expected_radius is None:
        x, y, radius, inliers = fit_solar_circle(boundary, data.shape)
    else:
        x, y, radius, inliers = fit_center_with_known_radius(data, boundary, expected_radius)
    return SolarCircle(x=x, y=y, radius=radius, mask=mask, boundary=boundary, inlier_points=inliers)


def compute_shift(
    timestamp: datetime,
    first_timestamp: datetime,
    last_timestamp: datetime,
    first_circle: SolarCircle,
    last_circle: SolarCircle,
) -> AlignmentRow:
    total_seconds = (last_timestamp - first_timestamp).total_seconds()
    if total_seconds <= 0:
        raise AlignmentError("Last reference timestamp must be later than first reference timestamp")
    alpha = (timestamp - first_timestamp).total_seconds() / total_seconds
    predicted_x = first_circle.x + alpha * (last_circle.x - first_circle.x)
    predicted_y = first_circle.y + alpha * (last_circle.y - first_circle.y)
    return AlignmentRow(
        filename="",
        timestamp=timestamp,
        alpha=alpha,
        predicted_x=predicted_x,
        predicted_y=predicted_y,
        shift_x=first_circle.x - predicted_x,
        shift_y=first_circle.y - predicted_y,
    )


def shift_image(data: np.ndarray, dx: float, dy: float, order: int) -> np.ndarray:
    if data.ndim == 3:
        shifted_channels = [
            shift_image(data[:, :, channel], dx, dy, order)
            for channel in range(data.shape[2])
        ]
        return np.stack(shifted_channels, axis=2)

    return ndimage.shift(
        data,
        shift=(dy, dx),
        order=order,
        mode="constant",
        cval=0.0,
        prefilter=order > 1,
    )


def parse_crop_factor(value: str) -> float:
    text = str(value).strip()
    if text.endswith("%"):
        text = text[:-1].strip()
        scale = 100.0
    else:
        scale = 1.0

    try:
        factor = float(text) / scale
    except ValueError as exc:
        raise argparse.ArgumentTypeError("--crop must be a percentage like 60% or a factor like 0.6") from exc

    if factor > 1.0:
        factor /= 100.0
    if not (0.0 < factor <= 1.0):
        raise argparse.ArgumentTypeError("--crop must be greater than 0 and less than or equal to 100%")
    return factor


def compute_crop_bounds(
    image_shape: tuple[int, int],
    center_x: float,
    center_y: float,
    crop_factor: float,
) -> tuple[int, int, int, int]:
    height, width = image_shape
    crop_width = max(1, int(round(width * crop_factor)))
    crop_height = max(1, int(round(height * crop_factor)))
    cx = int(round(center_x))
    cy = int(round(center_y))
    x0 = max(0, min(width - crop_width, cx - crop_width // 2))
    y0 = max(0, min(height - crop_height, cy - crop_height // 2))
    return x0, y0, crop_width, crop_height


def crop_image(data: np.ndarray, bounds: tuple[int, int, int, int]) -> np.ndarray:
    x0, y0, width, height = bounds
    if data.ndim == 3:
        return data[y0 : y0 + height, x0 : x0 + width, :]
    return data[y0 : y0 + height, x0 : x0 + width]


def normalize_bayer_pattern(pattern: Any) -> str | None:
    if pattern is None:
        return None
    cleaned = str(pattern).strip().upper()
    return cleaned if cleaned in BAYER_PATTERNS else None


def bayer_pattern_for_frame(header: dict[str, Any], args: argparse.Namespace, filename: str) -> str | None:
    pattern = normalize_bayer_pattern(args.bayer_pattern) if args.bayer_pattern else normalize_bayer_pattern(header.get("BAYERPAT"))
    if args.debayer == "never":
        return None
    if args.debayer == "always" and pattern is None:
        raise AlignmentError(
            f"{filename}: no valid Bayer pattern found. Use --bayer-pattern RGGB/BGGR/GRBG/GBRG."
        )
    return pattern


def debayer_with_opencv(data: np.ndarray, pattern: str) -> np.ndarray | None:
    if cv2 is None:
        return None
    finite = data[np.isfinite(data)]
    if finite.size == 0:
        raise AlignmentError("Cannot debayer image with no finite pixels")
    if float(np.nanmin(finite)) < 0.0 or float(np.nanmax(finite)) > 65535.0:
        return None

    code_name = BAYER_TO_CV2[pattern]
    code = getattr(cv2, code_name)
    data_u16 = np.rint(np.clip(data, 0.0, 65535.0)).astype(np.uint16)
    return cv2.cvtColor(data_u16, code).astype(np.float32)


def debayer_bilinear(data: np.ndarray, pattern: str) -> np.ndarray:
    height, width = data.shape
    masks = {channel: np.zeros((height, width), dtype=np.float32) for channel in "RGB"}
    layout = ((pattern[0], pattern[1]), (pattern[2], pattern[3]))
    for y_parity in (0, 1):
        for x_parity in (0, 1):
            masks[layout[y_parity][x_parity]][y_parity::2, x_parity::2] = 1.0

    kernel = np.array(
        (
            (1.0, 2.0, 1.0),
            (2.0, 4.0, 2.0),
            (1.0, 2.0, 1.0),
        ),
        dtype=np.float32,
    )
    channels: list[np.ndarray] = []
    for channel in "RGB":
        mask = masks[channel]
        known = np.asarray(data, dtype=np.float32) * mask
        numerator = ndimage.convolve(known, kernel, mode="mirror")
        denominator = ndimage.convolve(mask, kernel, mode="mirror")
        plane = np.divide(numerator, denominator, out=np.zeros_like(numerator), where=denominator > 0)
        plane[mask > 0] = data[mask > 0]
        channels.append(plane.astype(np.float32, copy=False))
    return np.stack(channels, axis=2)


def debayer_image(data: np.ndarray, pattern: str) -> np.ndarray:
    debayered = debayer_with_opencv(data, pattern)
    if debayered is not None:
        return debayered
    return debayer_bilinear(data, pattern)


def compute_global_points(frames: Iterable[FrameInfo], black: float | None, white: float | None) -> tuple[float, float]:
    global_min = math.inf
    global_max = -math.inf
    for frame in frames:
        data, _ = read_fits(frame.path)
        finite = data[np.isfinite(data)]
        if finite.size == 0:
            raise AlignmentError(f"{frame.path.name}: no finite pixel values")
        if black is None:
            global_min = min(global_min, float(np.nanmin(finite)))
        if white is None:
            global_max = max(global_max, float(np.nanmax(finite)))

    black_point = float(black) if black is not None else global_min
    white_point = float(white) if white is not None else global_max
    if not white_point > black_point:
        raise AlignmentError("white point must be greater than black point")
    return black_point, white_point


def convert_for_export(data: np.ndarray, black_point: float, white_point: float) -> np.ndarray:
    scaled = (np.asarray(data, dtype=np.float32) - black_point) / (white_point - black_point)
    scaled = np.clip(scaled, 0.0, 1.0)
    return np.rint(scaled * 65535.0).astype(np.uint16)


def save_output(image_u16: np.ndarray, path: Path, fmt: str) -> None:
    if image_u16.ndim == 3:
        if image_u16.shape[2] != 3:
            raise AlignmentError(f"Unsupported color image shape for export: {image_u16.shape}")
        if cv2 is None:
            raise AlignmentError("OpenCV is required to save 16-bit RGB PNG/TIFF output")
        output = image_u16[:, :, ::-1]
        if not cv2.imwrite(str(path), output):
            raise AlignmentError(f"Could not write output image: {path}")
        return

    pil_image = Image.fromarray(image_u16, mode="I;16")
    if fmt == "png":
        pil_image.save(path)
        return
    if fmt == "tiff":
        pil_image.save(path, compression="tiff_deflate")
        return
    raise AlignmentError(f"Unsupported output format: {fmt}")


def save_debug_image(path: Path, data: np.ndarray, circle: SolarCircle) -> None:
    normalized = normalize_for_detection(data)
    base = np.rint(normalized * 255.0).astype(np.uint8)
    rgb = Image.fromarray(base, mode="L").convert("RGB")
    draw = ImageDraw.Draw(rgb)

    yx = np.column_stack(np.nonzero(circle.boundary))
    for y, x in yx[:: max(1, len(yx) // 6000)]:
        rgb.putpixel((int(x), int(y)), (255, 64, 64))

    bbox = [
        circle.x - circle.radius,
        circle.y - circle.radius,
        circle.x + circle.radius,
        circle.y + circle.radius,
    ]
    draw.ellipse(bbox, outline=(64, 255, 96), width=2)
    cx = int(round(circle.x))
    cy = int(round(circle.y))
    draw.line((cx - 12, cy, cx + 12, cy), fill=(64, 160, 255), width=2)
    draw.line((cx, cy - 12, cx, cy + 12), fill=(64, 160, 255), width=2)
    draw.text((12, 12), f"x={circle.x:.2f} y={circle.y:.2f} r={circle.radius:.2f}", fill=(255, 255, 255))
    rgb.save(path)


def draw_cross(draw: ImageDraw.ImageDraw, x: float, y: float, color: tuple[int, int, int], size: int = 12) -> None:
    cx = int(round(x))
    cy = int(round(y))
    draw.line((cx - size, cy, cx + size, cy), fill=color, width=2)
    draw.line((cx, cy - size, cx, cy + size), fill=color, width=2)


def draw_arrow(
    draw: ImageDraw.ImageDraw,
    start: tuple[float, float],
    end: tuple[float, float],
    color: tuple[int, int, int],
) -> None:
    x0, y0 = start
    x1, y1 = end
    draw.line((x0, y0, x1, y1), fill=color, width=3)
    angle = math.atan2(y1 - y0, x1 - x0)
    head_len = 18.0
    head_angle = math.radians(28.0)
    for sign in (-1, 1):
        hx = x1 - head_len * math.cos(angle + sign * head_angle)
        hy = y1 - head_len * math.sin(angle + sign * head_angle)
        draw.line((x1, y1, hx, hy), fill=color, width=3)


def save_debug_drift_image(
    path: Path,
    first_data: np.ndarray,
    last_data: np.ndarray,
    first_circle: SolarCircle,
    last_circle: SolarCircle,
    duration_seconds: float,
) -> None:
    first = normalize_for_detection(first_data)
    last = normalize_for_detection(last_data)
    overlay = np.zeros((*first.shape, 3), dtype=np.uint8)
    overlay[:, :, 1] = np.rint(first * 220.0).astype(np.uint8)
    overlay[:, :, 0] = np.rint(last * 220.0).astype(np.uint8)
    overlay[:, :, 2] = np.rint(last * 220.0).astype(np.uint8)

    image = Image.fromarray(overlay, mode="RGB")
    draw = ImageDraw.Draw(image)

    first_bbox = [
        first_circle.x - first_circle.radius,
        first_circle.y - first_circle.radius,
        first_circle.x + first_circle.radius,
        first_circle.y + first_circle.radius,
    ]
    last_bbox = [
        last_circle.x - last_circle.radius,
        last_circle.y - last_circle.radius,
        last_circle.x + last_circle.radius,
        last_circle.y + last_circle.radius,
    ]
    draw.ellipse(first_bbox, outline=(80, 255, 80), width=2)
    draw.ellipse(last_bbox, outline=(255, 80, 255), width=2)
    draw_cross(draw, first_circle.x, first_circle.y, (80, 255, 80))
    draw_cross(draw, last_circle.x, last_circle.y, (255, 80, 255))
    draw_arrow(draw, (first_circle.x, first_circle.y), (last_circle.x, last_circle.y), (255, 220, 64))

    drift_x = last_circle.x - first_circle.x
    drift_y = last_circle.y - first_circle.y
    drift_per_min_x = drift_x / (duration_seconds / 60.0) if duration_seconds > 0 else 0.0
    drift_per_min_y = drift_y / (duration_seconds / 60.0) if duration_seconds > 0 else 0.0
    text_lines = [
        "First reference: green",
        "Last reference : magenta",
        f"Drift X: {drift_x:.2f} px ({drift_per_min_x:.3f} px/min)",
        f"Drift Y: {drift_y:.2f} px ({drift_per_min_y:.3f} px/min)",
        f"Duration: {duration_seconds:.1f} s",
    ]
    x_text, y_text = 12, 12
    for line in text_lines:
        draw.text((x_text + 1, y_text + 1), line, fill=(0, 0, 0))
        draw.text((x_text, y_text), line, fill=(255, 255, 255))
        y_text += 16

    image.save(path)


def write_csv(path: Path, rows: list[AlignmentRow]) -> None:
    with path.open("w", newline="", encoding="utf-8") as handle:
        writer = csv.writer(handle)
        writer.writerow(("filename", "timestamp", "alpha", "predicted_x", "predicted_y", "shift_x", "shift_y"))
        for row in rows:
            writer.writerow(
                (
                    row.filename,
                    row.timestamp.isoformat(timespec="milliseconds"),
                    f"{row.alpha:.8f}",
                    f"{row.predicted_x:.6f}",
                    f"{row.predicted_y:.6f}",
                    f"{row.shift_x:.6f}",
                    f"{row.shift_y:.6f}",
                )
            )


def print_reference(title: str, frame: FrameInfo, circle: SolarCircle) -> None:
    print(title)
    print(f"  file      : {frame.path.name}")
    print(f"  timestamp : {frame.timestamp.isoformat(timespec='milliseconds')}")
    print(f"  center X  : {circle.x:.3f}")
    print(f"  center Y  : {circle.y:.3f}")
    print(f"  radius    : {circle.radius:.3f}")


def output_extension(fmt: str) -> str:
    return ".tif" if fmt == "tiff" else ".png"


def run(args: argparse.Namespace) -> None:
    input_dir = Path(args.input).resolve()
    output_dir = Path(args.output).resolve()
    if not input_dir.exists() or not input_dir.is_dir():
        raise AlignmentError(f"Input directory does not exist: {input_dir}")

    first_path = resolve_reference(input_dir, args.first)
    last_path = resolve_reference(input_dir, args.last)
    output_dir.mkdir(parents=True, exist_ok=True)

    frames = collect_frames(input_dir)
    by_path = {frame.path.resolve(): frame for frame in frames}
    if first_path not in by_path:
        raise AlignmentError(f"First reference is not in the input directory: {first_path}")
    if last_path not in by_path:
        raise AlignmentError(f"Last reference is not in the input directory: {last_path}")

    first_frame = by_path[first_path]
    last_frame = by_path[last_path]
    if first_frame.timestamp >= last_frame.timestamp:
        raise AlignmentError("First reference must be earlier than last reference")

    first_data, _ = read_fits(first_frame.path)
    last_data, _ = read_fits(last_frame.path)
    first_circle = detect_solar_disk(first_data)
    last_circle_free = detect_solar_disk(last_data)
    free_radius_delta = abs(last_circle_free.radius - first_circle.radius) / first_circle.radius
    if free_radius_delta > 0.05 and args.radius_refit:
        last_circle = detect_solar_disk(last_data, expected_radius=first_circle.radius)
        radius_refit_used = True
    else:
        last_circle = last_circle_free
        radius_refit_used = False

    print_reference("First reference:", first_frame, first_circle)
    print_reference("Last reference:", last_frame, last_circle)

    if free_radius_delta > 0.05:
        print("WARNING: solar radius differs by more than 5 %")
        if radius_refit_used:
            print(
                "WARNING: last reference center was refit with the first reference radius "
                f"({first_circle.radius:.3f} px)"
            )

    total_seconds = (last_frame.timestamp - first_frame.timestamp).total_seconds()
    drift_x = last_circle.x - first_circle.x
    drift_y = last_circle.y - first_circle.y
    print(f"Total drift X : {drift_x:.3f} pixels")
    print(f"Total drift Y : {drift_y:.3f} pixels")
    print(f"Total duration: {total_seconds:.3f} seconds")
    print(f"Drift X       : {drift_x / (total_seconds / 60.0):.6f} pixels/minute")
    print(f"Drift Y       : {drift_y / (total_seconds / 60.0):.6f} pixels/minute")

    if args.debug:
        save_debug_image(input_dir / "debug_first.png", first_data, first_circle)
        save_debug_image(input_dir / "debug_last.png", last_data, last_circle)
        save_debug_drift_image(
            input_dir / "debug_drift.png",
            first_data,
            last_data,
            first_circle,
            last_circle,
            total_seconds,
        )

    first_pattern = bayer_pattern_for_frame(read_fits_header(first_frame.path)[0], args, first_frame.path.name)
    if first_pattern is None:
        print(f"Debayer       : {args.debayer} (grayscale output)")
    else:
        print(f"Debayer       : {args.debayer} ({first_pattern})")

    black_point, white_point = compute_global_points(frames, args.black_point, args.white_point)
    print(f"Black point   : {black_point:.6f}")
    print(f"White point   : {white_point:.6f}")
    crop_bounds = compute_crop_bounds(first_frame.shape, first_circle.x, first_circle.y, args.crop)
    print(
        "Crop          : "
        f"{args.crop * 100.0:.1f}% "
        f"({crop_bounds[2]}x{crop_bounds[3]} at x={crop_bounds[0]} y={crop_bounds[1]})"
    )

    rows: list[AlignmentRow] = []
    digits = max(4, len(str(len(frames))))
    ext = output_extension(args.format)
    for index, frame in enumerate(frames, start=1):
        row = compute_shift(frame.timestamp, first_frame.timestamp, last_frame.timestamp, first_circle, last_circle)
        row = AlignmentRow(
            filename=frame.path.name,
            timestamp=row.timestamp,
            alpha=row.alpha,
            predicted_x=row.predicted_x,
            predicted_y=row.predicted_y,
            shift_x=row.shift_x,
            shift_y=row.shift_y,
        )
        data, header = read_fits(frame.path)
        pattern = bayer_pattern_for_frame(header, args, frame.path.name)
        export_data = debayer_image(data, pattern) if pattern is not None else data
        shifted = shift_image(export_data, row.shift_x, row.shift_y, args.interpolation_order)
        cropped = crop_image(shifted, crop_bounds)
        exported = convert_for_export(cropped, black_point, white_point)
        output_path = output_dir / f"aligned_{index:0{digits}d}{ext}"
        save_output(exported, output_path, args.format)
        rows.append(row)
        print(f"[{index:02d}/{len(frames):02d}] {frame.path.name} -> dx={row.shift_x:.3f} dy={row.shift_y:.3f}")

    write_csv(output_dir / "alignment.csv", rows)


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        description="Recenters a FITS solar eclipse sequence using linear timestamp-based drift."
    )
    parser.add_argument("--input", required=True, help="Directory containing FITS files.")
    parser.add_argument("--first", required=True, help="First reference FITS filename.")
    parser.add_argument("--last", required=True, help="Last reference FITS filename.")
    parser.add_argument("--output", required=True, help="Output directory for aligned images and CSV.")
    parser.add_argument("--black-point", type=float, default=None, help="Fixed black point for all exports.")
    parser.add_argument("--white-point", type=float, default=None, help="Fixed white point for all exports.")
    parser.add_argument(
        "--crop",
        type=parse_crop_factor,
        default=parse_crop_factor("60%"),
        help="Centered crop size after alignment, as a percentage or factor. Default: 60%%.",
    )
    parser.add_argument("--debug", action="store_true", help="Write debug_first.png and debug_last.png.")
    parser.add_argument(
        "--debayer",
        choices=("auto", "always", "never"),
        default="auto",
        help="Debayer FITS images before shifting/export. Default: auto when BAYERPAT exists.",
    )
    parser.add_argument(
        "--bayer-pattern",
        choices=tuple(sorted(BAYER_PATTERNS)),
        default=None,
        help="Override FITS BAYERPAT, e.g. RGGB, BGGR, GRBG, or GBRG.",
    )
    parser.add_argument(
        "--no-radius-refit",
        dest="radius_refit",
        action="store_false",
        help="Do not refit the last reference with the first radius when free radii differ.",
    )
    parser.set_defaults(radius_refit=True)
    parser.add_argument(
        "--format",
        choices=("png", "tiff"),
        default="png",
        help="16-bit output format. Default: png.",
    )
    parser.add_argument(
        "--interpolation-order",
        type=int,
        choices=(0, 1, 2, 3, 4, 5),
        default=3,
        help="scipy.ndimage.shift spline order. 1=bilinear, 3=bicubic-like default.",
    )
    return parser


def main(argv: list[str] | None = None) -> int:
    parser = build_parser()
    args = parser.parse_args(argv)
    try:
        run(args)
    except AlignmentError as exc:
        print(f"ERROR: {exc}", file=sys.stderr)
        return 1
    return 0


if __name__ == "__main__":
    raise SystemExit(main())