#!/usr/bin/env python
"""Final Phase 1 classification: crossmatch RMS young sources vs CORNISH N+S,
compute physical diameter, assign UC/compact/classical. Produce UC HII list."""
import numpy as np
from astropy.table import Table
from astropy.coordinates import SkyCoord
import astropy.units as u

LOG=open('final_classify.log','w')
def p(*a):
    s=' '.join(str(x) for x in a);print(s);print(s,file=LOG);LOG.flush()

def sky(t,ra='RAJ2000',de='DEJ2000'):
    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=(u.hourangle,u.deg))
    return SkyCoord(np.array(r,float)*u.deg,np.array(d,float)*u.deg)

PCAS = 0.0048481  # pc per (arcsec*kpc)
m=Table.read('master_phase1.fits')
mc=sky(m)
m['GLON']=mc.galactic.l.deg; m['GLAT']=mc.galactic.b.deg
styp=np.array([str(x) for x in m['simple_type']])
dist=np.array(m['Dist'],float)

# --- CORNISH North ---
cn=Table.read('cornish_north_raw.fits')
ccn=sky(cn)
iN,dN,_=mc.match_to_catalog_sky(ccn)
mN=dN<5*u.arcsec
AScD=np.ma.filled(cn['AScD'].astype(float),np.nan)[iN]
gMaj=np.ma.filled(cn['gMaj'].astype(float),np.nan)[iN]
angN=np.where(np.isfinite(AScD),AScD,gMaj)
SN=np.ma.filled(cn['Sint'].astype(float),np.nan)[iN]

# --- CORNISH South ---
cs=Table.read('cornish_south_catalog.fits')
p('CORNISH-S Size unit',cs['Size'].unit,'maj unit',cs['maj'].unit,'S unit',cs['S'].unit)
from collections import Counter
p('CORNISH-S SourceType:',dict(Counter([str(x) for x in cs['SourceType']]).most_common(12)))
ccs=sky(cs)
iS,dS,_=mc.match_to_catalog_sky(ccs)
mS=dS<5*u.arcsec
Size=np.ma.filled(cs['Size'].astype(float),np.nan)[iS]
sMaj=np.ma.filled(cs['maj'].astype(float),np.nan)[iS]
angS=np.where(np.isfinite(Size)&(Size>0),Size,sMaj)
SS=np.ma.filled(cs['S'].astype(float),np.nan)[iS]
StypeS=np.array([str(x) for x in cs['SourceType']])[iS]

# combine: prefer whichever survey matched (N and S footprints disjoint)
ang=np.where(mN,angN,np.where(mS,angS,np.nan))
flux=np.where(mN,SN,np.where(mS,SS,np.nan))
matched=mN|mS
survey=np.where(mN,'N',np.where(mS,'S','-'))
sep=np.where(mN,dN.arcsec,np.where(mS,dS.arcsec,np.nan))
Dpc=ang*dist*PCAS

m['cornish_survey']=survey
m['cornish_sep_as']=sep
m['cornish_ang_as']=ang
m['cornish_flux_mJy']=flux
m['cornish_SourceType_S']=np.where(mS,StypeS,'')
m['Dpc']=Dpc
m['cornish_match']=matched

cls=np.array(['none']*len(m),dtype=object)
sel=matched&np.isfinite(Dpc)
cls[sel&(Dpc<=0.1)]='UCHII'
cls[sel&(Dpc>0.1)&(Dpc<=0.5)]='compact'
cls[sel&(Dpc>0.5)]='classical'
m['radio_class']=[str(x) for x in cls]
hii=styp=='HII region'

p('\n=== COMBINED CORNISH N+S classification ===')
p('Total master young sources:',len(m))
p('RMS HII regions:',int(hii.sum()),' YSO:',int((styp=="YSO").sum()),' HII/YSO:',int((styp=="HII/YSO").sum()))
p('Matched to CORNISH (N or S):',int(matched.sum()),' (N',int(mN.sum()),'S',int(mS.sum()),')')
for k in ['UCHII','compact','classical']:
    n=int((cls==k).sum());nh=int(((cls==k)&hii).sum())
    p(f'  {k:10s}: {n:4d} total, {nh:4d} RMS-HII')

uc=m[cls=='UCHII']
uc.sort('Dpc')
uc.write('uchii_candidates.fits',overwrite=True)
m.write('master_phase1_classified.fits',overwrite=True)
p('\nUC HII candidates ->',len(uc),'written uchii_candidates.fits')
p('UC HII fraction of RMS-HII-matched:',int(((cls=="UCHII")&hii).sum()),'/',int((matched&hii).sum()))
p('UC HII fraction of full embedded sample:',int((cls=="UCHII").sum()),'/',len(m))
LOG.close()