#!/usr/bin/env python
"""Phase 6 quantitative: measured large-scale accretion-flow rates (Wells+2024)
by evolutionary class. Are HII regions still accreting at clump scale, and at what
rate vs younger sources?"""
import numpy as np
from astropy.table import Table, join
LOG=open('phase6_flowrates.log','w')
def p(*a):
    s=' '.join(str(x) for x in a);print(s);print(s,file=LOG);LOG.flush()

t1=Table.read('wells2024_tabled1.fits')  # ALMAGAL, ID, Class, Mass, Lum
t2=Table.read('wells2024_tabled2.fits')  # ID, FlowRate Lin/Lout/Rin/Rout, Masscore
p('t1 cols',t1.colnames)
p('t2 cols',t2.colnames)
for c in ['FlowRateLin','FlowRateRin']:
    p(f't2 {c}: unit={t2[c].unit} desc={t2[c].description}')

# inflow rate per flow = max of left/right inflow (Msun/yr)
lin=np.ma.filled(t2['FlowRateLin'].astype(float),np.nan)
rin=np.ma.filled(t2['FlowRateRin'].astype(float),np.nan)
infl=np.nanmax(np.vstack([lin,rin]),axis=0)
t2['Mdot_in']=infl
# aggregate per source ID: total inflow rate (sum of flows)
ids=np.array([str(x) for x in t2['ID']])
agg={}
for i,idv in enumerate(ids):
    if np.isfinite(infl[i]):
        agg[idv]=agg.get(idv,0.0)+infl[i]
t1id=np.array([str(x) for x in t1['ID']])
cls=np.array([str(x) for x in t1['Class']])
mdot=np.array([agg.get(i,np.nan) for i in t1id])

p('\n=== Total large-scale inflow rate Mdot (Msun/yr) by evolutionary class ===')
order=['Quiescent','Protostellar','YSO','HII Region']
for c in order:
    m=(cls==c)&np.isfinite(mdot)
    if m.sum():
        v=mdot[m]
        p(f'  {c:13s} N={m.sum():2d}  median={np.median(v):.2e}  '
          f'q25-75={np.percentile(v,25):.1e}-{np.percentile(v,75):.1e}  max={v.max():.1e}')
# fraction with ANY detected inflow per class (in t1, whether they appear in agg)
p('\n=== Fraction of each class with a detected accretion flow ===')
for c in order:
    incls=cls==c; ndet=sum((str(t1['ID'][i]) in agg) for i in np.where(incls)[0])
    p(f'  {c:13s} {ndet}/{incls.sum()} ({100*ndet/max(incls.sum(),1):.0f}%)')

# our matched HII regions (from phase6b): the 12 RMS HII in Wells
MATCHED=['AG025.8012-0.1563','AG024.5073-0.2196','AG029.9558-0.0161','AG030.2960+0.0554',
         'AG305.1957+0.0330','AG308.6467+0.6472','AG309.9212+0.4780','AG311.6268+0.2907',
         'AG316.8113-0.0575','AG320.6753+0.2459','AG327.4026+0.4445','AG023.2657+0.0767']
alm=np.array([str(x) for x in t1['ALMAGAL']])
p('\n=== Our RMS HII regions in Wells: measured inflow ===')
got=[]
for nm in MATCHED:
    w=np.where(alm==nm)[0]
    if len(w):
        i=w[0]; md=agg.get(str(t1['ID'][i]),np.nan)
        got.append(md)
        p(f"  {nm:22s} Class={str(t1['Class'][i]):12s} Mdot={md if np.isfinite(md) else 'none':}")
got=[x for x in got if np.isfinite(x)]
if got:
    p(f'  -> {len(got)} with measured inflow; median Mdot={np.median(got):.2e} Msun/yr')
T=t1.copy(); T['Mdot_inflow']=mdot
T.write('wells_flowrates_byclass.fits',overwrite=True)
p('Wrote wells_flowrates_byclass.fits')
LOG.close()