Euclid Galaxy Clusters Analysis Tutorial

def check_mer_tile_requirement(coord, search_radius=2.0):
    """Check whether a sky coordinate is covered by exactly 4 Euclid MER science images.

    The Euclid MER mosaic is organized into tiles. A position inside a single
    tile returns exactly 4 images, one per photometric band (VIS, Y, J, H).
    Positions near tile boundaries may overlap two tiles and return more than 4
    images; those are rejected to avoid partial-coverage complications.

    Parameters
    ----------
    coord : `~astropy.coordinates.SkyCoord`
        Sky coordinate to check.
    search_radius : float, optional
        Radius in arcminutes within which to search for MER images. Default is 2.0.

    Returns
    -------
    valid : bool
        True if exactly 4 science images are found; False otherwise.
    mer_images : `~astropy.table.Table` or None
        Table of matching MER image records, or None if the query failed.
    """
    try:
        mer_images = Irsa.query_sia(pos=(coord, search_radius * u.arcmin), collection='euclid_DpdMerBksMosaic')
        mer_images = mer_images[(mer_images['facility_name'] == 'Euclid') & (mer_images['dataproduct_subtype'] == 'science')]

        if len(mer_images) == 4:
            print(f"  ✓ Found exactly 4 images for coordinates")
            return True, mer_images
        else:
            print(f"  ✗ Found {len(mer_images)} images (expected 4) - trying next location")
            return False, mer_images
    except (DALQueryError, Timeout, ConnectionError, OSError) as e:
        print(f"  ✗ Error querying MER tiles: {e}")
        return False, None

def find_control_field(cluster_df, cluster_ra, cluster_dec, min_distance_arcmin=15, max_attempts=100):
    """Find a random control field offset from the known cluster catalog.

    Generates candidate control fields by randomly offsetting 25–35 arcmin from
    a randomly chosen catalog cluster in a random direction. Each candidate is
    accepted only if it lies at least ``min_distance_arcmin`` away from every
    cluster in the catalog, ensuring the control field is a genuine blank field
    representative of the general galaxy population.

    RA offsets are corrected for the cosine of the declination so that angular
    separations are accurate. If no valid field is found within ``max_attempts``,
    a fallback position directly offset from the target cluster is returned.

    Parameters
    ----------
    cluster_df : `pandas.DataFrame`
        Cluster catalog with columns ``RAPZWav`` and ``DecPZWav`` (degrees).
    cluster_ra : float
        RA of the target cluster in degrees (used for the fallback only).
    cluster_dec : float
        Dec of the target cluster in degrees (used for the fallback only).
    min_distance_arcmin : float, optional
        Minimum angular separation in arcminutes from all catalog clusters.
        Default is 15.
    max_attempts : int, optional
        Maximum number of random candidates to try before falling back.
        Default is 100.

    Returns
    -------
    control_ra : float
        RA of the selected control field in degrees.
    control_dec : float
        Dec of the selected control field in degrees.
    """
    for attempt in range(max_attempts):
        anchor_cluster = cluster_df.sample(n=1).iloc[0]
        anchor_ra, anchor_dec = anchor_cluster['RAPZWav'], anchor_cluster['DecPZWav']

        offset_arcmin = np.random.uniform(25, 35)
        offset_deg = offset_arcmin / 60.0
        direction = np.random.uniform(0, 2 * np.pi)

        ra_offset = offset_deg * np.cos(direction) / np.cos(np.radians(anchor_dec))
        dec_offset = offset_deg * np.sin(direction)

        control_ra = (anchor_ra + ra_offset) % 360.0
        control_dec = np.clip(anchor_dec + dec_offset, -90.0, 90.0)

        control_coord = SkyCoord(ra=control_ra, dec=control_dec, unit='deg')
        cluster_coords = SkyCoord(ra=cluster_df['RAPZWav'], dec=cluster_df['DecPZWav'], unit='deg')
        distances_arcmin = control_coord.separation(cluster_coords).to(u.arcmin).value

        if np.all(distances_arcmin > min_distance_arcmin):
            return control_ra, control_dec

    # Fallback
    fallback_offset = (min_distance_arcmin + 10) / 60.0
    fallback_ra = (cluster_ra + fallback_offset) % 360.0
    fallback_dec = np.clip(cluster_dec + fallback_offset, -90.0, 90.0)
    return fallback_ra, fallback_dec

def _download_band(band, mer_images, field_coord, field_id, cache_dir, s3):
    """Download one photometric band from S3 and write a cutout FITS to the local cache.

    Parses the S3 path from the ``cloud_access`` JSON column of the MER image
    table, streams the remote FITS file without downloading the full tile, and
    writes a ``Cutout2D`` of size ``im_cutout`` centred on ``field_coord`` to a
    local cache file. If a cache file already exists it is reused without
    re-downloading.

    Parameters
    ----------
    band : str
        Photometric band name, one of ``'VIS'``, ``'Y'``, ``'J'``, ``'H'``.
    mer_images : `~astropy.table.Table`
        MER image table with columns ``energy_bandpassname`` and
        ``cloud_access``.
    field_coord : `~astropy.coordinates.SkyCoord`
        Centre of the cutout region.
    field_id : str
        Unique identifier used to name the cached FITS file on disk.
    cache_dir : str
        Path to the local directory in which cached files are stored.
    s3 : `s3fs.S3FileSystem`
        Authenticated S3 filesystem object used to open the remote file.

    Returns
    -------
    band : str
        The band name, echoed back so callers can build a ``{band: path}`` dict.
    cache_file : str
        Absolute path to the cached FITS file.
    """
    cache_file = os.path.join(cache_dir, f'{band}_{field_id}.fits')
    if not os.path.exists(cache_file):
        print(f"  Downloading {band} band...")

        # Parse the S3 path from the cloud_access JSON column
        row = mer_images[mer_images['energy_bandpassname'] == band][0]
        cloud = json.loads(row['cloud_access'])
        s3_obj = f"{cloud['aws']['bucket_name']}/{cloud['aws']['key']}"

        # Stream the remote FITS via S3 and write only the cutout region
        with s3.open(s3_obj, "rb") as f:
            with fits.open(f, memmap=False, lazy_load_hdus=True) as hdul:
                hdu0 = hdul[0]
                w0 = WCS(hdu0.header)
                # hdu0.section reads only the needed pixel rows from S3
                cut0 = Cutout2D(hdu0.section, position=field_coord, size=im_cutout, wcs=w0)
                fits.writeto(cache_file, cut0.data, header=cut0.wcs.to_header(), overwrite=True)
    else:
        print(f"  Using cached {band} band")
    return band, cache_file

def download_and_cache_field(mer_images, field_name, field_coord, field_id):
    """Stream cutout FITS images from S3 and cache them locally.

    Uses the ``cloud_access`` column in the MER image table to open each band's
    FITS file directly from the AWS S3 mirror without downloading the full tile.
    A ``Cutout2D`` region of size ``im_cutout`` centred on ``field_coord`` is
    written to a local cache file. On subsequent calls, cached files are reused.
    All four bands are downloaded in parallel using threads to reduce wall-clock
    time. After caching, the function re-opens each FITS and returns the 2-D
    pixel arrays together with the VIS-band WCS for downstream analysis.

    Parameters
    ----------
    mer_images : `~astropy.table.Table`
        MER image table as returned by :func:`check_mer_tile_requirement`,
        with columns ``energy_bandpassname``, ``access_url``, and
        ``cloud_access``.
    field_name : str
        Human-readable label (e.g. ``"cluster"`` or ``"control"``) used in
        progress messages.
    field_coord : `~astropy.coordinates.SkyCoord`
        Centre of the cutout.
    field_id : str
        Unique identifier used to name the cached FITS files on disk.

    Returns
    -------
    cutouts : dict
        Dictionary mapping band name (``'VIS'``, ``'Y'``, ``'J'``, ``'H'``)
        to the corresponding 2-D NumPy array.
    cutout_wcs : `~astropy.wcs.WCS`
        WCS object derived from the VIS-band cutout.
    """
    print(f"\nProcessing {field_name} field...")

    # Download all four bands in parallel; S3 reads are I/O-bound so threads help
    with ThreadPoolExecutor(max_workers=4) as executor:
        cached_files = dict(executor.map(
            lambda b: _download_band(b, mer_images, field_coord, field_id, cache_dir, s3),
            ['VIS', 'Y', 'J', 'H']
        ))

    # Read the cached cutout arrays; store the VIS WCS for downstream use
    cutouts = {}
    cutout_wcs = None
    for band in ['VIS', 'Y', 'J', 'H']:
        with fits.open(cached_files[band]) as hdu:
            cutouts[band] = hdu[0].data.copy()
            if band == 'VIS':
                cutout_wcs = WCS(hdu[0].header)

    return cutouts, cutout_wcs

def normalize_with_consistent_stretch(cluster_cutouts, control_cutouts, lower_percentile=1, upper_percentile=99):
    """Normalize cluster and control cutouts using a shared percentile stretch.
    The RGB composite is assembled as H→R, J→G, VIS→B, which maps redder
    (older) stellar populations toward red in the image.

    Parameters
    ----------
    cluster_cutouts : dict
        Band-keyed dictionary of 2-D arrays for the cluster field.
    control_cutouts : dict
        Band-keyed dictionary of 2-D arrays for the control field.
    lower_percentile : float, optional
        Lower percentile used for ``vmin``. Default is 1.
    upper_percentile : float, optional
        Upper percentile used for ``vmax``. Default is 99.

    Returns
    -------
    norm_cluster : dict
        Normalised band arrays for the cluster field.
    norm_control : dict
        Normalised band arrays for the control field.
    cluster_rgb : `~numpy.ndarray`
        (H, W, 3) RGB array for the cluster field.
    control_rgb : `~numpy.ndarray`
        (H, W, 3) RGB array for the control field.
    """
    norm_cluster = {}
    norm_control = {}

    for band in ['VIS', 'Y', 'J', 'H']:
        # Combine both fields to calculate global percentiles
        combined_data = np.concatenate([cluster_cutouts[band].flatten(), control_cutouts[band].flatten()])

        # Calculate global percentiles
        vmin = np.percentile(combined_data, lower_percentile)
        vmax = np.percentile(combined_data, upper_percentile)

        # Apply same normalization to both fields
        norm_cluster[band] = np.clip((cluster_cutouts[band] - vmin) / (vmax - vmin), 0, 1)
        norm_control[band] = np.clip((control_cutouts[band] - vmin) / (vmax - vmin), 0, 1)

        print(f"{band} band: vmin={vmin:.2e}, vmax={vmax:.2e}")

    # Create RGB composites
    cluster_rgb = np.dstack([norm_cluster['H'], norm_cluster['J'], norm_cluster['VIS']])
    control_rgb = np.dstack([norm_control['H'], norm_control['J'], norm_control['VIS']])

    return norm_cluster, norm_control, cluster_rgb, control_rgb

def downsample(arr, factor=4):
    """Block-average a 2-D or 3-D image array by an integer factor for faster display.

    Reshapes the array into non-overlapping blocks of size ``factor × factor``
    and takes the mean of each block. This preserves the overall brightness and
    contrast of the image while reducing the number of pixels that matplotlib
    must render, which significantly speeds up ``imshow`` for large arrays.

    Parameters
    ----------
    arr : `~numpy.ndarray`
        Input image array. Either 2-D (H, W) for single-band images or
        3-D (H, W, 3) for RGB composites.
    factor : int, optional
        Downsampling factor applied to both spatial axes. A factor of 4 reduces
        a 7200 × 7200 VIS cutout to 1800 × 1800 pixels. Default is 4.

    Returns
    -------
    out : `~numpy.ndarray`
        Block-averaged array with shape (H // factor, W // factor) for 2-D
        input or (H // factor, W // factor, 3) for 3-D input.

    Notes
    -----
    Rows and columns that do not fit evenly into ``factor``-sized blocks are
    trimmed before averaging to avoid partial-block artefacts.
    """
    h, w = arr.shape[:2]

    # Trim to the largest dimensions divisible by factor
    h_t = (h // factor) * factor
    w_t = (w // factor) * factor

    if arr.ndim == 2:
        # Reshape into (n_blocks_y, factor, n_blocks_x, factor) and average
        return arr[:h_t, :w_t].reshape(h_t // factor, factor,
                                       w_t // factor, factor).mean(axis=(1, 3))
    else:
        # Same reshape for each colour channel independently
        return arr[:h_t, :w_t].reshape(h_t // factor, factor,
                                       w_t // factor, factor, 3).mean(axis=(1, 3))

def query_galaxies_for_field(ra, dec, field_name, redshift_center, redshift_width=0.1):
    """Query galaxies within a redshift slice around a sky position.

    Submits a TAP/ADQL query to IRSA joining the Euclid Q1 MER photometric
    catalogue and the photo-z catalogue. Only objects satisfying all of the
    following criteria are returned:

    - Positive flux in all four bands (VIS, Y, J, H).
    - Photo-z classification = 2 (galaxy).
    - Fractional 90% credible interval width < 0.20, i.e.
      ``(phz_90_int2 - phz_90_int1) / (1 + phz_median) < 0.20``, selecting
      sources with well-constrained photometric redshifts.
    - ``phz_median`` within ``redshift_center ± redshift_width``.

    Parameters
    ----------
    ra : float
        Field centre RA in degrees.
    dec : float
        Field centre Dec in degrees.
    field_name : str
        Label used in progress messages.
    redshift_center : float
        Central photometric redshift of the cluster.
    redshift_width : float, optional
        Half-width of the redshift slice. Default is 0.1.

    Returns
    -------
    result : `~astropy.table.Table`
        Catalogue of galaxies matching all quality and spatial cuts.
    """
    print(f"\nQuerying galaxies for {field_name} field...")

    adql = (f"SELECT DISTINCT mer.object_id, mer.ra, mer.dec, "
            f"phz.flux_vis_unif, phz.flux_y_unif, phz.flux_j_unif, phz.flux_h_unif, "
            f"phz.phz_classification, phz.phz_median, phz.phz_90_int1, phz.phz_90_int2 "
            f"FROM {table_mer} AS mer "
            f"JOIN {table_phz} as phz "
            f"ON mer.object_id = phz.object_id "
            f"WHERE 1 = CONTAINS(POINT('ICRS', mer.ra, mer.dec), "
            f"BOX('ICRS', {ra}, {dec}, {cutout_deg}, {cutout_deg})) "
            f"AND phz.flux_vis_unif > 0 "
            f"AND phz.flux_y_unif > 0 "
            f"AND phz.flux_j_unif > 0 "
            f"AND phz.flux_h_unif > 0 "
            f"AND phz.phz_classification = 2 "
            f"AND ((phz.phz_90_int2 - phz.phz_90_int1) / (1 + phz.phz_median)) < 0.20 "
            f"AND phz.phz_median BETWEEN {redshift_center-redshift_width} AND {redshift_center+redshift_width}")

    result = Irsa.query_tap(adql).to_table()
    print(f"Found {len(result)} galaxies in {field_name} field (z = {redshift_center:.2f} ± {redshift_width:.2f})")

    return result

def apply_dbscan_clustering(galaxy_df, wcs, rgb_image, field_name, eps=500, min_samples=18):
    """Apply DBSCAN to detect galaxy overdensities in a redshift-selected sample. DBSCAN operates on the 2-D projected
    spatial distribution of galaxies at approximately the same redshift.

    **Parameter choices**

    ``eps=500`` and ``min_samples=18`` were tuned empirically for a 12-arcmin
    cutout at z~0.4 with a photo-z slice width of Δz=0.12. These values are
    data-dependent — feel free to adjust them to suit your science case:

    - Increase ``eps`` for higher-redshift or lower-richness clusters where
      members are more widely separated in projection.
    - Increase ``min_samples`` to suppress spurious small groups; decrease it
      to detect lower-richness structures.
    - For a physically motivated approach, estimate the virial radius at the
      cluster redshift, convert to pixels, and use that as ``eps``.

    Parameters
    ----------
    galaxy_df : `pandas.DataFrame`
        Galaxies pre-filtered to the cluster redshift slice, with columns
        ``ra`` and ``dec`` in degrees.
    wcs : `~astropy.wcs.WCS`
        WCS of the image cutout, used to project RA/Dec to pixel coordinates.
    rgb_image : `~numpy.ndarray`
        The RGB image array; its shape defines the valid pixel bounds.
    field_name : str
        Label used in printed diagnostics.
    eps : float, optional
        Maximum distance in pixels between neighbours in DBSCAN. Default is
        500 px (~50 arcsec at VIS resolution, ~300 kpc at z~0.4).
    min_samples : int, optional
        Minimum number of galaxies required to form a core point. Default is 18.

    Returns
    -------
    labels : `~numpy.ndarray`
        DBSCAN cluster labels for the valid galaxy subset; -1 indicates noise.
    valid_galaxy_coords : `~numpy.ndarray`
        (N, 2) pixel coordinate array for the galaxies that were clustered.
    n_clusters : int
        Number of clusters found (noise points excluded).
    n_noise : int
        Number of noise (unassigned) galaxies.
    """
    print(f"\nApplying DBSCAN clustering to {field_name} field...")

    # Convert galaxy coordinates to pixel coordinates
    galaxy_pixels = wcs.world_to_pixel_values(galaxy_df['ra'], galaxy_df['dec'])

    # Filter galaxies to only show those within the image bounds (needed due to query/cutout mismatch)
    image_height, image_width = rgb_image.shape[:2]
    valid_mask = ((galaxy_pixels[0] >= 0) & (galaxy_pixels[0] < image_width) &
                  (galaxy_pixels[1] >= 0) & (galaxy_pixels[1] < image_height))

    valid_galaxy_coords = np.column_stack([galaxy_pixels[0][valid_mask], galaxy_pixels[1][valid_mask]])

    print(f"  Total galaxies: {len(galaxy_df)}")
    print(f"  Galaxies within image bounds: {valid_mask.sum()}")

    # Apply DBSCAN clustering only to valid galaxies
    clustering = DBSCAN(eps=eps, min_samples=min_samples).fit(valid_galaxy_coords)
    labels = clustering.labels_

    # Count clusters (excluding noise points labeled as -1)
    n_clusters = len(set(labels)) - (1 if -1 in labels else 0)
    n_noise = list(labels).count(-1)

    print(f"{field_name} field: {n_clusters} clusters, {n_noise} noise points")

    return labels, valid_galaxy_coords, n_clusters, n_noise

def identify_cluster_members(galaxy_df, labels, galaxy_coords, field_name):
    """Separate a galaxy sample into cluster members and field galaxies.

    Re-applies the same image-bounds filter used in
    :func:`apply_dbscan_clustering` to align ``galaxy_df`` with the DBSCAN
    ``labels`` array, which only covers galaxies that fell within the image
    footprint. Galaxies with ``label != -1`` are cluster members; those with
    ``label == -1`` are noise (field galaxies).

    Parameters
    ----------
    galaxy_df : `pandas.DataFrame`
        Full galaxy catalogue for the field (before image-bounds filtering),
        with columns ``ra`` and ``dec`` in degrees.
    labels : `~numpy.ndarray`
        DBSCAN labels as returned by :func:`apply_dbscan_clustering`.
    galaxy_coords : `~numpy.ndarray`
        Pixel coordinates of the clustered galaxies (unused directly but kept
        for API consistency).
    field_name : str
        Either ``"Cluster"`` or ``"Control"``; selects the WCS and image used
        for the bounds filter.

    Returns
    -------
    cluster_members : `pandas.DataFrame`
        Galaxies assigned to a cluster group, with an added ``cluster_id``
        column containing the DBSCAN label.
    field_galaxies : `pandas.DataFrame`
        Galaxies classified as noise (``cluster_id = -1``).
    """
    print(f"\nAnalyzing {field_name} field galaxy populations...")

    # The labels array corresponds to the galaxies that were actually used in clustering
    # (those within image bounds). We need to create a mapping.

    # Get the WCS for coordinate conversion
    if field_name == "Cluster":
        wcs = cluster_cutout_wcs
        rgb_image = cluster_rgb
    else:
        wcs = control_cutout_wcs
        rgb_image = control_rgb

    # Convert galaxy coordinates to pixel coordinates
    galaxy_pixels = wcs.world_to_pixel_values(galaxy_df['ra'], galaxy_df['dec'])

    # Filter galaxies to only show those within the image bounds (same as in clustering)
    image_height, image_width = rgb_image.shape[:2]
    valid_mask = ((galaxy_pixels[0] >= 0) & (galaxy_pixels[0] < image_width) &
                  (galaxy_pixels[1] >= 0) & (galaxy_pixels[1] < image_height))

    # Get only the valid galaxies (those used in clustering)
    valid_galaxies = galaxy_df[valid_mask].copy()

    print(f"  Total galaxies in catalog: {len(galaxy_df)}")
    print(f"  Galaxies within image bounds: {len(valid_galaxies)}")
    print(f"  Labels array length: {len(labels)}")

    # Now the labels array should match the valid_galaxies DataFrame
    if len(valid_galaxies) != len(labels):
        print(f"  WARNING: Mismatch between valid galaxies ({len(valid_galaxies)}) and labels ({len(labels)})")
        # Use the minimum length to avoid errors
        min_len = min(len(valid_galaxies), len(labels))
        valid_galaxies = valid_galaxies.iloc[:min_len]
        labels = labels[:min_len]

    # Create a mask for cluster members (labels != -1)
    cluster_member_mask = labels != -1
    field_galaxy_mask = labels == -1

    # Get cluster members and field galaxies
    cluster_members = valid_galaxies[cluster_member_mask].copy()
    field_galaxies = valid_galaxies[field_galaxy_mask].copy()

    # Add cluster assignment information
    cluster_members['cluster_id'] = labels[cluster_member_mask]
    field_galaxies['cluster_id'] = -1  # Field galaxies

    print(f"  Cluster members: {len(cluster_members)} galaxies")
    print(f"  Field galaxies: {len(field_galaxies)} galaxies")

    # Count galaxies per cluster
    if len(cluster_members) > 0:
        cluster_counts = cluster_members['cluster_id'].value_counts().sort_index()

    return cluster_members, field_galaxies

def calculate_color_magnitude(df):
    """Convert uniform-aperture fluxes to AB magnitudes and compute Y-H color.

    Magnitudes are computed as ``m = -2.5 * log10(flux) + 23.9``, where the
    zero-point 23.9 corresponds to fluxes in microjanskies — the unit of the
    ``flux_*_unif`` columns in the Euclid Q1 photo-z catalogue. The resulting
    magnitudes are not K-corrected or corrected for Milky Way extinction.

    Parameters
    ----------
    df : `pandas.DataFrame`
        DataFrame containing columns ``flux_y_unif`` and ``flux_h_unif`` in
        microjanskies. Both must be positive (no non-detections).

    Returns
    -------
    df : `pandas.DataFrame`
        Copy of the input with three additional columns: ``H_mag``, ``Y_mag``
        (AB magnitudes), and ``Y_H_color``.
    """
    df = df.copy()

    # Convert fluxes to magnitudes (using -2.5 * log10(flux))
    # Note: These are instrumental magnitudes, not absolute magnitudes
    df['H_mag'] = -2.5 * np.log10(df['flux_h_unif'])+23.9
    df['Y_mag'] = -2.5 * np.log10(df['flux_y_unif'])+23.9

    # Calculate Y-H color
    df['Y_H_color'] = df['Y_mag'] - df['H_mag']

    return df

def remove_outliers_bounds(df, h_min=17, h_max=25, yh_min=-0.5, yh_max=1.5):
    """Filter a colour-magnitude table to a physically motivated range.

    Retains only rows whose H-band magnitude and Y-H colour fall within the
    specified boundaries. The defaults bracket the expected locus of cluster
    galaxies at z~0.4 in Euclid Q1 data: objects brighter than H=17 are likely
    stars or saturated; objects fainter than H=25 have unreliable photo-z and
    colours; Y-H outside [-0.5, 1.5] typically indicates bad photometry or
    non-galaxy contaminants.

    Parameters
    ----------
    df : `pandas.DataFrame`
        DataFrame with columns ``H_mag`` and ``Y_H_color``.
    h_min : float, optional
        Minimum H magnitude (brightest limit). Default is 17.
    h_max : float, optional
        Maximum H magnitude (faintest limit). Default is 25.
    yh_min : float, optional
        Minimum Y-H colour. Default is -0.5.
    yh_max : float, optional
        Maximum Y-H colour. Default is 1.5.

    Returns
    -------
    df_clean : `pandas.DataFrame`
        Filtered copy of the input DataFrame.
    """
    df_clean = df.copy()
    df_clean = df_clean[
        (df_clean['H_mag'] >= h_min) & (df_clean['H_mag'] <= h_max) &
        (df_clean['Y_H_color'] >= yh_min) & (df_clean['Y_H_color'] <= yh_max)
    ]
    return df_clean

def get_n_spectra(obj_ids, n=10):
    """
    Fetch up to n spectra for the given object_ids, stopping early once n are found.
    Uses one TAP query, groups by FITS file, caches per-object results in ./data/irsa_spectra/.
    Skips TAP rows with missing (masked) HDU indices.
    """
    obj_ids = [int(x) for x in obj_ids]
    if len(obj_ids) == 0 or n <= 0:
        return {}

    # Cache-first
    spectra = {}
    remaining = []
    for oid in obj_ids:
        cache_npz = os.path.join(spectra_cache_dir, f"{oid}.npz")
        if os.path.exists(cache_npz):
            d = np.load(cache_npz)
            spectra[oid] = {
                "wave": d["wave"] * u.angstrom,
                "flux": d["flux"] * (u.erg / u.s / u.cm**2 / u.angstrom),
                "error": d["error"] * (u.erg / u.s / u.cm**2 / u.angstrom),
                "object_id": oid,
            }
            if len(spectra) >= n:
                return dict(list(spectra.items())[:n])
        else:
            remaining.append(oid)

    if len(remaining) == 0:
        return dict(list(spectra.items())[:n])

    # TAP query once
    id_list = ",".join(map(str, remaining))
    adql_query = f"""
    SELECT objectid, path, hdu
    FROM {table_1dspectra}
    WHERE objectid IN ({id_list})
    """

    try:
        assoc = Irsa.query_tap(adql_query).to_table()
    except (DALQueryError, requests.exceptions.RequestException) as e:
        print("TAP query failed:", e)
        return dict(list(spectra.items())[:n])

    if len(assoc) == 0:
        print("No spectrum associations returned.")
        return dict(list(spectra.items())[:n])

    # Group by FITS file, skipping masked/invalid hdu
    groups = {}
    skipped_hdu = 0

    for row in assoc:
        obj_id = int(row["objectid"])
        uri = str(row["path"])

        # HDU can be masked for some rows -> skip
        hdu_val = row["hdu"]
        try:
            # This will fail for masked values
            hdu_index = int(hdu_val)
        except Exception:
            skipped_hdu += 1
            continue

        spec_fpath_key = uri.replace("api/spectrumdm/convert/euclid/", "").split("?")[0]
        s3_uri = f"s3://{BUCKET_NAME}/{spec_fpath_key}"
        groups.setdefault(s3_uri, []).append((obj_id, hdu_index))

    if skipped_hdu > 0:
        print(f"Skipped {skipped_hdu} TAP rows with missing/invalid HDU indices.")

    if len(groups) == 0:
        print("No usable (path,hdu) associations after filtering.")
        return dict(list(spectra.items())[:n])

    # Open files with progress bar
    for s3_uri, items in tqdm(list(groups.items()), desc="Opening FITS files", unit="file"):
        if len(spectra) >= n:
            break

        try:
            with fits.open(s3_uri, fsspec_kwargs={"anon": True}, lazy_load_hdus=True) as hdul:
                for (obj_id, hdu_index) in items:
                    if len(spectra) >= n:
                        break

                    try:
                        spec = QTable.read(hdul[hdu_index], format="fits")
                        header = hdul[hdu_index].header

                        fscale = header.get("FSCALE", 1.0)
                        wave = np.asarray(spec["WAVELENGTH"]) * u.angstrom
                        signal = np.asarray(spec["SIGNAL"])
                        var = np.asarray(spec["VAR"])
                        mask = np.asarray(spec["MASK"])

                        if not (len(wave) == len(signal) == len(var) == len(mask)):
                            continue

                        valid = (mask % 2 == 0) & (mask < 64)
                        wave = wave[valid]
                        flux = signal[valid] * fscale * u.erg / u.s / u.cm**2 / u.angstrom
                        error = np.sqrt(var[valid]) * fscale * flux.unit

                        spectra[obj_id] = {
                            "wave": wave,
                            "flux": flux,
                            "error": error,
                            "object_id": obj_id,
                        }

                        np.savez(
                            os.path.join(spectra_cache_dir, f"{obj_id}.npz"),
                            wave=wave.value,
                            flux=flux.value,
                            error=error.value,
                        )

                    except Exception:
                        continue

        except Exception as e:
            print(f"Failed to open {s3_uri}: {e}")
            continue

    return dict(list(spectra.items())[:n])

def preprocess_spectrum(spec):
    """Continuum-remove + robust-normalize a spectrum for visualization."""
    w = np.asarray(spec['wave'].value, float)
    f = np.asarray(spec['flux'].value, float)
    e = np.asarray(spec['error'].value, float) if 'error' in spec else None

    good = np.isfinite(w) & np.isfinite(f)
    if e is not None:
        good &= np.isfinite(e) & (e > 0)

    w = w[good]
    f = f[good]
    if e is not None:
        e = e[good]

    if w.size < 30:
        return None

    # sort by wavelength (just in case)
    s = np.argsort(w)
    w, f = w[s], f[s]
    if e is not None:
        e = e[s]

    # light smoothing to suppress pixel-to-pixel noise
    f_s = gaussian_filter1d(f, sigma=sigma_smooth) if (sigma_smooth and sigma_smooth > 0) else f

    # running median continuum
    k = int(cont_window)
    if k % 2 == 0:
        k += 1
    if k >= f_s.size:
        k = max(5, (f_s.size // 2) * 2 - 1)

    cont = median_filter(f_s, size=k, mode='nearest')
    resid = f_s - cont

    # robust normalization so objects are comparable in amplitude
    scale = np.nanpercentile(np.abs(resid), norm_percentile)
    if not np.isfinite(scale) or scale <= 0:
        scale = np.nanstd(resid)
    if not np.isfinite(scale) or scale <= 0:
        return None

    y = resid / scale
    return w, y

def build_stack(spectra_dict, w_grid):
    """Interpolate processed spectra onto a common grid and return matrix [nobj, ngrid]."""
    Ys = []
    for obj_id, spec in spectra_dict.items():
        out = preprocess_spectrum(spec)
        if out is None:
            continue
        w, y = out
        # interpolate to common grid; outside range -> NaN so it doesn't bias stats
        y_i = np.interp(w_grid, w, y, left=np.nan, right=np.nan)
        Ys.append(y_i)

    if len(Ys) == 0:
        return None

    Y = np.vstack(Ys)
    med = np.nanmedian(Y, axis=0)
    p16 = np.nanpercentile(Y, 16, axis=0)
    p84 = np.nanpercentile(Y, 84, axis=0)
    return Y, med, p16, p84

def lines_in_window(wmin, wmax, z):
    """Return dict of lines whose observed wavelength falls in [wmin, wmax]."""
    out = {}
    for name, rest in emission_lines.items():
        obs = rest * (1 + z)
        if wmin <= obs <= wmax:
            out[name] = obs
    return out

def label_emission_line(ax, x, label, y=0.92, rotation=90, color='k'):
    """Place a vertical emission-line label at wavelength x."""
    ax.text(
        x, y, label,
        transform=ax.get_xaxis_transform(),
        rotation=rotation,
        color=color,
        fontsize=10,
        ha='right',
        va='top',
        bbox=dict(facecolor='white', edgecolor='none', alpha=0.6, pad=1)
    )

def plot_panel(ax, spectra_dict, stack, color, title):
    # individual processed spectra (faint)
    for obj_id, spec in spectra_dict.items():
        out = preprocess_spectrum(spec)
        if out is None:
            continue
        w, y = out
        ax.plot(w, y, color=color, alpha=0.35, lw=1)

    # median + spread (this is the key improvement)
    if stack is not None:
        Y, med, p16, p84 = stack
        ax.plot(w_grid, med, color=color, lw=2.5, alpha=0.9, label=f"median (n={Y.shape[0]})")
        ax.fill_between(w_grid, p16, p84, color=color, alpha=0.15, label="16–84%")

    ax.axhline(0, color="k", lw=1, alpha=0.25)
    ax.set_ylabel("Continuum-subtracted, robust-normalized flux")
    ax.set_title(title)
    ax.grid(True, alpha=0.35)
    ax.legend(loc="upper right")

def search_ned_field(ra, dec, z_min, z_max, label, radius_arcmin=3, max_retries=3):
    """Search NED within radius_arcmin of (ra, dec), filter to z_min–z_max, and print a summary.

    Returns
    -------
    objects : astropy Table or None
        Rows in the redshift range, or None on failure.
    ned_ra_col, ned_dec_col : str or None
        Names of the RA and Dec columns in the returned table.
    ned_coord_unit : tuple or None
        Unit tuple suitable for SkyCoord construction.
    """
    ra_candidates  = ['RA(deg)', 'RA_deg', 'RA', 'RAJ2000']
    dec_candidates = ['DEC(deg)', 'DEC_deg', 'DEC', 'DEJ2000']

    def pick_col(table, candidates):
        for c in candidates:
            if c in table.colnames:
                return c
        return None

    center = SkyCoord(ra=ra, dec=dec, unit='deg')

    for attempt in range(max_retries):
        try:
            results = Ned.query_region(center, radius=radius_arcmin * u.arcmin)

            has_z = np.isfinite(np.asarray(results['Redshift'], float))
            objects_with_z = results[has_z]
            in_z_range = (objects_with_z['Redshift'] >= z_min) & (objects_with_z['Redshift'] <= z_max)
            objects = objects_with_z[in_z_range]

            print(f"{label}: {len(results)} total NED objects, {len(objects_with_z)} with redshift, "
                  f"{len(objects)} in z={z_min:.2f}–{z_max:.2f}")

            ned_ra_col  = pick_col(results, ra_candidates)
            ned_dec_col = pick_col(results, dec_candidates)
            if ned_ra_col is None or ned_dec_col is None:
                raise KeyError(f"Could not find RA/Dec columns. Available: {results.colnames}")

            ra_unit_str = str(getattr(results[ned_ra_col], 'unit', '')).lower()
            use_deg = 'deg' in ra_unit_str or 'deg' in ned_ra_col.lower()
            ned_coord_unit = ('deg', 'deg') if use_deg else (u.hourangle, u.deg)

            if len(objects) > 0:
                for obj in objects:
                    obj_coord = SkyCoord(ra=obj[ned_ra_col], dec=obj[ned_dec_col], unit=ned_coord_unit)
                    sep = center.separation(obj_coord).to(u.arcmin).value
                    print(f"  {obj['Object Name']} - {obj['Type']} - z={obj['Redshift']:.3f} - {sep:.1f}'")
                types = np.array([str(t) for t in objects['Type']], dtype=str)
                unique, counts = np.unique(types, return_counts=True)
                print("  Types:", dict(zip(unique, counts)))
            else:
                print("  No objects found in redshift range")

            return objects, ned_ra_col, ned_dec_col, ned_coord_unit

        except (Timeout, ConnectionError) as e:
            if attempt < max_retries - 1:
                time.sleep(5)
            else:
                print(f"{label}: NED search failed after {max_retries} attempts")
        except (RemoteServiceError, OSError, KeyError) as e:
            print(f"{label}: NED search error: {e}")
            break

    return None, None, None, None

def overlay_ned_sources(
    cluster_objects, control_objects,
    ned_ra_col, ned_dec_col, ned_coord_unit,
    cluster_rgb, control_rgb,
    cluster_cutout_wcs, control_cutout_wcs,
    cluster_labels, cluster_galaxy_coords,
    control_labels, control_galaxy_coords,
):
    if cluster_objects is None or len(cluster_objects) == 0:
        print("No cluster NED objects to overlay.")
        return

    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(20, 8),
                                   subplot_kw={'projection': cluster_cutout_wcs})

    for ax, rgb, labels, coords, ned_objs, wcs, title in [
        (ax1, cluster_rgb, cluster_labels, cluster_galaxy_coords,
         cluster_objects, cluster_cutout_wcs, 'Cluster Field with NED Sources'),
        (ax2, control_rgb, control_labels, control_galaxy_coords,
         control_objects if control_objects is not None else [], control_cutout_wcs, 'Control Field'),
    ]:
        ax.imshow(rgb, origin='lower', alpha=0.3)

        unique_labels = set(labels)
        colors = plt.cm.Spectral(np.linspace(0, 1, len(unique_labels)))
        for k, col in zip(unique_labels, colors):
            if k == -1:
                continue  # noise points not shown
            mask = (labels == k)
            xy = coords[mask]
            if len(xy) > 0:
                ax.scatter(xy[:, 0], xy[:, 1], c='red', marker='o', s=20, alpha=0.8,
                           edgecolors='white', linewidth=0.5)

        for obj in ned_objs:
            ned_coord = SkyCoord(ra=obj[ned_ra_col], dec=obj[ned_dec_col], unit=ned_coord_unit)
            ned_pixel = wcs.world_to_pixel(ned_coord)
            if (0 <= ned_pixel[0] < rgb.shape[1] and 0 <= ned_pixel[1] < rgb.shape[0]):
                ax.scatter(ned_pixel[0], ned_pixel[1], facecolors='none', marker='o', s=100,
                           alpha=0.9, edgecolors='blue', linewidth=3)

        ax.set_title(title, fontsize=14, fontweight='bold')
        ax.set_xlabel('RA')
        ax.set_ylabel('Dec')

    handles = [
        plt.Line2D([0], [0], marker='o', color='none', linestyle='None',
                   markersize=8, markeredgecolor='blue', markeredgewidth=2),
        plt.Line2D([0], [0], marker='o', color='red', linestyle='None',
                   markersize=8, alpha=0.8),
    ]
    fig.legend(handles, ['NED Galaxy', 'Euclid Cluster Members'],
               loc='upper center', bbox_to_anchor=(0.5, 0.95), ncol=2, fontsize=10)

    plt.tight_layout()
    plt.show()