#!/usr/bin/env python
"""Phase 6 infall kinematics. Extract optically-thick (H2CO 303-202) and
optically-thin (C18O 2-1) spectra at each ionizing source from ALMA-IMF line
cubes (memmap; reads only the needed spectral column). Infall = blue-asymmetric
thick line: dv = (V_thick - V_thin)/dV_thin < ~-0.25."""
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
import astropy.units as u
import warnings; warnings.filterwarnings('ignore')
LOG=open('phase6_infall.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'
C=299792.458

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))}

TARGS=[('W51-IRS2','G049.4903-00.3694'),('W51-IRS2','G049.4885-00.3799'),
       ('G010.62','G010.6235-00.3834'),('G010.62','G010.6311-00.3864'),
       ('G338.93','G338.9237+00.5618')]

RESTGHZ={'h2co':218.22219,'c18o':219.56035}  # GHz; H2CO 3(03)-2(02), C18O 2-1
SPW={'h2co':'spw3','c18o':'spw4'}

def find_cube(field,line):
    spw=SPW[line]
    for sub in ['linecubes_12m','B6/cleanest','fullcubes_12m','B6']:
        d=os.path.join(R,field,sub)
        if os.path.isdir(d):
            c=glob.glob(os.path.join(d,f'*B6_{spw}_*.image.pbcor.fits'))
            if not c: c=glob.glob(os.path.join(d,f'*B6_{spw}_*.image.fits'))
            c=[x for x in c if 'JvM' in x] or c
            if c: return sorted(c,key=len)[0]
    return None

def spectrum(f,sc,line):
    """velocity[km/s] (rel. line rest freq), intensity at nearest pixel, memmap."""
    h=fits.open(f,memmap=True)[0]; hdr=h.header
    w=WCS(hdr); wc=w.celestial
    x,y=wc.world_to_pixel(sc); ix,iy=int(round(float(x))),int(round(float(y)))
    nx=hdr['NAXIS1']; ny=hdr['NAXIS2']
    if not (0<=ix<nx and 0<=iy<ny): return None,None
    spec=np.squeeze(h.data[...,iy,ix]).astype(float)
    nch=spec.shape[0]
    nu=w.spectral.pixel_to_world(np.arange(nch)).to(u.Hz).value
    rest=RESTGHZ[line]*1e9
    vel=C*(rest-nu)/rest
    return vel,spec

def moments(vel,spec,vwin=None):
    s=spec.copy()
    noise=np.nanmedian(np.abs(s-np.nanmedian(s)))/0.6745
    msk=np.isfinite(s)
    if vwin is not None:
        msk &= (vel>vwin[0])&(vel<vwin[1])
    good=msk&(s>3*noise)
    if good.sum()<3: return np.nan,np.nan,np.nan,noise
    w=s[good]; v=vel[good]
    Vc=np.sum(w*v)/np.sum(w)
    dV=np.sqrt(np.sum(w*(v-Vc)**2)/np.sum(w))
    return Vc, dV, np.nanmax(s[msk])/noise, noise

rows=[]
for field,nm in TARGS:
    if nm not in name2i: continue
    sc=mc[name2i[nm]].icrs
    fthick=find_cube(field,'h2co'); fthin=find_cube(field,'c18o')
    p(f'\n=== {nm} ({field}) ===')
    p(f'  thick(H2CO): {os.path.basename(fthick) if fthick else None}')
    p(f'  thin(C18O):  {os.path.basename(fthin) if fthin else None}')
    if not (fthick and fthin): p('  missing cube'); continue
    try:
        vt,st=spectrum(fthick,sc,'h2co'); vn,sn=spectrum(fthin,sc,'c18o')
        if vt is None or vn is None: p('  off-field'); continue
        # systemic from optically-thin C18O within +-150 km/s of LSR-ish window
        Vthin,dVn,snrN,_=moments(vn,sn,vwin=(-150,150))
        # window thick line around systemic to avoid other lines in the spw
        vw=(Vthin-60,Vthin+60) if np.isfinite(Vthin) else (-150,150)
        Vthick,dVt,snrT,_=moments(vt,st,vwin=vw)
        dv=(Vthick-Vthin)/dVn if np.isfinite(dVn) and dVn>0 else np.nan
        infall = np.isfinite(dv) and dv<-0.25 and snrT>5 and snrN>5
        verdict=('INFALL (blue-asym)' if infall else
                 ('expansion/red' if (np.isfinite(dv) and dv>0.25 and snrT>5) else
                  ('no asym / weak' if snrN>5 else 'thin line undetected')))
        rows.append((nm,field,Vthick,Vthin,dVn,dv,snrT,snrN,verdict))
        p(f'  V_thick={Vthick:.1f} V_thin={Vthin:.1f} dV_thin={dVn:.1f} km/s  '
          f'dv={dv:+.2f}  SNR(thick/thin)={snrT:.0f}/{snrN:.0f}  => {verdict}')
    except Exception as e:
        p('  ERR',repr(e)[:90])

if rows:
    T=Table(rows=rows,names=['Name','field','V_thick','V_thin','dV_thin','dv_asym','SNR_thick','SNR_thin','verdict'])
    T.write('phase6_infall.fits',overwrite=True)
    p('\nWrote phase6_infall.fits')
LOG.close()