#!/usr/bin/env python
"""Characterize the radio (free-free) emission of each RMS HII region from CORNISH:
position, size (resolved vs point-like), physical diameter, electron density n_e,
emission measure EM -> classify hyper-compact (HC) / ultra-compact (UC) / compact,
and assess whether a star can be localized at all."""
import numpy as np
from astropy.table import Table
from astropy.coordinates import SkyCoord
import astropy.units as u
LOG=open('radio_characterize.log','w')
def p(*a):
    s=' '.join(str(x) for x in a);print(s);print(s,file=LOG);LOG.flush()
PCAS=0.0048481  # pc per (arcsec*kpc)
TE=1e4; NU=5.0  # CORNISH 5 GHz

def sky(t,ra,de,unit):
    r=t[ra];d=t[de]
    if r.dtype.kind in ('S','U'):
        return SkyCoord([x.decode() if isinstance(x,bytes) else str(x) for x in r],
                        [x.decode() if isinstance(x,bytes) else str(x) for x in d],unit=unit)
    return SkyCoord(np.array(r,float)*u.deg,np.array(d,float)*u.deg)

def physics(Sint_mJy,maj_as,min_as,d_kpc):
    """Tb, EM, n_e from integrated flux + Gaussian FWHM size + distance."""
    Tb=1.222e3*Sint_mJy/(NU**2*maj_as*min_as)        # K, mean brightness temp
    tau=-np.log(np.clip(1-Tb/TE,1e-6,None))           # allow optically thick
    EM=tau/(8.235e-2*TE**-1.35*NU**-2.1)              # pc cm^-6
    geo=np.sqrt(maj_as*min_as)*d_kpc*PCAS             # physical FWHM diam, pc
    L=np.clip(geo,1e-4,None)
    ne=np.sqrt(EM/L)                                  # cm^-3
    return Tb,EM,ne,geo

m=Table.read('master_phase1_classified.fits')
hii=np.array([str(x) for x in m['simple_type']])=='HII region'
m=m[hii].copy()
mc=sky(m,'RAJ2000','DEJ2000',(u.hourangle,u.deg))
dist=np.array(m['Dist'],float)

# ---- CORNISH North ----
cn=Table.read('cornish_north_raw.fits'); ccn=sky(cn,'RAJ2000','DEJ2000',(u.deg,u.deg))
iN,dN,_=mc.match_to_catalog_sky(ccn); mN=dN<5*u.arcsec
# ---- CORNISH South ---- (RAJ2000 is sexagesimal HOURS; use decimal GLON/GLAT instead)
cs=Table.read('cornish_south_catalog.fits')
ccs=SkyCoord(np.array(cs['GLON'],float)*u.deg,np.array(cs['GLAT'],float)*u.deg,frame='galactic').icrs
iS,dS,_=mc.match_to_catalog_sky(ccs); mS=dS<5*u.arcsec

N=len(m)
cra=np.full(N,np.nan); cdec=np.full(N,np.nan); flux=np.full(N,np.nan)
maj=np.full(N,np.nan); minr=np.full(N,np.nan); decon=np.full(N,np.nan)
resolved=np.zeros(N,bool); srctype=np.array(['']*N,dtype=object); beam=np.full(N,np.nan)

# North: gMaj/gMin = fitted Gaussian FWHM; AScD = deconvolved scale (only >1.8")
gMajN=np.ma.filled(cn['gMaj'].astype(float),np.nan)[iN]
gMinN=np.ma.filled(cn['gMin'].astype(float),np.nan)[iN]
AScDN=np.ma.filled(cn['AScD'].astype(float),np.nan)[iN]
SintN=np.ma.filled(cn['Sint'].astype(float),np.nan)[iN]
craN=ccn.ra.deg[iN]; cdecN=ccn.dec.deg[iN]
# South: maj/min (FWHM), Size = deconvolved, S=flux, SourceType
sMaj=np.ma.filled(cs['maj'].astype(float),np.nan)[iS]
sMin=np.ma.filled(cs['min'].astype(float),np.nan)[iS]
SizeS=np.ma.filled(cs['Size'].astype(float),np.nan)[iS]
SS=np.ma.filled(cs['S'].astype(float),np.nan)[iS]
StypeS=np.array([str(x) for x in cs['SourceType']])[iS]
craS=ccs.ra.deg[iS]; cdecS=ccs.dec.deg[iS]

for i in range(N):
    if mN[i]:
        cra[i],cdec[i]=craN[i],cdecN[i]; flux[i]=SintN[i]
        maj[i],minr[i]=gMajN[i],gMinN[i]; decon[i]=AScDN[i]; beam[i]=1.5
        resolved[i]=np.isfinite(AScDN[i]) and AScDN[i]>1.8
        srctype[i]='cornishN'
    elif mS[i]:
        cra[i],cdec[i]=craS[i],cdecS[i]; flux[i]=SS[i]
        maj[i],minr[i]=sMaj[i],sMin[i]; decon[i]=SizeS[i]; beam[i]=2.5
        resolved[i]=np.isfinite(SizeS[i]) and SizeS[i]>0
        srctype[i]=StypeS[i]

matched=mN|mS
# physical size: deconvolved if resolved, else beam upper limit
ang=np.where(resolved&np.isfinite(decon),decon,beam)
Dpc=ang*dist*PCAS
Tb,EM,ne,geoFWHM=physics(flux,np.where(np.isfinite(maj),maj,beam),
                          np.where(np.isfinite(minr),minr,beam),dist)

# classification by physical diameter
cls=np.array(['none']*N,dtype=object)
sel=matched&np.isfinite(Dpc)
cls[sel&(Dpc<=0.05)]='HC HII'      # hyper-compact
cls[sel&(Dpc>0.05)&(Dpc<=0.1)]='UC HII'
cls[sel&(Dpc>0.1)&(Dpc<=0.5)]='compact'
cls[sel&(Dpc>0.5)]='classical'

m['cornish_RA']=cra; m['cornish_DE']=cdec; m['cornish_flux_mJy']=flux
m['radio_FWHM_as']=np.sqrt(maj*minr); m['radio_decon_as']=decon
m['radio_resolved']=resolved; m['Dpc_radio']=Dpc
m['ne_cm3']=ne; m['EM_pccm6']=EM; m['Tb_K']=Tb
m['hii_class']=[str(x) for x in cls]; m['matched_cornish']=matched
m.write('radio_characterized.fits',overwrite=True)

p('=== Radio characterization of RMS HII regions (CORNISH N+S) ===')
p(f'matched to CORNISH: {int(matched.sum())}/{N}')
p(f'  resolved (deconvolved size measured): {int((matched&resolved).sum())}')
p(f'  unresolved / point-like (<=beam):     {int((matched&~resolved).sum())}')
from collections import Counter
p('\n-- physical-size class --')
for k in ['HC HII','UC HII','compact','classical']:
    s=cls==k
    if s.sum():
        p(f'  {k:9s}: {int(s.sum()):3d}  median D={np.nanmedian(Dpc[s]):.3f}pc '
          f'ne={np.nanmedian(ne[s]):.1e} EM={np.nanmedian(EM[s]):.1e} '
          f'frac_resolved={np.mean(resolved[s]):.2f}')
p('\n-- star localizability --')
pt=matched&~resolved
p(f'  point-like (star *might* be near radio peak): {int(pt.sum())} ({100*pt.sum()/matched.sum():.0f}% of matched)')
p(f'  resolved/extended (star NOT localizable from radio centroid): {int((matched&resolved).sum())} '
  f'({100*(matched&resolved).sum()/matched.sum():.0f}%)')
p('Wrote radio_characterized.fits')
LOG.close()