python_reduction.py

import os
import shutil

import numpy as np
from astropy.io import fits
from astropy.stats import sigma_clip, sigma_clipped_stats
from scipy import ndimage
from skimage.registration import phase_cross_correlation

REGISTRATION_BORDERS = {
    "H": 150,
    "J": 150,
}


def ensure_fits_path(path):
    if os.path.exists(path):
        return path
    if not os.path.splitext(path)[1]:
        alt = f"{path}.fits"
        if os.path.exists(alt):
            return alt
    return path


def read_list_file(path):
    with open(path) as handle:
        return [line.strip() for line in handle if line.strip()]


def resolve_paths(base_dir, items):
    return [ensure_fits_path(os.path.abspath(os.path.join(base_dir, item))) for item in items]


def read_fits(path):
    with fits.open(path) as hdul:
        data = np.asarray(hdul[0].data, dtype=np.float64)
        header = hdul[0].header.copy()
    return data, header


def write_fits(path, data, header=None):
    fits.PrimaryHDU(np.asarray(data, dtype=np.float32), header=header).writeto(
        path, overwrite=True, output_verify="silentfix+ignore"
    )


def load_bad_pixel_mask(path):
    mask_path = ensure_fits_path(path)
    data, _ = read_fits(mask_path)
    return np.asarray(data != 0, dtype=bool)


def fill_bad_pixels(data, bad_mask):
    if bad_mask is None:
        return data.copy()

    mask = np.asarray(bad_mask, dtype=bool)
    if mask.shape != data.shape:
        raise RuntimeError(
            f"Bad pixel mask shape {mask.shape} does not match image shape {data.shape}."
        )
    if not np.any(mask):
        return data.copy()
    if np.all(mask):
        raise RuntimeError("Bad pixel mask marks every pixel as bad.")

    _, indices = ndimage.distance_transform_edt(mask, return_indices=True)
    filled = data.copy()
    filled[mask] = data[tuple(indices[:, mask])]
    return filled


def estimate_mode(data):
    mean, median, _ = sigma_clipped_stats(data, sigma=3.0)
    mode = 3.0 * median - 2.0 * mean
    if not np.isfinite(mode):
        return float(np.nanmedian(data))
    return float(mode)


def combine_median(paths):
    stack = [read_fits(path)[0] for path in paths]
    return np.median(np.stack(stack, axis=0), axis=0)


def normalize_flat_frame(flat_frame):
    valid = np.isfinite(flat_frame) & (flat_frame > 0)
    if not np.any(valid):
        raise RuntimeError("Flat frame has no positive finite pixels for normalization.")
    scale = float(np.median(flat_frame[valid]))
    if scale == 0 or not np.isfinite(scale):
        raise RuntimeError("Flat frame normalization scale is invalid.")
    return flat_frame / scale


def combine_scaled_sigma_clip(images, sigma_value=1.0):
    if len(images) == 1:
        return images[0].copy()

    modes = np.array([estimate_mode(image) for image in images], dtype=float)
    valid = np.isfinite(modes) & (np.abs(modes) > 0)
    target_mode = np.median(modes[valid]) if np.any(valid) else 1.0

    scaled = []
    for image, mode in zip(images, modes):
        if np.isfinite(mode) and mode != 0:
            scaled.append(image * (target_mode / mode))
        else:
            scaled.append(image.copy())

    stack = np.stack(scaled, axis=0)
    if len(images) < 3:
        return np.median(stack, axis=0)

    clipped = sigma_clip(
        stack,
        sigma_lower=sigma_value,
        sigma_upper=sigma_value,
        axis=0,
        masked=True,
    )
    combined = np.ma.median(clipped, axis=0)
    if np.ma.isMaskedArray(combined):
        combined = combined.filled(np.nanmedian(stack, axis=0))
    return np.asarray(combined, dtype=np.float64)


def combine_sum_minmax(images, nlow=1, nhigh=1):
    stack = np.stack(images, axis=0)
    if len(images) <= nlow + nhigh:
        return np.sum(stack, axis=0)

    sorted_stack = np.sort(stack, axis=0)
    trimmed = sorted_stack[nlow : len(images) - nhigh]
    return np.sum(trimmed, axis=0)


def simple_cosmic_ray_clean(data, sigma_threshold=5.0, filter_size=5):
    median_image = ndimage.median_filter(data, size=filter_size, mode="nearest")
    residual = data - median_image
    sigma = 1.4826 * np.nanmedian(np.abs(residual - np.nanmedian(residual)))
    if not np.isfinite(sigma) or sigma == 0:
        return data.copy()
    cosmic_mask = residual > sigma_threshold * sigma
    cleaned = data.copy()
    cleaned[cosmic_mask] = median_image[cosmic_mask]
    return cleaned


def parse_iraf_section(section, shape):
    text = section.strip().strip("[]")
    xpart, ypart = text.split(",")
    x1, x2 = [int(value) for value in xpart.split(":")]
    y1, y2 = [int(value) for value in ypart.split(":")]
    ny, nx = shape
    x1 = max(1, min(nx, x1))
    x2 = max(1, min(nx, x2))
    y1 = max(1, min(ny, y1))
    y2 = max(1, min(ny, y2))
    xslice = slice(min(x1, x2) - 1, max(x1, x2))
    yslice = slice(min(y1, y2) - 1, max(y1, y2))
    return yslice, xslice


def build_registration_section(shape, band, border=None):
    ny, nx = shape
    margin = REGISTRATION_BORDERS.get(band, 150) if border is None else border
    margin = max(0, int(margin))
    x1 = margin + 1
    x2 = max(x1, nx - margin)
    y1 = margin + 1
    y2 = max(y1, ny - margin)
    return f"[{x1}:{x2},{y1}:{y2}]"


def register_images(images, section):
    reference = images[0]
    yslice, xslice = parse_iraf_section(section, reference.shape)
    reference_crop = reference[yslice, xslice]

    shifts = [(0.0, 0.0)]
    aligned = [reference.copy()]

    for image in images[1:]:
        moving_crop = image[yslice, xslice]
        shift, _, _ = phase_cross_correlation(
            reference_crop, moving_crop, upsample_factor=10
        )
        aligned_image = ndimage.shift(
            image,
            shift=shift,
            order=1,
            mode="constant",
            cval=0.0,
            prefilter=False,
        )
        shifts.append((float(shift[1]), float(shift[0])))
        aligned.append(aligned_image)
    return aligned, shifts


def median_of_medians(images):
    return float(np.median([np.median(image) for image in images]))


def sum_exptime(headers):
    return float(sum(float(header.get("EXPTIME", 0.0)) for header in headers))


def write_shift_file(path, shifts):
    with open(path, "w") as handle:
        for xshift, yshift in shifts:
            handle.write(f"{xshift:.3f} {yshift:.3f}\n")


def build_master_frame(list_path, output_path, bad_mask, normalize=False):
    source_paths = resolve_paths(os.path.dirname(list_path), read_list_file(list_path))
    master = combine_median(source_paths)
    master = fill_bad_pixels(master, bad_mask)
    if normalize:
        master = normalize_flat_frame(master)
    _, header = read_fits(source_paths[0])
    write_fits(output_path, master, header)
    return master


def correct_science_frame(data, dark_frame, flat_frame):
    corrected = data - dark_frame
    if flat_frame is not None:
        safe_flat = flat_frame.copy()
        safe_flat[~np.isfinite(safe_flat)] = 1.0
        safe_flat[np.abs(safe_flat) < 1e-12] = 1.0
        corrected = corrected / safe_flat
    return corrected


def write_final_products(sn_folder, band, aligned_images, headers, imed, date, shifts):
    obj_names = read_list_file(os.path.join(sn_folder, f"listobj{band}"))
    write_shift_file(os.path.join(sn_folder, f"shifts{band}.txt"), shifts)
    for obj_name, image, header in zip(obj_names, aligned_images, headers):
        write_fits(os.path.join(sn_folder, f"x{obj_name}.fits"), image, header)

    final_nsc = combine_sum_minmax(aligned_images, nlow=1, nhigh=1)
    write_fits(os.path.join(sn_folder, f"final{band}nsc.fits"), final_nsc, headers[0])

    med_final = float(np.median(final_nsc))
    final_image = final_nsc + (imed - med_final)
    final_header = headers[0].copy()
    final_header["EXPTIME"] = sum_exptime(headers)
    final_path = os.path.join(sn_folder, f"final{band}.fits")
    write_fits(final_path, final_image, final_header)

    sn_name = os.path.basename(sn_folder)
    root_output = os.path.join(
        os.path.dirname(sn_folder), f"SN{sn_name}_{date}_{band}.fits"
    )
    shutil.copyfile(final_path, root_output)

    return {
        "imed": imed,
        "med_final": med_final,
        "final_path": final_path,
        "root_output": root_output,
        "exptime": final_header["EXPTIME"],
    }


def process_band(sn_folder, band, bad_mask, dark_frame, flat_frame, date):
    listorig_path = os.path.join(sn_folder, f"listorig{band}")
    listobj_path = os.path.join(sn_folder, f"listobj{band}")
    if not os.path.exists(listorig_path) or not os.path.exists(listobj_path):
        return None

    orig_paths = resolve_paths(sn_folder, read_list_file(listorig_path))
    obj_names = read_list_file(listobj_path)

    raw_images = []
    headers = []
    for orig_path, obj_name in zip(orig_paths, obj_names):
        data, header = read_fits(orig_path)
        copied_path = os.path.join(sn_folder, f"c{obj_name}.fits")
        write_fits(copied_path, data, header)
        fixed = fill_bad_pixels(data, bad_mask)
        corrected = correct_science_frame(fixed, dark_frame, flat_frame)
        raw_images.append(corrected)
        headers.append(header)
        write_fits(os.path.join(sn_folder, f"fdc{obj_name}.fits"), corrected, header)

    imed = median_of_medians(raw_images)

    sky_images = []
    for index in range(len(raw_images)):
        other_images = [image for j, image in enumerate(raw_images) if j != index]
        sky = combine_scaled_sigma_clip(other_images, sigma_value=1.0)
        sky_images.append(sky)
        write_fits(os.path.join(sn_folder, f"sky{band}{index + 1}.fits"), sky, headers[index])

    obj_images = []
    for obj_name, image, sky, header in zip(obj_names, raw_images, sky_images, headers):
        sky_subtracted = image - sky
        threshold = float(np.mean(sky_subtracted) - 1000.0)
        sky_subtracted = sky_subtracted.copy()
        sky_subtracted[sky_subtracted < threshold] = threshold
        obj_images.append(sky_subtracted)
        write_fits(os.path.join(sn_folder, f"{obj_name}.fits"), sky_subtracted, header)

    cr_images = []
    for obj_name, image, header in zip(obj_names, obj_images, headers):
        cleaned = simple_cosmic_ray_clean(image)
        cr_images.append(cleaned)
        write_fits(os.path.join(sn_folder, f"cr{obj_name}.fits"), cleaned, header)

    section = build_registration_section(cr_images[0].shape, band)
    aligned_images, shifts = register_images(cr_images, section=section)
    return write_final_products(sn_folder, band, aligned_images, headers, imed, date, shifts)


def rerun_alignment_band(sn_folder, band, date):
    listobj_path = os.path.join(sn_folder, f"listobj{band}")
    if not os.path.exists(listobj_path):
        return None

    obj_names = read_list_file(listobj_path)
    cr_paths = [ensure_fits_path(os.path.join(sn_folder, f"cr{name}")) for name in obj_names]
    if not all(os.path.exists(path) for path in cr_paths):
        raise RuntimeError(
            f"Cannot resume alignment for {os.path.basename(sn_folder)} {band}: missing crobj files."
        )

    aligned_images = []
    headers = []
    for path in cr_paths:
        data, header = read_fits(path)
        aligned_images.append(data)
        headers.append(header)

    section = build_registration_section(aligned_images[0].shape, band)
    aligned_images, shifts = register_images(aligned_images, section=section)

    fdc_paths = [ensure_fits_path(os.path.join(sn_folder, f"fdc{name}")) for name in obj_names]
    if not all(os.path.exists(path) for path in fdc_paths):
        raise RuntimeError(
            f"Cannot resume alignment for {os.path.basename(sn_folder)} {band}: missing fdcobj files."
        )
    fdc_images = [read_fits(path)[0] for path in fdc_paths]
    imed = median_of_medians(fdc_images)
    return write_final_products(sn_folder, band, aligned_images, headers, imed, date, shifts)


def run_python_reduction(output_folder):
    output_folder = os.path.abspath(output_folder)
    if not os.path.isdir(output_folder):
        raise RuntimeError(f"Output folder does not exist: {output_folder}")

    mmdd = os.path.basename(output_folder)
    parent = os.path.dirname(output_folder)
    bpm_path = os.path.join(parent, "bpm2012.fits")
    bad_mask = load_bad_pixel_mask(bpm_path)

    listdark_path = os.path.join(output_folder, "listdark")
    if not os.path.exists(listdark_path):
        raise RuntimeError(f"Missing listdark in {output_folder}")

    print(f"Building master dark in {output_folder}")
    dark_frame = build_master_frame(
        listdark_path, os.path.join(output_folder, "dark.fits"), bad_mask
    )

    flat_h = None
    flat_j = None
    listflat_h = os.path.join(output_folder, "listflatH")
    listflat_j = os.path.join(output_folder, "listflatJ")
    if os.path.exists(listflat_h):
        print("Building master flat H")
        flat_h = build_master_frame(
            listflat_h, os.path.join(output_folder, "flatH.fits"), bad_mask, normalize=True
        )
    if os.path.exists(listflat_j):
        print("Building master flat J")
        flat_j = build_master_frame(
            listflat_j, os.path.join(output_folder, "flatJ.fits"), bad_mask, normalize=True
        )

    date_candidates = []
    for item in read_list_file(listdark_path):
        parts = os.path.normpath(item).split(os.sep)
        if len(parts) >= 2:
            date_candidates.append(parts[-2])
    date = date_candidates[0] if date_candidates else mmdd

    sn_folders = sorted(
        path
        for path in (
            os.path.join(output_folder, name) for name in os.listdir(output_folder)
        )
        if os.path.isdir(path)
    )

    for sn_folder in sn_folders:
        print(f"Processing {os.path.basename(sn_folder)}")
        process_band(sn_folder, "H", bad_mask, dark_frame, flat_h, date)
        process_band(sn_folder, "J", bad_mask, dark_frame, flat_j, date)


def rerun_alignment_stage(output_folder):
    output_folder = os.path.abspath(output_folder)
    if not os.path.isdir(output_folder):
        raise RuntimeError(f"Output folder does not exist: {output_folder}")

    listdark_path = os.path.join(output_folder, "listdark")
    if not os.path.exists(listdark_path):
        raise RuntimeError(f"Missing listdark in {output_folder}")

    date_candidates = []
    for item in read_list_file(listdark_path):
        parts = os.path.normpath(item).split(os.sep)
        if len(parts) >= 2:
            date_candidates.append(parts[-2])
    date = date_candidates[0] if date_candidates else os.path.basename(output_folder)

    sn_folders = sorted(
        path
        for path in (
            os.path.join(output_folder, name) for name in os.listdir(output_folder)
        )
        if os.path.isdir(path)
    )
    for sn_folder in sn_folders:
        print(f"Re-running alignment for {os.path.basename(sn_folder)}")
        rerun_alignment_band(sn_folder, "H", date)
        rerun_alignment_band(sn_folder, "J", date)