#!/usr/bin/env python
"""Phase 6 (imaging): explicit accretion signatures at the ionizing source from
ALMA-IMF moment-0 maps. Hot-core (CH3CN, esp v8=1) => embedded accreting protostar;
RRL (H41a/H30a) => ionized gas; N2H+ => cold dense gas; SiO => outflow/shock."""
import os, glob, numpy as np
from astropy.table import Table
from astropy.io import fits
from astropy.wcs import WCS
from astropy.coordinates import SkyCoord
from astropy.nddata import Cutout2D
import astropy.units as u
LOG=open('phase6_lines.log','w')
def p(*a):
    s=' '.join(str(x) for x in a);print(s);print(s,file=LOG);LOG.flush()
R='/orange/adamginsburg/ALMAIMF_Images'

m=Table.read('master_phase1_classified.fits')
hii=np.array([str(x) for x in m['simple_type']])=='HII region'
m=m[hii]
mc=SkyCoord(np.array(m['GLON'],float)*u.deg,np.array(m['GLAT'],float)*u.deg,frame='galactic')
name2i={str(m['Name'][i]):i for i in range(len(m))}

# 8 sources measured in common-beam alpha
TARGS=[('W51-IRS2','G049.4885-00.3799'),('W51-IRS2','G049.4903-00.3694'),
       ('W51-E','G049.4885-00.3799'),('G010.62','G010.6235-00.3834'),
       ('G010.62','G010.6260-00.3744'),('G010.62','G010.6311-00.3864'),
       ('G327.29','G327.2852-00.5735'),('G328.25','G328.2523-00.5320C'),
       ('G338.93','G338.9237+00.5618')]
# line -> (band, filename token)
LINEDEF=[('H41a_RRL','B3','.h41a.'),('CH3CN_hotcore','B3','.ch3cn.'),
         ('CH3CNv8_HOT','B3','ch3cnv8'),('N2Hp_colddense','B3','.n2hp.'),
         ('H2CO_dense','B6','h2co_303-202'),('SiO_outflow','B6','.sio.')]

def find_m0(field,band,tok):
    d=os.path.join(R,field,band,'moments')
    if not os.path.isdir(d): return None
    c=glob.glob(os.path.join(d,f'*{tok}m0.fits'))
    if not c: c=glob.glob(os.path.join(d,f'*{tok}*m0.fits'))
    pb=[f for f in c if 'pbcor' in os.path.basename(f)]
    c=pb if pb else c
    return c[0] if c else None

def peak_sn(f,sc):
    """peak within 2.5" of the ionizing position, in S/N units (robust noise
    from the wider field)."""
    try:
        h=fits.open(f)[0]; d=np.squeeze(h.data); w=WCS(h.header).celestial
        wide=Cutout2D(d,sc,20*u.arcsec,wcs=w).data.astype(float)
        noise=np.nanmedian(np.abs(wide-np.nanmedian(wide)))/0.6745
        core=Cutout2D(d,sc,2.5*u.arcsec,wcs=w).data.astype(float)
        if not np.isfinite(core).any() or not np.isfinite(noise) or noise<=0: return np.nan
        return float(np.nanmax(core)/noise)
    except Exception:
        return np.nan

rows=[]
for field,nm in TARGS:
    if nm not in name2i: continue
    sc=mc[name2i[nm]].icrs
    vals={}
    for lab,band,tok in LINEDEF:
        f=find_m0(field,band,tok)
        vals[lab]=peak_sn(f,sc) if f else np.nan   # S/N within 2.5"
    SN=5.0
    det=lambda k: np.isfinite(vals[k]) and vals[k]>SN
    hot = det('CH3CN_hotcore') or det('CH3CNv8_HOT')   # hot core => embedded accreting protostar
    rrl = det('H41a_RRL')                              # ionized gas
    verdict = ('HOT CORE + ionized (active, cluster-confused)' if (hot and rrl) else
               ('HOT CORE (embedded/accreting)' if hot else
                ('ionized only (post-accretion: free-free, no hot core)' if rrl else 'weak/none (>5sigma)')))
    rows.append((nm,field,vals['H41a_RRL'],vals['CH3CN_hotcore'],vals['CH3CNv8_HOT'],
                 vals['N2Hp_colddense'],vals['H2CO_dense'],vals['SiO_outflow'],verdict))
    p(f"{nm:22s} H41a={vals['H41a_RRL']:.3g} CH3CN={vals['CH3CN_hotcore']:.3g} "
      f"v8={vals['CH3CNv8_HOT']:.3g} N2H+={vals['N2Hp_colddense']:.3g} "
      f"H2CO={vals['H2CO_dense']:.3g} SiO={vals['SiO_outflow']:.3g}  => {verdict}")

if rows:
    T=Table(rows=rows,names=['Name','field','H41a','CH3CN','CH3CNv8','N2Hp','H2CO','SiO','verdict'])
    T.write('phase6_line_signatures.fits',overwrite=True)
    p('\n=== Phase 6 summary ===')
    v=[r[8] for r in rows]
    from collections import Counter
    for k,c in Counter(v).most_common(): p(f'  {c}  {k}')
    p('Wrote phase6_line_signatures.fits')
LOG.close()