#!/usr/bin/env python
"""Rigorous ALMA-on-O-star P&G reservoir test (precise SIMBAD positions).
Per target: panel grid [2MASS-K | WISE-3.4 | ALMA bands] with O-star marked.
Photometry at the exact star position: aperture peak + annulus background + noise
-> CONSERVATIVE dust-mass (all flux treated as dust = upper bound; free-free only
lowers it) + 3sigma upper limit -> max sustainable accretion time vs P&G ~4 Myr.
Where 2 ALMA bands exist, the spectral index separates free-free (a<=~1) from dust (a~3-4).
"""
import warnings, os
import numpy as np
warnings.filterwarnings('ignore')
import matplotlib; matplotlib.use('Agg')
import matplotlib.pyplot as plt
from astropy.table import Table
from astropy.coordinates import SkyCoord
from astropy.io import fits
from astropy.wcs import WCS
from astropy.nddata import Cutout2D
import astropy.units as u
import astropy.constants as const
from astroquery.skyview import SkyView

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

# precise SIMBAD/literature ICRS positions
TARGETS = [
    ('M8_Her36',        SkyCoord(270.91792, -24.37858, unit='deg'), 1.34, 'O7.5V revealed'),
    ('M20_HD164492A',   SkyCoord(270.59833, -23.03083, unit='deg'), 1.42, 'O7.5Vz revealed'),
    ('NGC6357_Pis24-1', SkyCoord(261.18125, -34.19917, unit='deg'), 1.70, 'O3.5If revealed'),
    ('W33Main',         SkyCoord(273.55750, -17.92920, unit='deg'), 2.40, 'EMBEDDED HII (not revealed)'),
]
FOV = 20 * u.arcsec
os.makedirs('pg_ostar_figs', exist_ok=True)

def gasmass(S_Jy, dist_kpc, freq_ghz, T=30.0):
    nu = freq_ghz*1e9; kappa = 0.0099*(freq_ghz/230.)**1.8
    B = (2*const.h.value*nu**3/const.c.value**2)/(np.exp(const.h.value*nu/(const.k_B.value*T))-1)*1e3
    return (S_Jy*1e-23)*(dist_kpc*3.086e21)**2/(kappa*B)/1.989e33

def photometry(path, c, beam_as):
    """aperture(1.5*beam) peak bg-subtracted + 3sig, using annulus 4-8*beam."""
    h = fits.open(path)[0]; data = np.squeeze(h.data).astype(float); w = WCS(h.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): return None
    ps = np.abs(w.proj_plane_pixel_scales()[0].to(u.arcsec).value)
    yy, xx = np.mgrid[0:ny, 0:nx]; rr = np.hypot(xx-xi, yy-yi)*ps
    ap = rr < 1.5*beam_as; ann = (rr > 4*beam_as) & (rr < 8*beam_as) & np.isfinite(data)
    if ann.sum() < 20: ann = (rr > 2*beam_as) & (rr < 6*beam_as) & np.isfinite(data)
    bg = np.nanmedian(data[ann]); noise = 1.4826*np.nanmedian(np.abs(data[ann]-bg))
    peak = np.nanmax(data[ap]) - bg
    return dict(peak=peak, noise=noise, snr=peak/noise if noise>0 else np.nan, ps=ps)

def cut(path, c, hdu=0):
    h = fits.open(path)[hdu]; d = np.squeeze(h.data).astype(float); w = WCS(h.header).celestial
    try: co = Cutout2D(d, c, FOV, wcs=w, mode='partial'); return co.data, co.wcs
    except Exception: return None, None

print(f"{'target':18s}{'spectype':16s}{'nu':>5s}{'beam':>6s}{'Speak':>9s}{'SNR':>6s}{'Mdust(<3sig UL)':>18s}{'t_acc(yr)':>11s}")
for name, c, dist, spt in TARGETS:
    panels = []
    for surv, lab in [('2MASS-K', '2MASS-K'), ('WISE 3.4', 'WISE-3.4um')]:
        try:
            im = SkyView.get_images(position=c, survey=[surv], radius=FOV, pixels=150)[0][0]
            panels.append((lab, np.squeeze(im.data).astype(float), WCS(im.header).celestial))
        except Exception:
            pass
    cov = np.where(c.separation(aic).arcsec < 0.6*ahalf)[0]
    chosen = {}
    for j in cov:
        key = 'B3' if afreq[j] < 130 else ('B6' if afreq[j] < 280 else 'B7')
        if key not in chosen or abeam[j] < abeam[chosen[key]]:
            chosen[key] = j
    meas = []
    for key, j in sorted(chosen.items()):
        p = str(ai['path'][j]); d, w = cut(p, c)
        if d is not None: panels.append((f"ALMA {afreq[j]:.0f}GHz", d, w))
        ph = photometry(p, c, abeam[j])
        if ph: meas.append((afreq[j], abeam[j], ph))
    # figure
    n = len(panels); fig, axes = plt.subplots(1, n, figsize=(3.0*n, 3.3)); axes = np.atleast_1d(axes)
    for ax, (lab, d, w) in zip(axes, panels):
        v = d[np.isfinite(d)]; vmin, vmax = (np.nanpercentile(v, [10, 99.5]) if v.size else (0, 1))
        ax.imshow(d, origin='lower', cmap='inferno', vmin=vmin, vmax=vmax)
        xs, ys = w.world_to_pixel(c); ax.plot(xs, ys, '+', color='cyan', ms=16, mew=2)
        ax.set_title(lab, fontsize=9); ax.set_xticks([]); ax.set_yticks([])
    fig.suptitle(f"{name}  {spt}  d={dist}kpc  (cyan + = O-star; FOV {FOV})", fontsize=10)
    fig.tight_layout(rect=[0, 0, 1, 0.92]); fig.savefig(f'pg_ostar_figs/{name}_rig.png', dpi=135); plt.close(fig)
    # photometry table + masses (conservative: all flux as dust)
    for nu, bm, ph in meas:
        det = ph['snr'] >= 3
        S = ph['peak'] if det else 3*ph['noise']
        M = gasmass(max(S, 3*ph['noise']) if det else 3*ph['noise'], dist, nu)
        tacc = M/1e-4  # yr at 1e-4 Msun/yr
        tag = f"{M:.3f}{'' if det else ' UL'}"
        print(f"{name:18s}{spt:16s}{nu:5.0f}{bm:6.2f}{ph['peak']*1e3:8.2f}m{ph['snr']:6.1f}{tag:>18s}{tacc:>11.0f}")
    # spectral index if 2 bands
    if len(meas) >= 2:
        (n1, _, p1), (n2, _, p2) = meas[0], meas[-1]
        if p1['peak'] > 0 and p2['peak'] > 0:
            a = np.log(p2['peak']/p1['peak'])/np.log(n2/n1)
            kind = 'FREE-FREE' if a < 1.5 else 'DUST'
            print(f"{'':18s}spectral index {n1:.0f}->{n2:.0f} GHz: alpha={a:+.2f}  => {kind}")
print('figs -> pg_ostar_figs/*_rig.png   (t_acc = reservoir / 1e-4 Msun/yr; P&G needs ~4e6 yr)')