TDE Light Curve

def show_gallery(files, max_images=9):
    """
    Display a gallery of FITS images.

    Parameters
    ----------
    files : list of str
        List of S3 URIs to FITS files.
    max_images : int, optional
        Maximum number of images to display (default: 9).
    """
    # Limit the number of images to display
    n_images = min(len(files), max_images)

    # Choose number of columns: up to 3, or equal to n_images if fewer than 4
    ncols = n_images if n_images < 4 else 3

    # Compute number of rows based on total images and columns
    nrows = (n_images + ncols - 1) // ncols

    # Create the subplot grid
    fig, axes = plt.subplots(nrows, ncols, figsize=(4 * ncols, 4 * nrows))
    axes = np.atleast_1d(axes).ravel()  # flatten in case of 1D output

    # Loop through each file and display the image
    for i, f in enumerate(files[:n_images]):
        # Open FITS file directly from S3 (anonymous read)
        with fits.open(f, fsspec_kwargs={"anon": True}, memmap=False) as hdul:
            data = hdul[1].data  # Extract image data
            # Compute robust display scaling (5th–99th percentile)
            vmin, vmax = np.nanpercentile(data, [5, 99])
            # Show image in grayscale with good contrast
            axes[i].imshow(data, origin="lower", cmap="gray", vmin=vmin, vmax=vmax)
            del data
            # Title: just the base filename
            axes[i].set_title(f.split("/")[-1], fontsize=8)
            # Remove tick marks and labels
            axes[i].axis("off")

    # Hide any unused subplot axes
    for j in range(i + 1, len(axes)):
        axes[j].axis("off")

    # Adjust layout for neat display
    plt.tight_layout()
    plt.show()

def get_s3_fpath(cloud_access):
    """Extract the S3 URI from the cloud_access JSON string in an SIA result."""
    cloud_info = json.loads(cloud_access) # converts str to dict
    bucket_name = cloud_info['aws']['bucket_name']
    key = cloud_info['aws']['key']

    return f's3://{bucket_name}/{key}'

def Roman_TDS_image_search(host_galaxy, radius, bandname):
    """
    Get OpenUniverse2024 Roman TDS images for the TDE host galaxy.

    Parameters
    ----------
    host_galaxy : pandas.DataFrame
        Must include 'galaxy_id', 'ra', and 'dec' columns.
    radius : astropy.units.Quantity
        Search radius.
    bandname : string
        Roman bandname for which to search images (e.g. 'J129', 'H158').

    Returns
    -------
    astropy.table.Table
        SIA result table filtered to TDE images in the specified band,
        with an added 's3_uri' column for the image file paths.
    """
    row = host_galaxy.iloc[0]
    galaxy_id = int(row["galaxy_id"])
    ra_center, dec_center = row["ra"], row["dec"]
    print(f"Accessing images in band={bandname} for galaxy_id={galaxy_id} at RA={ra_center:.3f}, Dec={dec_center:.3f} ...", end="")

    coords = SkyCoord(ra_center, dec_center, unit='deg')
    sia_results = Irsa.query_sia(pos=(coords, radius.to(u.deg)), collection=OU_ROMAN_SIA_COLLECTION)
    filtered_results = sia_results[['TDS_simple_model' in r['obs_id'] and bandname in r['energy_bandpassname']
                                    for r in sia_results]]
    filtered_results['s3_uri'] = [get_s3_fpath(r['cloud_access']) for r in filtered_results]

    print(f"done. Found {len(filtered_results)} images.")
    return filtered_results

def select_images_by_mjd_quantiles(images, n_select=9):
    """
    Select up to `n_select` images using rank-quantiles on time axis.

    Parameters
    ----------
    images : astropy Table
        Table of images (SIA results).
    n_select : int, optional
        Maximum number of images to select. Default is 9.

    Returns
    -------
    astropy Table
        Subset of input `images` corresponding to the selected quantiles.
    """
    sorted_images = images.copy()  # avoid modifying the original table
    sorted_images.sort('t_min') # sort by time axis, t_min = t_max for TDS images, so we can use either
    num_quantiles = min(n_select, len(sorted_images))
    quantile_indices = np.rint(  # round off to nearest integer
        np.linspace(0, len(sorted_images) - 1, num_quantiles)
    ).astype(int)
    return sorted_images[quantile_indices]

def run_aperture_photometry(host_galaxy, bandname, image_column="image_filenames", aperture_radius=1.0):
    """
    Perform circular aperture photometry on a list of FITS images.

    Parameters
    ----------
    host_galaxy : pandas.DataFrame
        Must contain columns 'ra', 'dec', and a nested column with FITS image paths
        (each a list of FITS image paths).
    bandname : string
        Bandname for which to do photometry.
    image_column : string, optional
        Name of the dataframe column containing lists of FITS image paths.
        Default is "image_filenames".
    aperture_radius : float, optional
        Aperture radius in arcsec. Default is 1.0 arcseconds which is ~9 pixels for Roman .

    Returns
    -------
    None
        Results are added directly to `host_galaxy` as new columns:
        f'mjd_obs_{{bandname}}', f'flux_{{bandname}}', f'flux_err_{{bandname}}',
        and 'aperture_radius_pix'.
    """

    row = host_galaxy.iloc[0]
    filenames = row[image_column]
    print(f"Performing photometry in {bandname} for {len(filenames)} sampled images "
          f"for host galaxy at RA={row['ra']:.3f}, Dec={row['dec']:.3f} ...", end="")

    mjd_list, flux_list, flux_err_list, aperture_radius_pix_list = [], [], [], []

    for fname in filenames:

        #opening a gzipped fits file, don't recommend changing the next line.
        with fits.open(fname, fsspec_kwargs={"anon": True}, memmap=False) as hdul:
            data = hdul[1].data
            header = hdul[1].header

            # Build a WCS object so photutils can convert between sky and pixel coordinates.
            wcs = WCS(header)

            # Simple circular aperture centered on the host galaxy position
            sky_position = SkyCoord(row["ra"], row["dec"], unit="deg", frame="icrs")
            aperture = SkyCircularAperture(sky_position, r=aperture_radius * u.arcsec)
            pixel_aperture = aperture.to_pixel(wcs)

            # Perform aperture photometry
            phot_table = aperture_photometry(data, aperture, wcs=wcs)

            # Check output (optional)
            #print(phot_table)

            # Background estimate from a local annulus around the source
            annulus = SkyCircularAnnulus(sky_position, r_in=3 * aperture_radius * u.arcsec,
                                                       r_out=6 * aperture_radius * u.arcsec)
            annulus_stats = ApertureStats(data, annulus, wcs=wcs)
            background = annulus_stats.median

            # Subtract background from aperture sum
            aperture_area = pixel_aperture.area
            flux = phot_table['aperture_sum'][0] - background * aperture_area

            # Approximate uncertainty from background rms in annulus
            flux_err = annulus_stats.std * np.sqrt(aperture_area)

            # Observation mid-time from MJD-OBS
            mjd_obs = header.get('MJD-OBS', None)

            #store related info
            mjd_list.append(mjd_obs)
            flux_list.append(flux)
            flux_err_list.append(flux_err)
            aperture_radius_pix_list.append(float(pixel_aperture.r))

    print("done.")

    # Add as nested columns
    host_galaxy[f"mjd_obs_{bandname}"] = [mjd_list]
    host_galaxy[f"flux_{bandname}"] = [flux_list]
    host_galaxy[f"flux_err_{bandname}"] = [flux_err_list]
    host_galaxy["aperture_radius_pix"] = [aperture_radius_pix_list]

def plot_single_band_light_curve(df, bandname, start_mjd):
    """
    Plot the TDE host galaxy light curve for a single band, with error bars.

    Parameters
    ----------
    df : pandas.DataFrame
        Must contain f'mjd_obs_{{bandname}}', f'flux_{{bandname}}',
        and f'flux_err_{{bandname}}' columns.
    bandname : str
        Photometric band name used for column labels.
    start_mjd : float
        MJD of the TDE start, used to set the x-axis origin.

    Returns
    -------
    matplotlib.figure.Figure
        The generated figure.
    """
    row = df.iloc[0]
    galaxy_id = row["galaxy_id"]

    # Extract nested arrays
    #need to go to numpy so we can check for non-finite values
    times = np.array(row[f"mjd_obs_{bandname}"], dtype=float) - start_mjd
    fluxes = np.array(row[f"flux_{bandname}"], dtype=float)
    flux_errs = np.array(row[f"flux_err_{bandname}"], dtype=float)

    # Drop NaNs / non-finite values
    mask = np.isfinite(times) & np.isfinite(fluxes)
    mask &= np.isfinite(flux_errs)
    times, fluxes, flux_errs = times[mask], fluxes[mask], flux_errs[mask]
    if len(times) == 0:
        print(f"⚠️ No valid flux points.")
        return None

    #sort on time
    sort_idx = np.argsort(times)
    times, fluxes, flux_errs = times[sort_idx], fluxes[sort_idx], flux_errs[sort_idx]

    # Normalize to median flux
    median_flux = np.median(fluxes)
    fluxes = fluxes / median_flux
    flux_errs = flux_errs / median_flux

    # Plot
    fig, ax = plt.subplots(figsize=(6, 4))
    ax.errorbar(
        times, fluxes, yerr=flux_errs, fmt="o", capsize=3, label=f"Galaxy {galaxy_id}"
    )
    ax.plot(times, fluxes, "-", alpha=0.6, color=ax.get_lines()[-1].get_color())

    ax.set_xlabel("Days since start of TDE")
    ax.set_ylabel("Normalized Flux")
    ax.set_title(f"Light Curve for Galaxy {galaxy_id} ({bandname})")
    ax.legend()

    # Restrict x-axis to data range with small padding
    margin = 0.05 * (times.max() - times.min()) if len(times) > 1 else 0.1
    ax.set_xlim(times.min() - margin, times.max() + margin)

    plt.show()
    return fig

def plot_multiband_light_curve(df, bands, start_mjd):
    """
    Plot the TDE host galaxy light curve for multiple bands on a single axes.

    Parameters
    ----------
    df : pandas.DataFrame
        Must contain f'mjd_obs_{{band}}', f'flux_{{band}}', and
        f'flux_err_{{band}}' columns for each band in `bands`.
    bands : list of str
        Roman band names to plot (e.g. ['J129', 'H158']).
    start_mjd : float
        MJD of the TDE start, used to set the x-axis origin.

    Returns
    -------
    matplotlib.figure.Figure
        The generated figure.
    """
    colors = itertools.cycle(plt.cm.tab10.colors)

    fig, ax = plt.subplots(figsize=(7, 5))
    row = df.iloc[0]
    xmin, xmax = np.inf, -np.inf

    for bandname in bands:
        # Extract nested arrays
        #need to go to numpy so we can check for non-finite values
        times = np.array(row[f"mjd_obs_{bandname}"], dtype=float) - start_mjd
        fluxes = np.array(row[f"flux_{bandname}"], dtype=float)
        flux_errs = np.array(row[f"flux_err_{bandname}"], dtype=float)

        # Drop invalid
        mask = np.isfinite(times) & np.isfinite(fluxes) & np.isfinite(flux_errs)
        times, fluxes, flux_errs = times[mask], fluxes[mask], flux_errs[mask]

        if len(times) == 0:  #empty photometry
            continue

        #sort on time
        sort_idx = np.argsort(times)
        times, fluxes, flux_errs = times[sort_idx], fluxes[sort_idx], flux_errs[sort_idx]

        # Normalize to median flux
        median_flux = np.median(fluxes)
        fluxes = fluxes / median_flux
        flux_errs = flux_errs / median_flux

        color = next(colors)
        ax.errorbar(times, fluxes, yerr=flux_errs, fmt="o", capsize=3,
                    color=color, label=bandname)
        ax.plot(times, fluxes, "-", alpha=0.6, color=color)

        xmin = min(xmin, times.min())
        xmax = max(xmax, times.max())

    ax.set_xlabel("Days since start of TDE")
    ax.set_ylabel("Normalized Flux")
    ax.set_title("TDE Host Galaxy Light Curve")
    ax.legend(fontsize="small")

    if xmin == np.inf or xmax == -np.inf:
        print("⚠️ No valid photometry data found for any band. Cannot plot light curve.")
        plt.close(fig)
        return fig

    # Restrict x-axis to data range with padding
    margin = 0.05 * (xmax - xmin) if xmax > xmin else 0.1
    ax.set_xlim(xmin - margin, xmax + margin)

    plt.show()
    return fig

def make_cutout(fname, ra, dec, size=100):
    """
    Create a North-up cutout around (RA, Dec) from a Roman TDS FITS image on S3.

    Parameters
    ----------
    fname : str
        Full or partial S3 path to the .fits.gz image.
    ra, dec : float
        Target coordinates in degrees.
    size : int or float, optional
        Cutout size in pixels. Default = 100.

    Returns
    -------
    cutout : astropy.nddata.Cutout2D or None
        Cutout centered on (RA, Dec) and rotated so North points up.
        Returns None if the target is outside the field.
    """

    coord = SkyCoord(ra=ra * u.deg, dec=dec * u.deg)

    with fits.open(fname, fsspec_kwargs={"anon": True}, memmap=False) as hdu:
        img = hdu[1].data
        header = hdu[1].header
        wcs = WCS(header)

        try:
            # Step 1: extract cutout centered on target using the native WCS
            cutout = Cutout2D(img, coord, size, wcs=wcs, mode="partial")
        except Exception as e:
            print(f"Error creating cutout for {fname}: {e}")
            return None
        del img

        # Step 2: determine rotation angle to place North at the top.
        # Invert the CD matrix to find the North direction in pixel space.
        # This handles chip orientation automatically without needing SCA_NUM logic.
        det = header['CD1_1'] * header['CD2_2'] - header['CD1_2'] * header['CD2_1']
        north_col = -header['CD1_2'] / det
        north_row =  header['CD1_1'] / det
        angle = (-np.degrees(np.arctan2(north_col, north_row))) % 360

        # Step 3: rotate only the small cutout (not the full image) to align North up
        cutout.data = rotate(cutout.data, angle=angle, reshape=False, cval=np.nan)

        return cutout

def cutout_gallery(image_filenames, mjd_list, ra, dec, aperture_radius_pix_list, size=100, ncols=4,
                   galaxy_id=None):
    """
    Display a gallery of cutouts centered on (RA, Dec) for a list of Roman TDS images.

    Parameters
    ----------
    image_filenames : list of str
        List of S3 image filenames.
    mjd_list : list-like
        Observation MJD values aligned with `image_filenames`.
    ra, dec : float
        Target coordinates in degrees.
    aperture_radius_pix_list : list of float
        Aperture radius values in pixels, aligned with `image_filenames`.
    size : int or float, optional
        Cutout size in pixels. Default = 100.
    ncols : int, optional
        Number of columns in the gallery grid. Default = 4.
    galaxy_id : int or str, optional
        Galaxy ID for labeling the plot.

    Returns
    -------
    fig : matplotlib.figure.Figure
        The displayed figure object.
    """
    # Initialize lists to store image cutouts, observation times, and aligned pixel radii.
    cutouts, mjds, radii_pix = [], [], []

    # Loop over all image filenames and generate cutouts and only keep MJDs and pixel radii for those with valid cutouts
    for fname, mjd, radius_pix in zip(image_filenames, mjd_list, aperture_radius_pix_list):
        cutout = make_cutout(fname, ra, dec, size=size)
        if cutout is not None:
            cutouts.append(cutout)
            mjds.append(mjd)
            radii_pix.append(radius_pix)

    # Stop if no valid cutouts were created
    if not cutouts:
        raise ValueError("No valid cutouts could be created.")

    # Set up grid dimensions for displaying the gallery
    n_images = len(cutouts)
    nrows = int(np.ceil(n_images / ncols))
    fig, axes = plt.subplots(nrows, ncols, figsize=(3 * ncols, 3 * nrows))
    axes = np.atleast_1d(axes).ravel()

    # Build figure title if information is available
    if galaxy_id is not None:
        fig.suptitle(
            f"Cutouts of TDE host galaxy {galaxy_id}",
            fontsize=14, y=0.98
        )
    elif galaxy_id is not None:
        fig.suptitle(
            f"Cutouts for host galaxy {galaxy_id}",
            fontsize=14, y=0.98
        )

    # Display each cutout image with contrast scaling and MJD label
    for ax, cutout, mjd, radius_pix in zip(axes, cutouts, mjds, radii_pix):
        img = cutout.data
        vmin, vmax = np.nanpercentile(img, [5, 99])
        ax.imshow(img, origin="lower", cmap="gray", vmin=vmin, vmax=vmax)
        # Draw the aperture circle at the image center.
        # The galaxy is centered by Cutout2D and stays centered after the North-up rotation.
        x_center = img.shape[1] / 2.0
        y_center = img.shape[0] / 2.0
        if np.isfinite(radius_pix) and radius_pix > 0:
            aperture_circle = plt.Circle((x_center, y_center), radius_pix,
                                         edgecolor="cyan", facecolor="none", linewidth=1.3)
            ax.add_patch(aperture_circle)

        ax.set_title(f"MJD {mjd:.2f}", fontsize=8)
        ax.axis("off")

    # Hide any extra axes
    for ax in axes[len(cutouts):]:
        ax.axis("off")

    plt.tight_layout()
    plt.show()
    return