Coordinate Transformation Utilities

from typing import Optional, Union, List, Tuple
import numpy as np

def calc_pixmap(
    wcs_from,
    wcs_to,
    shape: Optional[Tuple[int, int]] = None,
    disable_bbox: str = "to"
) -> np.ndarray:
    """
    Calculate pixel-to-pixel mapping between two coordinate systems.

    Transforms pixel coordinates from the 'from' image to pixel coordinates
    in the 'to' image using their respective WCS objects.

    Parameters:
    - wcs_from: WCS object for source coordinate system
                Must have array_shape property or bounding_box if shape not provided
    - wcs_to: WCS object for destination coordinate system
    - shape: Output array shape (Ny, Nx) in numpy order
             If None, uses wcs_from.array_shape or derives from bounding_box
    - disable_bbox: Controls bounding box usage ("to", "from", "both", "none")
                   "to" disables bounding box for destination WCS
                   "from" disables bounding box for source WCS
                   "both" disables both, "none" uses both

    Returns:
    3D numpy array of shape (Ny, Nx, 2) where:
    - [..., 0] contains X-coordinates in destination system
    - [..., 1] contains Y-coordinates in destination system

    Raises:
    - ValueError: If output shape cannot be determined from inputs

    Notes:
    - Output frames of both WCS objects must have same units
    - Coordinates outside bounding boxes may be set to NaN depending on disable_bbox
    - Used to generate pixmap arrays for Drizzle.add_image()
    """

def estimate_pixel_scale_ratio(
    wcs_from,
    wcs_to,
    refpix_from: Optional[Union[np.ndarray, Tuple, List]] = None,
    refpix_to: Optional[Union[np.ndarray, Tuple, List]] = None
) -> float:
    """
    Compute ratio of pixel scales between two WCS objects.

    Pixel scale is estimated as square root of pixel area (assumes square pixels).
    If reference pixel positions are not provided, uses center of bounding box,
    center of pixel_shape, or (0, 0) as fallback.

    Parameters:
    - wcs_from: Source WCS object (must have pixel_shape property)
    - wcs_to: Destination WCS object
    - refpix_from: Reference pixel coordinates in source image for scale calculation
                   Can be numpy array, tuple, or list of coordinates
    - refpix_to: Reference pixel coordinates in destination image for scale calculation
                 Can be numpy array, tuple, or list of coordinates

    Returns:
    Float representing ratio of destination to source pixel scales
    (pixel_scale_to / pixel_scale_from)

    Notes:
    - Useful for determining scale parameter in drizzling operations
    - Reference pixels default to image centers when not specified
    - Uses approximate algorithm for efficiency
    """

import numpy as np
from drizzle.utils import calc_pixmap
from astropy import wcs

# Create sample WCS objects (in real usage, these come from FITS headers)
# Here we simulate simple WCS objects
class SimpleWCS:
    def __init__(self, pixel_shape, scale=1.0, offset=(0, 0)):
        self.array_shape = pixel_shape
        self.scale = scale
        self.offset = offset

    def pixel_to_world_values(self, x, y):
        return x * self.scale + self.offset[0], y * self.scale + self.offset[1]

    def world_to_pixel_values(self, world_x, world_y):
        return (world_x - self.offset[0]) / self.scale, (world_y - self.offset[1]) / self.scale

# Source image: 100x100 pixels
wcs_from = SimpleWCS(pixel_shape=(100, 100), scale=1.0, offset=(0, 0))

# Destination: 150x150 pixels with different scale and offset
wcs_to = SimpleWCS(pixel_shape=(150, 150), scale=0.8, offset=(10, 5))

# Calculate pixel mapping
pixmap = calc_pixmap(wcs_from, wcs_to)

print(f"Pixel map shape: {pixmap.shape}")  # (100, 100, 2)
print(f"Sample mapping - pixel (50,50) maps to: ({pixmap[50,50,0]:.2f}, {pixmap[50,50,1]:.2f})")

from drizzle.utils import calc_pixmap

# Calculate mapping with custom output shape
pixmap = calc_pixmap(
    wcs_from,
    wcs_to,
    shape=(200, 180),  # Custom shape instead of using wcs_from.array_shape
    disable_bbox="both"  # Disable both bounding boxes
)

print(f"Custom shape mapping: {pixmap.shape}")

# Different bounding box configurations
pixmap_from_disabled = calc_pixmap(wcs_from, wcs_to, disable_bbox="from")
pixmap_to_disabled = calc_pixmap(wcs_from, wcs_to, disable_bbox="to")  # Default
pixmap_none_disabled = calc_pixmap(wcs_from, wcs_to, disable_bbox="none")

import numpy as np
from drizzle.utils import estimate_pixel_scale_ratio

# Calculate pixel scale ratio between two coordinate systems
ratio = estimate_pixel_scale_ratio(wcs_from, wcs_to)
print(f"Pixel scale ratio (to/from): {ratio:.4f}")

# Calculate ratio at specific reference points
ref_from = (25, 25)  # Reference point in source image
ref_to = (30, 28)    # Reference point in destination image

ratio_at_refs = estimate_pixel_scale_ratio(
    wcs_from,
    wcs_to,
    refpix_from=ref_from,
    refpix_to=ref_to
)
print(f"Ratio at reference points: {ratio_at_refs:.4f}")

# Using numpy arrays for reference points
ref_from_array = np.array([25.5, 25.5])
ref_to_array = np.array([30.2, 28.1])

ratio_arrays = estimate_pixel_scale_ratio(
    wcs_from,
    wcs_to,
    refpix_from=ref_from_array,
    refpix_to=ref_to_array
)
print(f"Ratio with array inputs: {ratio_arrays:.4f}")

import numpy as np
from drizzle.resample import Drizzle
from drizzle.utils import calc_pixmap, estimate_pixel_scale_ratio

# Setup coordinate systems
wcs_input = SimpleWCS(pixel_shape=(120, 120), scale=1.0)
wcs_output = SimpleWCS(pixel_shape=(200, 200), scale=0.75)

# Calculate transformations
pixmap = calc_pixmap(wcs_input, wcs_output)
scale_ratio = estimate_pixel_scale_ratio(wcs_input, wcs_output)

# Create sample data
data = np.random.random((120, 120)).astype(np.float32)

# Use calculated parameters in drizzling
drizzler = Drizzle(out_shape=(200, 200))
nmiss, nskip = drizzler.add_image(
    data=data,
    exptime=10.0,
    pixmap=pixmap,
    scale=scale_ratio  # Use calculated scale ratio
)

print(f"Drizzling completed: missed={nmiss}, skipped={nskip}")

# Example with actual astropy WCS objects (requires astropy)
try:
    from astropy import wcs
    from astropy.io import fits

    # Load WCS from FITS headers (example)
    # header1 = fits.getheader('input_image.fits')
    # header2 = fits.getheader('output_image.fits')
    # wcs1 = wcs.WCS(header1)
    # wcs2 = wcs.WCS(header2)

    # For demonstration, create simple WCS
    w1 = wcs.WCS(naxis=2)
    w1.wcs.crpix = [50, 50]
    w1.wcs.cdelt = [0.1, 0.1]
    w1.wcs.crval = [0, 0]
    w1.array_shape = (100, 100)

    w2 = wcs.WCS(naxis=2)
    w2.wcs.crpix = [75, 75]
    w2.wcs.cdelt = [0.08, 0.08]
    w2.wcs.crval = [1, 1]
    w2.array_shape = (150, 150)

    # Calculate pixel mapping
    pixmap = calc_pixmap(w1, w2)
    scale_ratio = estimate_pixel_scale_ratio(w1, w2)

    print(f"WCS pixel map shape: {pixmap.shape}")
    print(f"Scale ratio: {scale_ratio:.4f}")

except ImportError:
    print("Astropy not available for this example")

from drizzle.utils import calc_pixmap, estimate_pixel_scale_ratio

# Handle cases where shape cannot be determined
class BadWCS:
    def __init__(self):
        self.array_shape = None
        # No bounding_box property either

try:
    bad_wcs = BadWCS()
    good_wcs = SimpleWCS((100, 100))
    pixmap = calc_pixmap(bad_wcs, good_wcs)  # No shape info available
except ValueError as e:
    print(f"Shape determination error: {e}")

# Handle coordinate transformation errors
try:
    # WCS objects that don't implement required methods
    class IncompleteWCS:
        array_shape = (50, 50)
        # Missing required transformation methods

    incomplete_wcs = IncompleteWCS()
    pixmap = calc_pixmap(incomplete_wcs, good_wcs)
except AttributeError as e:
    print(f"WCS method error: {e}")