#!/usr/bin/env python
"""All-sky IRAS PSC >1e4 Lsun@<=3kpc gate (3kpc-complete: union flux floor so cold
SEDs aren't lost), split off-plane vs on-plane, cross-match on-plane vs RMS.

Gate (distance-independent): F_bol >= F_bol(1e4 Lsun,3kpc)=3.55e-11 W/m^2.
Color (Wood&Churchwell89 embedded): log(F25/F12)>=0.57 & log(F60/F25)>=0.30.
"""
import numpy as np
from astroquery.vizier import Vizier
from astropy.table import vstack, unique
from astropy.coordinates import SkyCoord
import astropy.units as u

Lsun = 3.828e26; kpc = 3.0856776e19
F_THRESH = 1e4 * Lsun / (4 * np.pi * (3 * kpc) ** 2)
nu = 2.99792458e8 / (np.array([12., 25., 60., 100.]) * 1e-6)

def fbol(f12, f25, f60, f100):
    Fnu = np.vstack([f12, f25, f60, f100]).T * 1e-26
    return np.trapezoid(Fnu[:, ::-1], nu[::-1], axis=1)

Vizier.ROW_LIMIT = -1
v = Vizier(columns=['IRAS', 'RA1950', 'DE1950', 'Fnu_12', 'Fnu_25', 'Fnu_60', 'Fnu_100',
                    'q_Fnu_12', 'q_Fnu_25', 'q_Fnu_60', 'q_Fnu_100'], row_limit=-1)
# UNION floor: warm sources via F60, cold via F100 -> 3kpc-complete
parts = []
for con in [{'Fnu_60': '>20'}, {'Fnu_100': '>25'}]:
    r = v.query_constraints(catalog='II/125/main', **con)
    parts.append(r[0]); print('query', con, len(r[0]))
t = unique(vstack(parts), keys='IRAS')
print('union rows:', len(t))

sc = SkyCoord(t['RA1950'], t['DE1950'], unit=(u.hourangle, u.deg),
              frame='fk4', equinox='B1950')
g = sc.galactic
t['GLON'] = g.l.deg; t['GLAT'] = g.b.deg
t['_ra'] = sc.icrs.ra.deg; t['_dec'] = sc.icrs.dec.deg

def col(c): return np.nan_to_num(np.array(t[c], float))
f12, f25, f60, f100 = col('Fnu_12'), col('Fnu_25'), col('Fnu_60'), col('Fnu_100')
q25 = np.nan_to_num(np.array(t['q_Fnu_25'], float))
q60 = np.nan_to_num(np.array(t['q_Fnu_60'], float))
Fb = fbol(f12, f25, f60, f100)
t['F_bol'] = Fb; t['L_at3kpc'] = Fb * 4 * np.pi * (3 * kpc) ** 2 / Lsun
with np.errstate(divide='ignore', invalid='ignore'):
    c1 = np.log10(np.where(f12 > 0, f25 / f12, np.nan))
    c2 = np.log10(np.where(f25 > 0, f60 / f25, np.nan))
red = (c1 >= 0.57) & (c2 >= 0.30)
gate = (Fb >= F_THRESH) & (q25 >= 2) & (q60 >= 2) & red
print(f'all-sky gate+red candidates: {gate.sum()}')

cand = t[gate]
offp = cand[np.abs(cand['GLAT']) > 5]
onp = cand[(np.abs(cand['GLAT']) <= 5) & (np.abs(cand['GLON']) > 10) &
           (np.abs(cand['GLON'] - 360) > 10)]
print(f'  off-plane |b|>5:            {len(offp)}')
print(f'  on-plane |b|<5 & |l|>10:    {len(onp)}')

# cross-match on-plane vs RMS
from astropy.table import Table
rms = Table.read('rms_table1_raw.fits')
rc = SkyCoord(rms['RAJ2000'], rms['DEJ2000'], unit=(u.hourangle, u.deg))
oc = SkyCoord(onp['_ra'], onp['_dec'], unit=u.deg)
idx, d2d, _ = oc.match_to_catalog_sky(rc)
matched = d2d.arcsec < 60
onp['rms_sep'] = d2d.arcsec
onp['rms_name'] = [str(rms['Name'][i]) for i in idx]
onp['in_rms'] = matched
print(f'  on-plane matched to RMS (<60"): {matched.sum()} / {len(onp)}')
miss = onp[~matched]
print(f'  on-plane gate-passers NOT in RMS: {len(miss)}')

offp.sort('F_bol'); offp = offp[::-1]
miss.sort('F_bol'); miss = miss[::-1]
offp.write('offplane_iras_v2_cand.fits', overwrite=True)
miss.write('onplane_iras_notRMS.fits', overwrite=True)
print('\nwrote offplane_iras_v2_cand.fits', len(offp), '| onplane_iras_notRMS.fits', len(miss))
print('\n-- off-plane (should match prior 13 +/- cold additions) --')
for r in offp:
    print(f"  IRAS{r['IRAS']:12s} l={r['GLON']:7.2f} b={r['GLAT']:+6.2f} L@3kpc={r['L_at3kpc']:.2e}")
print('\n-- on-plane gate-passers NOT in RMS, top 30 by F_bol --')
for r in miss[:30]:
    print(f"  IRAS{r['IRAS']:12s} l={r['GLON']:7.2f} b={r['GLAT']:+6.2f} "
          f"L@3kpc={r['L_at3kpc']:.2e} (nrstRMS {r['rms_sep']:.0f}\")")