#!/usr/bin/env python
"""P&G accretion-reservoir test on REVEALED O stars: measure ALMA continuum at the
star position. Detection -> circumstellar dust/ionized gas mass; non-detection -> 3sigma
mass upper limit. Answers: is there ANY reservoir left to keep accreting?

For each target, pick the best covering continuum image (on-source, smallest beam),
measure peak in a 1-beam aperture at the star + image noise, convert to gas mass."""
import warnings, numpy as np
warnings.filterwarnings('ignore')
from astropy.table import Table
from astropy.coordinates import SkyCoord
from astropy.io import fits
from astropy.wcs import WCS
import astropy.units as u
import astropy.constants as const

ai = Table.read('alma_footprint_index.fits')
ic = SkyCoord(np.array(ai['crval1'], float), np.array(ai['crval2'], float), unit='deg')
half = np.minimum(np.array(ai['naxis1'], float), np.array(ai['naxis2'], float)) * \
       np.array(ai['cdelt_as'], float) / 2.0
freq = np.array(ai['freq_ghz'], float); beam = np.array(ai['beam_as'], float)

TARGETS = [
    ('M8/Her36',       SkyCoord('18:03:40.3', '-24:22:43', unit=(u.hourangle, u.deg)), 1.34),
    ('M20/HD164492A',  SkyCoord('18:02:23.5', '-23:01:51', unit=(u.hourangle, u.deg)), 1.42),
    ('NGC6357/Pis24-1',SkyCoord('17:24:43.5', '-34:11:57', unit=(u.hourangle, u.deg)), 1.70),
    ('W33Main',        SkyCoord('18:14:13.5', '-17:55:47', unit=(u.hourangle, u.deg)), 2.40),
]

def gas_mass(S_Jy, dist_kpc, freq_ghz, T=30.0):
    """Optically-thin dust -> total (gas+dust) mass. kappa OH94 thin-ice scaled, beta=1.8."""
    nu = freq_ghz * 1e9
    kappa = 0.0099 * (freq_ghz / 230.0) ** 1.8     # cm^2/g of gas (dust/gas=1/100, OH94)
    B = (2 * const.h.value * nu**3 / const.c.value**2) / \
        (np.exp(const.h.value * nu / (const.k_B.value * T)) - 1)   # W/m^2/Hz/sr
    B_cgs = B * 1e3                                  # erg/s/cm^2/Hz/sr
    S_cgs = S_Jy * 1e-23                             # erg/s/cm^2/Hz
    d_cm = dist_kpc * 3.086e21
    M_g = S_cgs * d_cm**2 / (kappa * B_cgs)
    return M_g / 1.989e33                            # Msun

print(f"{'target':18s}{'band/GHz':>9s}{'beam':>6s}{'peak@*':>9s}{'noise':>8s}{'SNR':>6s} "
      f"{'M_gas(3sig if<3)':>18s}")
for nm, c, dist in TARGETS:
    sep = c.separation(ic).arcsec
    cov = np.where(sep < 0.6 * half)[0]              # well inside
    if len(cov) == 0:
        print(f"{nm:18s}  no on-source image"); continue
    # prefer B6/B7 (dust); fall back any; among those, best (smallest) beam, closest
    score = []
    for j in cov:
        b = freq[j]
        pref = 0 if 211 <= b <= 373 else (1 if 125 <= b <= 211 else 2)
        score.append((pref, beam[j], sep[j], j))
    score.sort()
    j = score[0][3]
    path = str(ai['path'][j])
    try:
        hd = fits.open(path)[0]
        data = np.squeeze(hd.data).astype(float)
        w = WCS(hd.header).celestial
        x, y = w.world_to_pixel(c)
        xi, yi = int(round(float(x))), int(round(float(y)))
        ny, nx = data.shape
        if not (0 <= xi < nx and 0 <= yi < ny):
            print(f"{nm:18s}  star off image {path.split('/')[-1][:40]}"); continue
        # peak in 1-beam box around star
        pixscale = abs(ai['cdelt_as'][j])
        r = max(1, int(round(beam[j] / pixscale)))
        sub = data[max(0,yi-r):yi+r+1, max(0,xi-r):xi+r+1]
        peak = np.nanmax(sub)
        noise = np.nanstd(data[np.isfinite(data)])    # crude image rms
        # robust noise via MAD
        d = data[np.isfinite(data)]
        noise = 1.4826 * np.nanmedian(np.abs(d - np.nanmedian(d)))
        snr = peak / noise if noise > 0 else np.nan
        if snr >= 3:
            M = gas_mass(peak, dist, freq[j])
            mstr = f"{M:.2f} Msun (det)"
        else:
            M = gas_mass(3 * noise, dist, freq[j])
            mstr = f"<{M:.2f} (3sig UL)"
        print(f"{nm:18s}{freq[j]:9.0f}{beam[j]:6.2f}{peak*1e3:8.2f}m{noise*1e3:7.2f}m{snr:6.1f} {mstr:>18s}")
    except (FileNotFoundError, OSError) as e:
        print(f"{nm:18s}  unreadable: {str(e)[:40]}")