#!/usr/bin/env python
"""Crossmatch RMS HII regions with CORNISH-North, compute physical diameter,
classify UC / compact / classical. Robust to sexagesimal vs decimal coords."""
import numpy as np
from astropy.table import Table
from astropy.coordinates import SkyCoord
import astropy.units as u

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

def to_skycoord(t, racol='RAJ2000', decol='DEJ2000'):
    ra = t[racol]; de = t[decol]
    if ra.dtype.kind in ('S','U'):
        # sexagesimal strings; RA in hours, Dec in deg
        return SkyCoord([x.decode() if isinstance(x,bytes) else str(x) for x in ra],
                        [x.decode() if isinstance(x,bytes) else str(x) for x in de],
                        unit=(u.hourangle, u.deg))
    else:
        return SkyCoord(np.array(ra,float)*u.deg, np.array(de,float)*u.deg)

m = Table.read('master_phase1.fits')
cn = Table.read('cornish_north_raw.fits')
p('CORNISH RA dtype', cn['RAJ2000'].dtype, 'sample', repr(cn['RAJ2000'][0]))

mc = to_skycoord(m)
cc = to_skycoord(cn)
m['GLON'] = mc.galactic.l.deg; m['GLAT'] = mc.galactic.b.deg
p('Master GLON range', round(float(m['GLON'].min()),1), round(float(m['GLON'].max()),1))

in_foot = (m['GLON']>9.5)&(m['GLON']<65.5)&(np.abs(m['GLAT'])<1.05)
styp = np.array([str(x) for x in m['simple_type']])
hii = styp=='HII region'
p('\nMaster HII regions total:', hii.sum())
p('Master HII regions in CORNISH-N footprint:', int((hii&in_foot).sum()))
p('All young in footprint:', int(in_foot.sum()))

idx, d2d, _ = mc.match_to_catalog_sky(cc)
matched = d2d < 5*u.arcsec
p('Master sources w/ CORNISH match <5":', int(matched.sum()))
p('  HII type matched:', int((matched&hii).sum()))
p('  HII & in footprint matched:', int((matched&hii&in_foot).sum()))

AScD = np.ma.filled(cn['AScD'].astype(float), np.nan)[idx]   # arcsec deconvolved
gMaj = np.ma.filled(cn['gMaj'].astype(float), np.nan)[idx]
dist = np.array(m['Dist'], float)                            # kpc
# use AScD where available, else gMaj (convolved) as upper limit proxy
ang = np.where(np.isfinite(AScD), AScD, gMaj)
Dpc = ang * dist * 0.0048481
m['cornish_sep_as'] = d2d.arcsec
m['cornish_AScD_as'] = AScD
m['cornish_gMaj_as'] = gMaj
m['cornish_ang_as'] = ang
m['cornish_Sint_mJy'] = np.ma.filled(cn['Sint'].astype(float), np.nan)[idx]
m['Dpc'] = Dpc
m['cornish_match'] = matched
m['in_cornishN'] = in_foot

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]

p('\n--- Radio size classification (CORNISH-N matched) ---')
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')
p('Matched, no size:', int((matched&~np.isfinite(Dpc)).sum()))
# of HII regions in footprint, how many matched
p('\nHII-in-footprint matched fraction:',
  int((matched&hii&in_foot).sum()), '/', int((hii&in_foot).sum()))

m.write('master_phase1_classified.fits', overwrite=True)
p('Wrote master_phase1_classified.fits')
LOG.close()