#!/usr/bin/env python
"""ALMA-on-O-star P&G reservoir test for the REVEALED O stars.
Per target: cutout grid [2MASS-K | ALMA bands | GLOSTAR 5.8GHz], O-star marked.
GLOSTAR (assumed optically-thin free-free) predicts ALMA free-free
  S_ff(nu) = S_5.8GHz * (nu/5.8GHz)^-0.1
which is subtracted from the ALMA peak -> dust excess -> circumstellar mass.
"""
import warnings, glob, 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.hips2fits import hips2fits

GLOSTAR_DIR = '/orange/adamginsburg/galactic_plane_surveys/glostar/glostar_continuum'
NU_GLO = 5.8  # GHz
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)

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),
]
FOV = 30 * u.arcsec
os.makedirs('pg_ostar_figs', exist_ok=True)

# GLOSTAR tiles: cache WCS + shape (CRVAL is offset from center, so match by pixel-in-range)
glo_files = glob.glob(os.path.join(GLOSTAR_DIR, 'glostar_B_*FREQ0.fits'))
glo = []
for f in glo_files:
    h = fits.getheader(f)
    glo.append((WCS(h).celestial, int(h['NAXIS1']), int(h['NAXIS2']), f))

def glostar_tile(c):
    """Return path of the GLOSTAR tile most centrally covering c, else None."""
    best = None
    for w, n1, n2, f in glo:
        x, y = w.world_to_pixel(c)
        if 0 <= x < n1 and 0 <= y < n2:
            d = (x - n1/2)**2 + (y - n2/2)**2     # prefer most central
            if best is None or d < best[0]:
                best = (d, f)
    return best[1] if best else None

def cutout(path, center, fov, hdu=0):
    h = fits.open(path)[hdu]
    data = np.squeeze(h.data).astype(float)
    w = WCS(h.header).celestial
    try:
        co = Cutout2D(data, center, fov, wcs=w, mode='partial')
        return co.data, co.wcs
    except (ValueError, Exception):
        return None, None

def peakflux(path, center, hdu=0, rad_as=1.5):
    """peak Jy/beam within rad_as of center, + robust noise (sigma-clipped)."""
    h = fits.open(path)[hdu]
    data = np.squeeze(h.data).astype(float)
    w = WCS(h.header).celestial
    x, y = w.world_to_pixel(center)
    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, None
    ps = np.abs(w.proj_plane_pixel_scales()[0].to(u.arcsec).value)
    r = max(1, int(round(rad_as / ps)))
    sub = data[max(0, yi-r):yi+r+1, max(0, xi-r):xi+r+1]
    peak = np.nanmax(sub)
    d = data[np.isfinite(data)]
    med = np.nanmedian(d)
    for _ in range(3):
        s = 1.4826 * np.nanmedian(np.abs(d - med)); d = d[np.abs(d - med) < 3*s]; med = np.nanmedian(d)
    noise = 1.4826 * np.nanmedian(np.abs(data[np.isfinite(data)] - med))
    return peak, noise

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

print(f"{'target':16s}{'nu':>6s}{'S_ALMA':>9s}{'S_glo5.8':>10s}{'S_ff_pred':>10s}{'S_dust':>9s}{'M_dust':>10s}")
for name, c, dist in TARGETS:
    # ALMA: best dust band (B6/7) on source
    cov = np.where(c.separation(aic).arcsec < 0.6*ahalf)[0]
    panels = []  # (label, data, wcs, marker)
    # 2MASS K
    try:
        hdu2 = hips2fits.query(hips='CDS/P/2MASS/K', width=200, height=200,
                               ra=c.ra, dec=c.dec, fov=FOV, projection='TAN', format='fits')
        panels.append(('2MASS-K', np.squeeze(hdu2[0].data).astype(float), WCS(hdu2[0].header).celestial))
    except Exception as e:
        print(f"  {name} 2MASS fail {str(e)[:30]}")
    # ALMA panels: pick best B3-ish and best B6/7
    alma_meas = []
    chosen = {}
    for j in cov:
        b = afreq[j]
        key = 'low' if b < 180 else 'high'
        if key not in chosen or ai['beam_as'][j] < ai['beam_as'][chosen[key]]['beam_as'] if False else key not in chosen:
            chosen[key] = j
    for key, j in chosen.items():
        p = str(ai['path'][j]); d, w = cutout(p, c, FOV)
        if d is not None:
            panels.append((f"ALMA {afreq[j]:.0f}GHz", d, w))
        pk, ns = peakflux(p, c)
        if pk is not None:
            alma_meas.append((afreq[j], pk, ns))
    # GLOSTAR: tile that actually covers the star (pixel-in-range), else outside coverage
    Sglo = None
    gp = glostar_tile(c)
    if gp is not None:
        Sglo, gns = peakflux(gp, c)
        d, w = cutout(gp, c, FOV)
        if d is not None:
            panels.append(('GLOSTAR 5.8GHz', d, w))
    # figure grid
    n = len(panels); fig, axes = plt.subplots(1, n, figsize=(3.2*n, 3.4))
    if n == 1: axes = [axes]
    for ax, (lab, d, w) in zip(axes, panels):
        ax2 = plt.subplot(1, n, list(panels).index((lab, d, w))+1, projection=w) if False else ax
        vmin, vmax = np.nanpercentile(d[np.isfinite(d)], [5, 99.5]) if np.isfinite(d).any() 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=14, mew=2)
        ax.set_title(lab, fontsize=9); ax.set_xticks([]); ax.set_yticks([])
    fig.suptitle(f"{name}  (d={dist} kpc)  cyan + = O-star", fontsize=11)
    fig.tight_layout(rect=[0, 0, 1, 0.93])
    fig.savefig(f'pg_ostar_figs/{name}.png', dpi=130); plt.close(fig)
    # free-free subtraction per ALMA band
    for nu, Salma, ns in alma_meas:
        if Sglo is not None and Sglo > 0:
            Sff = Sglo * (nu/NU_GLO)**-0.1
            Sdust = Salma - Sff
            md = gasmass(max(Sdust, 0), dist, nu) if Sdust > 3*ns else gasmass(3*ns, dist, nu)
            tag = f"{md:.3f}" + ("" if Sdust > 3*ns else "(UL)")
            print(f"{name:16s}{nu:6.0f}{Salma*1e3:8.2f}m{Sglo*1e3:9.2f}m{Sff*1e3:9.2f}m{Sdust*1e3:8.2f}m{tag:>10s}")
        else:
            md = gasmass(Salma, dist, nu)
            print(f"{name:16s}{nu:6.0f}{Salma*1e3:8.2f}m{'noGLO':>10s}{'--':>10s}{'--':>9s}{md:>9.3f}*")
print('figs -> pg_ostar_figs/   (* = no GLOSTAR, ALMA flux as dust upper bound incl free-free)')