########################################################
# Started Logging At: 2026-06-15 14:27:09
########################################################
import warnings
import math
from pathlib import Path

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.patches import Circle, Rectangle

from astropy.table import Table, vstack
from astropy.coordinates import SkyCoord, match_coordinates_sky
import astropy.units as u
from astropy.nddata import Cutout2D
from astropy.visualization import simple_norm
from astropy.io import fits
from astropy.wcs import WCS

warnings.filterwarnings('ignore', category=UserWarning)
warnings.filterwarnings('ignore', category=fits.verify.VerifyWarning)

get_ipython().run_line_magic('matplotlib', 'inline')
get_ipython().run_line_magic('config', "InlineBackend.figure_formats = ['png']")
get_ipython().run_line_magic('config', "InlineBackend.rc = {'figure.dpi': 96}")
# ──────────────────────────────────────────────────────────────
# MIRI Vega zero-point fluxes (Jy)
# F770W, F1500W from STScI calibration docs
# F1130W from Rayleigh-Jeans scaling of F770W Vega spectrum
# F2550W from STScI calibration docs
# ──────────────────────────────────────────────────────────────
MIRI_VEGA_ZP_JY = {
    'f770w':  64.13,
    'f1130w': 29.40,
    'f1500w': 17.60,
    'f2550w':  6.72,
}

WAVELENGTHS_UM = {
    'f187n': 1.874,
    'f210m': 2.096,
    'f335m': 3.362,
    'f470n': 4.707,
    'f480m': 4.817,
    'f770w': 7.639,
    'f1130w': 11.309,
    'f1500w': 15.065,
    'f2550w': 25.363,
}

SICKLE_BASE = Path('/orange/adamginsburg/jwst/sickle')
CATS_DIR    = SICKLE_BASE / 'catalogs'

# MIRI merged catalogs (one per observation pointing, Vega mags)
MIRI_CAT_FILES = {
    'o001': CATS_DIR / 'o001_miri_cmd_matched.fits',
    'o002': CATS_DIR / 'o002_miri_cmd_matched.fits',
    'o003': CATS_DIR / 'o003_miri_cmd_matched.fits',
}

# NIRCam merged catalog (iter3, includes prog 3958 + 2221 + 1182 where available)
NIRCAM_CAT_FILE = CATS_DIR / 'iterative_merged_indivexp_photometry_tables_merged_iter3.fits'

# Best i2d images — NIRCam (satstar-replaced data_i2d)
IMAGE_FILES = {
    'f187n': SICKLE_BASE / 'F187N/pipeline/jw03958-o007_t001_nircam_clear-f187n-nrcb_data_i2d.fits',
    'f210m': SICKLE_BASE / 'F210M/pipeline/jw03958-o007_t001_nircam_clear-f210m-nrcb_data_i2d.fits',
    'f335m': SICKLE_BASE / 'F335M/pipeline/jw03958-o007_t001_nircam_clear-f335m-nrcb_data_i2d.fits',
    'f470n': SICKLE_BASE / 'F470N/pipeline/jw03958-o007_t001_nircam_clear-f470n-nrcb_data_i2d.fits',
    'f480m': SICKLE_BASE / 'F480M/pipeline/jw03958-o007_t001_nircam_clear-f480m-nrcb_data_i2d.fits',
    # MIRI F770W data_i2d (satstar-replaced, obs 001 + 002 only)
    'f770w_o001': SICKLE_BASE / 'F770W/pipeline/jw03958-o001_t001_miri_clear-f770w-mirimage_data_i2d.fits',
    'f770w_o002': SICKLE_BASE / 'F770W/pipeline/jw03958-o002_t001_miri_clear-f770w-mirimage_data_i2d.fits',
    # MIRI F1130W/F1500W obs-level mosaics
    'f1130w_o001': SICKLE_BASE / 'F1130W/pipeline/jw03958-o001_t001_miri_f1130w_i2d.fits',
    'f1130w_o002': SICKLE_BASE / 'F1130W/pipeline/jw03958-o002_t001_miri_f1130w_i2d.fits',
    'f1130w_o003': SICKLE_BASE / 'F1130W/pipeline/jw03958-o003_t001_miri_f1130w_i2d.fits',
    'f1500w_o001': SICKLE_BASE / 'F1500W/pipeline/jw03958-o001_t001_miri_f1500w_i2d.fits',
    'f1500w_o002': SICKLE_BASE / 'F1500W/pipeline/jw03958-o002_t001_miri_f1500w_i2d.fits',
    'f1500w_o003': SICKLE_BASE / 'F1500W/pipeline/jw03958-o003_t001_miri_f1500w_i2d.fits',
}

OUT_DIR = SICKLE_BASE / 'sed_figures_top100_f1500w'
OUT_DIR.mkdir(exist_ok=True)
print(f'Output → {OUT_DIR}')
# ──────────────────────────────────────────────────────────────
# Load and combine MIRI catalogs from all 3 observations
# ──────────────────────────────────────────────────────────────
miri_tables = []
for obs, path in MIRI_CAT_FILES.items():
    t = Table.read(path)
    t['obs'] = obs
    # Rename F2550W columns if present
    for col in list(t.colnames):
        if col.lower().startswith('mag_'):
            t.rename_column(col, col.lower())
    miri_tables.append(t)

# Add missing columns so vstack works even if o001/o002 lack mag_f2550w
all_mag_cols = set()
for t in miri_tables:
    all_mag_cols |= {c for c in t.colnames if c.startswith('mag_')}
for t in miri_tables:
    for col in all_mag_cols:
        if col not in t.colnames:
            t[col] = np.nan

miri_all = vstack(miri_tables, join_type='outer')
miri_coords = SkyCoord(miri_all['ra'] * u.deg, miri_all['dec'] * u.deg)

print(f'Combined MIRI catalog: {len(miri_all)} sources from 3 obs')
for obs in ['o001', 'o002', 'o003']:
    mask = miri_all['obs'] == obs
    n15 = np.sum(np.isfinite(miri_all['mag_f1500w'][mask]))
    print(f'  {obs}: {mask.sum()} sources, {n15} with F1500W')
# ──────────────────────────────────────────────────────────────
# De-duplicate: sources within 0.3" of each other across obs
# Keep the row with the best (smallest) F1500W mag.
# ──────────────────────────────────────────────────────────────
MATCH_RADIUS = 0.3 * u.arcsec

# Coerce mag column to plain float array (handles masked arrays)
def _to_float_arr(col):
    return np.array([float(v) if (v is not None and v is not np.ma.masked) else np.nan
                     for v in col])

mag_f15_all = _to_float_arr(miri_all['mag_f1500w'])

# Self-match (nthneighbor=2 skips self)
idx_near, sep_near, _ = match_coordinates_sky(miri_coords, miri_coords, nthneighbor=2)
close_pairs = sep_near < MATCH_RADIUS

dropped = np.zeros(len(miri_all), dtype=bool)
for i in np.where(close_pairs)[0]:
    j = int(idx_near[i])
    if dropped[i] or dropped[j]:
        continue
    mi, mj = mag_f15_all[i], mag_f15_all[j]
    if np.isfinite(mi) and np.isfinite(mj):
        dropped[j if mi <= mj else i] = True
    elif np.isfinite(mi):
        dropped[j] = True
    elif np.isfinite(mj):
        dropped[i] = True
    else:
        dropped[j] = True

miri_dedup    = miri_all[~dropped]
mag_f15_dedup = mag_f15_all[~dropped]
print(f'After de-duplication: {len(miri_dedup)} unique MIRI sources')
f1500_valid = np.isfinite(mag_f15_dedup)
print(f'  With F1500W detection: {f1500_valid.sum()}')
# ──────────────────────────────────────────────────────────────
# Select top 100 brightest F1500W sources
# ──────────────────────────────────────────────────────────────
miri_f15    = miri_dedup[f1500_valid]
mags_f15    = mag_f15_dedup[f1500_valid]
sort_idx    = np.argsort(mags_f15)      # smallest mag = brightest
top100      = miri_f15[sort_idx[:100]]
top100_mags = mags_f15[sort_idx[:100]]

print(f'Top 100 F1500W: mag range [{top100_mags.min():.2f}, {top100_mags.max():.2f}]')
print(f'RA range:  [{float(top100["ra"].min()):.4f}, {float(top100["ra"].max()):.4f}]')
print(f'Dec range: [{float(top100["dec"].min()):.4f}, {float(top100["dec"].max()):.4f}]')
obs_counts = {obs: int(np.sum(top100['obs'] == obs)) for obs in ['o001','o002','o003']}
print('Obs breakdown:', obs_counts)
# ──────────────────────────────────────────────────────────────
# Load NIRCam merged catalog (prog 3958 + 2221 + 1182)
# ──────────────────────────────────────────────────────────────
nircam_cat   = Table.read(NIRCAM_CAT_FILE)
nircam_scs   = nircam_cat['skycoord_ref']
nircam_ra    = np.array([sc.ra.deg  for sc in nircam_scs])
nircam_dec   = np.array([sc.dec.deg for sc in nircam_scs])
nircam_coords = SkyCoord(nircam_ra * u.deg, nircam_dec * u.deg)

NIRCAM_FILTERS = ['f187n', 'f210m', 'f335m', 'f470n', 'f480m']
print(f'NIRCam catalog: {len(nircam_cat)} sources')
print(f'Filters: {NIRCAM_FILTERS}')
print(f'Flux columns available: {[c for c in nircam_cat.colnames if c.startswith("flux_jy")][:10]}')
# ──────────────────────────────────────────────────────────────
# Cross-match top-100 MIRI sources with NIRCam catalog
# ──────────────────────────────────────────────────────────────
top100_coords = SkyCoord(top100['ra'] * u.deg, top100['dec'] * u.deg)
idx, sep, _ = match_coordinates_sky(top100_coords, nircam_coords)

XMATCH_RADIUS = 0.5 * u.arcsec  # MIRI PSF is ~0.25" at F770W, ~0.5" at F1500W
has_nircam = sep < XMATCH_RADIUS
print(f'MIRI top-100 matched to NIRCam within {XMATCH_RADIUS}: {has_nircam.sum()} sources')
print(f'No NIRCam match: {(~has_nircam).sum()} sources')
# ──────────────────────────────────────────────────────────────
# Load image data (lazy, WCS only; data loaded on demand in cutout)
# ──────────────────────────────────────────────────────────────
def _get_sci_ext(path):
    """Return name of science extension."""
    with fits.open(path) as h:
        names = [x.name for x in h]
    return 'SCI' if 'SCI' in names else names[1]

def load_image(path):
    path = str(path)
    ext  = _get_sci_ext(path)
    with fits.open(path) as h:
        data = h[ext].data.copy().astype(float)
        wcs  = WCS(h[ext].header, naxis=2)
    return data, wcs

print('Loading images (this takes a minute)...')
image_data = {}
for key, path in IMAGE_FILES.items():
    if not Path(path).exists():
        print(f'  MISSING: {key}')
        continue
    try:
        image_data[key] = load_image(path)
        print(f'  OK {key}: {image_data[key][0].shape}')
    except Exception as e:
        print(f'  ERROR {key}: {e}')
# ──────────────────────────────────────────────────────────────
# Flux-conversion helpers
# ──────────────────────────────────────────────────────────────
def vega_mag_to_jy(mag, filt):
    """Convert Vega magnitude to Jy using MIRI Vega zero-points."""
    zp = MIRI_VEGA_ZP_JY.get(filt.lower())
    mag = float(mag) if mag is not None else np.nan
    if zp is None or not np.isfinite(mag):
        return np.nan
    return zp * 10.0 ** (-mag / 2.5)

def _safe_float(val):
    if val is None or val is np.ma.masked:
        return np.nan
    try:
        v = float(val)
        return v if np.isfinite(v) else np.nan
    except (TypeError, ValueError):
        return np.nan

def get_source_fluxes(i_top100):
    """Return dict of filt -> (wave_um, flux_jy, eflux_jy) for source i."""
    row    = top100[i_top100]
    fluxes = {}

    # MIRI (from merged catalog, Vega mags)
    for filt in MIRI_VEGA_ZP_JY:
        col = f'mag_{filt}'
        if col not in top100.colnames:
            continue
        f = vega_mag_to_jy(_safe_float(row[col]), filt)
        if np.isfinite(f) and f > 0:
            fluxes[filt] = (WAVELENGTHS_UM[filt], f, np.nan)

    # NIRCam (from iterative merged catalog, already in Jy)
    if has_nircam[i_top100]:
        nc_row = nircam_cat[int(idx[i_top100])]
        for filt in NIRCAM_FILTERS:
            fcol = f'flux_jy_{filt}'
            ecol = f'eflux_jy_{filt}'
            if fcol not in nircam_cat.colnames:
                continue
            f  = _safe_float(nc_row[fcol])
            ef = _safe_float(nc_row[ecol]) if ecol in nircam_cat.colnames else np.nan
            if np.isfinite(f) and f > 0:
                fluxes[filt] = (WAVELENGTHS_UM[filt], f, ef)

    return fluxes
# ──────────────────────────────────────────────────────────────
# Cutout helpers
# ──────────────────────────────────────────────────────────────
CUTOUT_SIZES = {
    'nircam': 2.0 * u.arcsec,
    'miri':   4.0 * u.arcsec,
}

# Mapping: abstract filter → list of image_data keys to try in order
FILTER_IMG_KEYS = {
    'f187n':  ['f187n'],
    'f210m':  ['f210m'],
    'f335m':  ['f335m'],
    'f470n':  ['f470n'],
    'f480m':  ['f480m'],
    'f770w':  ['f770w_o001', 'f770w_o002'],
    'f1130w': ['f1130w_o001', 'f1130w_o002', 'f1130w_o003'],
    'f1500w': ['f1500w_o001', 'f1500w_o002', 'f1500w_o003'],
}

def get_cutout(coord, filt):
    """Return (cutout_data, pix_scale_arcsec) for the best available obs."""
    size = CUTOUT_SIZES['miri'] if filt.startswith('f7') or filt.startswith('f1') else CUTOUT_SIZES['nircam']
    for key in FILTER_IMG_KEYS.get(filt, []):
        if key not in image_data:
            continue
        data, wcs = image_data[key]
        try:
            x, y = wcs.world_to_pixel(coord)
            if not (0 <= x < data.shape[1] and 0 <= y < data.shape[0]):
                continue
            cut = Cutout2D(data, (x, y), size, wcs=wcs, mode='partial', fill_value=np.nan)
            pix_scale = np.sqrt(wcs.proj_plane_pixel_area().to(u.arcsec**2).value)
            return cut.data, pix_scale
        except Exception:
            continue
    return None, None
# ──────────────────────────────────────────────────────────────
# Main SED + cutout plot function
# ──────────────────────────────────────────────────────────────
DISPLAY_FILTERS = ['f187n', 'f210m', 'f335m', 'f470n', 'f480m', 'f770w', 'f1130w', 'f1500w']

def plot_source(i_top100, save_dir=None):
    row   = top100[i_top100]
    coord = SkyCoord(float(row['ra']) * u.deg, float(row['dec']) * u.deg)
    obs   = str(row['obs'])
    rank  = i_top100 + 1
    mag15 = float(top100_mags[i_top100])

    fluxes = get_source_fluxes(i_top100)

    n_filt = len(DISPLAY_FILTERS)
    n_top  = math.ceil(n_filt / 2)
    n_bot  = n_filt - n_top

    cut_w, cut_h = 0.10, 0.20
    h_pad = 0.005
    row_margin, row_gap, sed_h = 0.03, 0.04, 0.30

    bot_bot = row_margin
    bot_top = bot_bot + cut_h
    sed_bot = bot_top + row_gap
    sed_top = sed_bot + sed_h
    top_bot = sed_top + row_gap
    top_top = top_bot + cut_h
    fig_h_f = top_top + row_margin

    fig_w = max(14, n_top * 1.6)
    fig_h = max(fig_w * fig_h_f * 0.65, 7)
    fig = plt.figure(figsize=(fig_w, fig_h))

    # SED panel
    row_block_w = n_top * cut_w + (n_top - 1) * h_pad
    sed_w  = min(0.42, row_block_w * 0.9)
    ax_sed = fig.add_axes([(1 - sed_w) / 2, sed_bot, sed_w, sed_h])

    if fluxes:
        sorted_f  = sorted(fluxes.values(), key=lambda x: x[0])
        waves     = [f[0] for f in sorted_f]
        fvals     = [f[1] for f in sorted_f]
        ferrs     = [f[2] if np.isfinite(f[2]) else 0.0 for f in sorted_f]
        ax_sed.errorbar(waves, fvals, yerr=ferrs,
                        fmt='o-', capsize=3, lw=1.2, ms=5, color='C0')
        ax_sed.set_xscale('log')
        ax_sed.set_yscale('log')

    ax_sed.set_xlabel(r'Wavelength [$\mu$m]', fontsize=9)
    ax_sed.set_ylabel(r'Flux Density [Jy]',   fontsize=9)
    nc_str = (f'  NIRCam sep={sep[i_top100].to(u.mas).value:.0f} mas'
              if has_nircam[i_top100] else '  no NIRCam match')
    ax_sed.set_title(
        f'Rank {rank:03d}  |  F1500W = {mag15:.2f} Vega  |  MIRI {obs}\n'
        f'RA={float(row["ra"]):.5f}  Dec={float(row["dec"]):.5f}{nc_str}',
        fontsize=8
    )

    # Cutout panels
    top_bands = DISPLAY_FILTERS[:n_top]
    bot_bands = DISPLAY_FILTERS[n_top:]

    for band_list, row_bottom in [(top_bands, top_bot), (bot_bands, bot_bot)]:
        nb = len(band_list)
        rw = nb * cut_w + (nb - 1) * h_pad
        x0 = (1 - rw) / 2
        for j, filt in enumerate(band_list):
            ax = fig.add_axes([x0 + j * (cut_w + h_pad), row_bottom, cut_w, cut_h])
            cut_d, pix_sc = get_cutout(coord, filt)
            if cut_d is not None:
                finite = np.isfinite(cut_d)
                if finite.any():
                    norm = simple_norm(cut_d[finite], stretch='log', percent=99.5)
                    ax.imshow(cut_d, origin='lower', cmap='inferno', norm=norm)
                    aper_as = 0.4 if filt.startswith('f7') or filt.startswith('f1') else 0.2
                    aper_px = aper_as / pix_sc
                    cx, cy = cut_d.shape[1] / 2, cut_d.shape[0] / 2
                    ax.add_patch(Circle((cx, cy), aper_px,
                                        edgecolor='cyan', facecolor='none', lw=0.7))
            ax.set_title(filt.upper(), fontsize=7)
            ax.set_xticks([])
            ax.set_yticks([])

    if save_dir is not None:
        out = Path(save_dir) / f'sed_rank{rank:03d}_{obs}_F1500W{mag15:.2f}.png'
        fig.savefig(out, dpi=120, bbox_inches='tight')
        plt.close(fig)
        return str(out)
    return fig
# Quick preview of the top-5 sources
for i in range(min(5, len(top100))):
    fig = plot_source(i)
    plt.show()
    plt.close()
# ──────────────────────────────────────────────────────────────
# Save all 100 SEDs
# ──────────────────────────────────────────────────────────────
for i in range(len(top100)):
    out = plot_source(i, save_dir=OUT_DIR)
    if (i + 1) % 10 == 0:
        print(f'  {i+1}/100  → {out}')

print(f'\nAll SEDs saved to {OUT_DIR}')
# ──────────────────────────────────────────────────────────────
# Summary table of top-100 sources with all catalog fluxes
# ──────────────────────────────────────────────────────────────
from astropy.table import Column

rows = []
for i in range(len(top100)):
    row  = top100[i]
    d    = {'rank': i+1, 'ra': float(row['ra']), 'dec': float(row['dec']),
            'obs': str(row['obs']),
            'mag_f1500w': float(row['mag_f1500w']),
            'has_nircam': bool(has_nircam[i])}
    # MIRI fluxes in Jy
    for filt in ['f770w', 'f1130w', 'f1500w', 'f2550w']:
        col = f'mag_{filt}'
        mag = float(row[col]) if col in row.colnames else np.nan
        d[f'flux_jy_{filt}'] = vega_mag_to_jy(mag, filt)
    # NIRCam fluxes in Jy
    if has_nircam[i]:
        nc_row = nircam_cat[idx[i]]
        for filt in NIRCAM_FILTERS:
            fcol = f'flux_jy_{filt}'
            val  = float(nc_row[fcol]) if fcol in nc_row.colnames and nc_row[fcol] is not None else np.nan
            d[f'flux_jy_{filt}'] = val
    else:
        for filt in NIRCAM_FILTERS:
            d[f'flux_jy_{filt}'] = np.nan
    rows.append(d)

summary = Table(rows)
out_cat = OUT_DIR / 'top100_f1500w_summary.fits'
summary.write(out_cat, overwrite=True)
print(f'Summary table → {out_cat}')
summary[:5]
#[Out]# <Table length=5>
#[Out]#  rank         ra                 dec         ... flux_jy_f470n flux_jy_f480m
#[Out]# int64      float64             float64       ...    float64       float64   
#[Out]# ----- ------------------ ------------------- ... ------------- -------------
#[Out]#     1 266.57825217882845 -28.802398932631192 ...           0.0           0.0
#[Out]#     2 266.56232243105245   -28.8073646620411 ...           nan           nan
#[Out]#     3 266.57818149494943  -28.79817218587603 ...           nan           nan
#[Out]#     4 266.58149819615204 -28.811690264110194 ...           nan           nan
#[Out]#     5  266.5685100411142 -28.804432020702322 ...           nan           nan