import numpy as np
from astropy.io import fits
from astropy.stats import sigma_clipped_stats
from photutils.segmentation import detect_threshold, detect_sources
import matplotlib.pyplot as plt

# ---------------------------
# User settings
# ---------------------------
input_fits = "example_image.fits"  # Path to your FITS image
mask_output_fits = "source_mask.fits"
nsigma = 3.0  # Detection threshold in sigma
min_area = 5  # Minimum number of pixels for a source

# ---------------------------
# Load the image
# ---------------------------
data = fits.getdata(input_fits)

# Estimate background stats
mean, median, std = sigma_clipped_stats(data, sigma=3.0)
print(f"Background: mean={mean:.3f}, median={median:.3f}, std={std:.3f}")

# ---------------------------
# Detect sources
# ---------------------------
threshold = detect_threshold(data, nsigma=nsigma)
segm = detect_sources(data, threshold, npixels=min_area)

if segm is None:
    print("No sources detected.")
else:
    # ---------------------------
    # Create a boolean mask
    # ---------------------------
    mask = segm.data > 0  # True where a source is detected

    # Save mask as FITS
    fits.writeto(mask_output_fits, mask.astype(np.uint8), overwrite=True)
    print(f"Mask saved to {mask_output_fits}")

    # ---------------------------
    # Show results
    # ---------------------------
    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(10, 5))
    ax1.imshow(data, origin='lower', cmap='gray', vmin=median-2*std, vmax=median+5*std)
    ax1.set_title("Original Image")
    ax2.imshow(mask, origin='lower', cmap='gray')
    ax2.set_title("Source Mask")
    plt.show()