High-Performance C Extensions

from typing import Optional, Tuple, Any
import numpy as np

def tdriz(
    input: np.ndarray,
    weights: np.ndarray,
    pixmap: np.ndarray,
    output: np.ndarray,
    counts: np.ndarray,
    context: Optional[np.ndarray],
    uniqid: int,
    xmin: int,
    xmax: int,
    ymin: int,
    ymax: int,
    scale: float,
    pixfrac: float,
    kernel: str,
    in_units: str,
    expscale: float,
    wtscale: float,
    fillstr: str
) -> Tuple[str, float, float]:
    """
    Core drizzling algorithm implementation in C for maximum performance.

    This function directly implements the drizzle algorithm at the C level,
    providing the computational engine for the Python Drizzle class.

    Parameters:
    - input: Input image array (float32, 2D)
    - weights: Pixel weighting array (float32, 2D, same shape as input)
    - pixmap: Coordinate mapping array (float64, 3D with shape (Ny, Nx, 2))
    - output: Output image array (float32, 2D, modified in-place)
    - counts: Output weight array (float32, 2D, same shape as output, modified in-place)
    - context: Context bit-field array (int32, 2D, same shape as output, modified in-place)
               Can be None to disable context tracking
    - uniqid: Unique identifier for this input image (1-based, used for context bit field)
    - xmin: Minimum X coordinate of bounding box in input image (0-based)
    - xmax: Maximum X coordinate of bounding box in input image (0-based)
    - ymin: Minimum Y coordinate of bounding box in input image (0-based)
    - ymax: Maximum Y coordinate of bounding box in input image (0-based)
    - scale: Pixel scale factor for flux scaling
    - pixfrac: Pixel fraction (0-1) controlling flux distribution area
    - kernel: Drizzling kernel name ("square", "gaussian", "point", "turbo", "lanczos2", "lanczos3")
    - in_units: Input units ("cps" for counts per second, "counts" for total counts)
    - expscale: Exposure scaling factor
    - wtscale: Weight scaling factor
    - fillstr: Fill value string for unsampled pixels ("INDEF", "NaN", or numeric string)

    Returns:
    Tuple of (version_string, nmiss, nskip):
    - version_string: C library version identifier
    - nmiss: Number of input pixels that did not contribute to output (float)
    - nskip: Number of input lines that were skipped (float)

    Notes:
    - All arrays must be C-contiguous for optimal performance
    - Output, counts, and context arrays are modified in-place
    - Context tracking uses bit fields; bit (uniqid-1) is set for contributing pixels
    - Kernel choice affects flux distribution and conservation properties
    - Performance-critical function optimized for large astronomical datasets
    """

def tblot(
    image: np.ndarray,
    pixmap: np.ndarray,
    output: np.ndarray,
    scale: float,
    kscale: float,
    interp: str,
    exptime: float,
    misval: float,
    sinscl: float
) -> None:
    """
    Core blotting algorithm implementation in C.

    Performs inverse drizzling by resampling input image onto output grid
    using interpolation methods. Optimized for performance with large images.

    Parameters:
    - image: Input image to be blotted (float32, 2D)
    - pixmap: Coordinate mapping array (float64, 3D with shape (Ny, Nx, 2))
              Maps from output pixel coordinates to input pixel coordinates
    - output: Output image array (float32, 2D, modified in-place)
              Must be pre-allocated with correct size
    - scale: Pixel ratio scale factor
    - kscale: Kernel scale factor (typically 1.0)
    - interp: Interpolation method string:
              "nearest" - nearest neighbor
              "linear" - bilinear interpolation
              "poly3" - cubic polynomial
              "poly5" - quintic polynomial
              "sinc" - sinc interpolation
              "lan3" - 3rd order Lanczos
              "lan5" - 5th order Lanczos
    - exptime: Exposure time for flux scaling
    - misval: Value to use for missing/invalid pixels
    - sinscl: Scaling factor for sinc interpolation (used when interp="sinc")

    Returns:
    None (modifies output array in-place)

    Notes:
    - Output array must be pre-allocated and zeroed
    - Not flux-conserving (unlike tdriz)
    - Best used with well-sampled input images
    - pixmap maps from output coordinates to input coordinates (inverse of tdriz)
    - Performance optimized for large astronomical image processing
    """

def test_cdrizzle(
    data: np.ndarray,
    weights: np.ndarray,
    pixmap: np.ndarray,
    output_data: np.ndarray,
    output_counts: np.ndarray,
    output_context: np.ndarray
) -> None:
    """
    Test function for C drizzling implementation.

    Runs C-level unit tests for the cdrizzle implementation, providing validation
    and performance testing of the core drizzling algorithms.

    Parameters:
    - data: Input test data array (float32, 2D)
    - weights: Input weight array (float32, 2D, same shape as data)
    - pixmap: Pixel mapping array (float64, 3D with shape (height, width, 2))
    - output_data: Output data array (float32, 2D, modified in-place)
    - output_counts: Output counts array (float32, 2D, same shape as output_data, modified in-place)
    - output_context: Context bit-field array (int32, 2D, same shape as output_data, modified in-place)

    Returns:
    None

    Notes:
    - Primarily for internal testing and validation
    - Calls internal C function utest_cdrizzle() for validation
    - All arrays must have correct dtypes and be C-contiguous
    - Not typically used in production workflows
    """

def invert_pixmap(
    pixmap: np.ndarray,
    xyout: np.ndarray,
    bbox: Optional[np.ndarray]
) -> np.ndarray:
    """
    Invert pixel mapping transformation using golden ratio search.

    Takes output coordinates and finds corresponding input coordinates by
    iteratively searching through the pixmap transformation array.

    Parameters:
    - pixmap: 3D coordinate mapping array (float64, shape (height, width, 2))
              Contains forward transformation from input to output coordinates
    - xyout: 1D array (float64, 2 elements) containing output coordinates [x_out, y_out]
    - bbox: 2D bounding box array (float64, shape (2, 2)) as [[xmin, xmax], [ymin, ymax]]
            or None to use full pixmap extent

    Returns:
    1D array (float64, 2 elements) containing inverted input coordinates [x_in, y_in]

    Notes:
    - Used internally for coordinate system conversions
    - Helps convert between drizzling and blotting coordinate mappings
    - Uses golden ratio search algorithm for iterative inversion
    - Essential for computing reverse transformations
    """

def clip_polygon(
    p: np.ndarray,
    q: np.ndarray
) -> List[Tuple[float, float]]:
    """
    Clip polygon p to the window defined by polygon q.

    Implements the Sutherland-Hodgman polygon clipping algorithm for geometric
    operations on pixel regions and image boundaries.

    Parameters:
    - p: 2D array (float64, shape (n_vertices, 2)) representing first polygon vertices
         Each row contains [x, y] coordinates of a vertex
    - q: 2D array (float64, shape (m_vertices, 2)) representing clipping window polygon
         Each row contains [x, y] coordinates of a clipping window vertex

    Returns:
    List of tuples, each containing (x, y) coordinates as floats representing
    vertices of the clipped polygon

    Notes:
    - Both input polygons are automatically oriented counter-clockwise before clipping
    - Uses Sutherland-Hodgman clipping algorithm for robust polygon intersection
    - Primarily used for geometric operations on image pixel regions
    - Returns empty list if polygons do not intersect
    """

import numpy as np
from drizzle import cdrizzle

# Prepare input arrays
input_data = np.random.random((100, 100)).astype(np.float32)
weights = np.ones((100, 100), dtype=np.float32)

# Create pixel mapping (identity transformation for example)
y, x = np.indices((100, 100), dtype=np.float64)
pixmap = np.dstack([x, y])

# Prepare output arrays
output_shape = (120, 120)
output_data = np.zeros(output_shape, dtype=np.float32)
output_counts = np.zeros(output_shape, dtype=np.float32)
context = np.zeros(output_shape, dtype=np.int32)

# Call C drizzling function directly
version, nmiss, nskip = cdrizzle.tdriz(
    input=input_data,
    weights=weights,
    pixmap=pixmap,
    output=output_data,
    counts=output_counts,
    context=context,
    uniqid=1,  # First image
    xmin=0,
    xmax=99,
    ymin=0,
    ymax=99,
    scale=1.0,
    pixfrac=1.0,
    kernel="square",
    in_units="cps",
    expscale=1.0,
    wtscale=1.0,
    fillstr="INDEF"
)

print(f"C library version: {version}")
print(f"Missed pixels: {nmiss}, Skipped lines: {nskip}")
print(f"Output data range: {output_data.min():.3f} to {output_data.max():.3f}")

import numpy as np
from drizzle import cdrizzle

# Create well-sampled input image (e.g., drizzled result)
input_image = np.random.random((150, 150)).astype(np.float32)

# Create coordinate mapping from output to input
output_shape = (100, 100)
y, x = np.indices(output_shape, dtype=np.float64)

# Scale coordinates to map to input image
x_scaled = x * 1.5  # Scale factor
y_scaled = y * 1.5
pixmap = np.dstack([x_scaled, y_scaled])

# Prepare output array
blotted_image = np.zeros(output_shape, dtype=np.float32)

# Call C blotting function directly
cdrizzle.tblot(
    image=input_image,
    pixmap=pixmap,
    output=blotted_image,
    scale=1.5,
    kscale=1.0,
    interp="poly5",
    exptime=1.0,
    misval=0.0,
    sinscl=1.0
)

print(f"Blotted image range: {blotted_image.min():.3f} to {blotted_image.max():.3f}")

import numpy as np
import time
from drizzle import cdrizzle

# Create large test arrays for performance testing
size = 2048
input_data = np.random.random((size, size)).astype(np.float32)
weights = np.ones_like(input_data)

# Create pixel mapping
y, x = np.indices((size, size), dtype=np.float64)
pixmap = np.dstack([x, y])

# Prepare outputs
output_data = np.zeros_like(input_data)
output_counts = np.zeros_like(input_data)

# Time the C function
start_time = time.time()

version, nmiss, nskip = cdrizzle.tdriz(
    input=input_data,
    weights=weights,
    pixmap=pixmap,
    output=output_data,
    counts=output_counts,
    context=None,  # Disable context for speed
    uniqid=1,
    xmin=0, xmax=size-1,
    ymin=0, ymax=size-1,
    scale=1.0,
    pixfrac=1.0,
    kernel="square",
    in_units="cps",
    expscale=1.0,
    wtscale=1.0,
    fillstr="0"
)

elapsed = time.time() - start_time
pixels_per_sec = (size * size) / elapsed

print(f"Processed {size}x{size} image in {elapsed:.2f} seconds")
print(f"Performance: {pixels_per_sec/1e6:.2f} Mpixels/sec")

import numpy as np
from drizzle import cdrizzle

def test_kernel(kernel_name, input_data, pixmap):
    """Test different drizzling kernels."""
    output_shape = (150, 150)
    output_data = np.zeros(output_shape, dtype=np.float32)
    output_counts = np.zeros(output_shape, dtype=np.float32)
    weights = np.ones_like(input_data)

    version, nmiss, nskip = cdrizzle.tdriz(
        input=input_data,
        weights=weights,
        pixmap=pixmap,
        output=output_data,
        counts=output_counts,
        context=None,
        uniqid=1,
        xmin=0, xmax=input_data.shape[1]-1,
        ymin=0, ymax=input_data.shape[0]-1,
        scale=1.0,
        pixfrac=1.0,
        kernel=kernel_name,
        in_units="cps",
        expscale=1.0,
        wtscale=1.0,
        fillstr="0"
    )

    return output_data, nmiss, nskip

# Test data
input_data = np.random.random((100, 100)).astype(np.float32)
y, x = np.indices((100, 100), dtype=np.float64)
pixmap = np.dstack([x * 1.2, y * 1.2])  # Slight magnification

# Test different kernels
kernels = ["square", "gaussian", "point", "turbo", "lanczos2", "lanczos3"]

for kernel in kernels:
    try:
        result, nmiss, nskip = test_kernel(kernel, input_data, pixmap)
        print(f"Kernel '{kernel}': output range [{result.min():.3f}, {result.max():.3f}], "
              f"missed={nmiss}, skipped={nskip}")
    except Exception as e:
        print(f"Kernel '{kernel}' failed: {e}")

import numpy as np
from drizzle import cdrizzle

# Handle array type mismatches
try:
    bad_input = np.ones((50, 50), dtype=np.int32)  # Wrong dtype
    weights = np.ones((50, 50), dtype=np.float32)
    y, x = np.indices((50, 50), dtype=np.float64)
    pixmap = np.dstack([x, y])
    output = np.zeros((50, 50), dtype=np.float32)
    counts = np.zeros((50, 50), dtype=np.float32)

    cdrizzle.tdriz(
        input=bad_input,  # int32 instead of float32
        weights=weights,
        pixmap=pixmap,
        output=output,
        counts=counts,
        context=None,
        uniqid=1,
        xmin=0, xmax=49, ymin=0, ymax=49,
        scale=1.0, pixfrac=1.0,
        kernel="square", in_units="cps",
        expscale=1.0, wtscale=1.0, fillstr="0"
    )
except (TypeError, ValueError) as e:
    print(f"Array type error: {e}")

# Handle invalid kernel names
try:
    input_data = np.ones((50, 50), dtype=np.float32)
    # ... other arrays setup ...
    cdrizzle.tdriz(
        input=input_data, weights=weights, pixmap=pixmap,
        output=output, counts=counts, context=None,
        uniqid=1, xmin=0, xmax=49, ymin=0, ymax=49,
        scale=1.0, pixfrac=1.0,
        kernel="invalid_kernel",  # Bad kernel name
        in_units="cps", expscale=1.0, wtscale=1.0, fillstr="0"
    )
except Exception as e:
    print(f"Invalid kernel error: {e}")

# Handle array shape mismatches
try:
    input_data = np.ones((50, 50), dtype=np.float32)
    bad_weights = np.ones((40, 40), dtype=np.float32)  # Wrong shape
    # ... rest of setup ...
    cdrizzle.tdriz(
        input=input_data,
        weights=bad_weights,  # Shape mismatch
        # ... other parameters ...
    )
except Exception as e:
    print(f"Shape mismatch error: {e}")