#!/usr/bin/env python
"""Off-plane (|b|>5) IRAS PSC search for >1e4 Lsun YSO candidates, RMS-like.

Gate is distance-independent: F_bol >= F_bol(1e4 Lsun, 3 kpc) = 3.55e-11 W/m^2
=> source is either >1e4 Lsun at <=3 kpc, or more luminous & farther.
Color cut (Wood & Churchwell 1989, embedded/red): log(F25/F12)>=0.57 & log(F60/F25)>=0.30.
Output candidates for vetting (AGB/PNe/galaxy contamination expected off-plane).
"""
import numpy as np
from astroquery.vizier import Vizier
from astropy.table import Table, vstack
from astropy.coordinates import SkyCoord
import astropy.units as u

Lsun = 3.828e26          # W
kpc = 3.0856776e19       # m
F_THRESH = 1e4 * Lsun / (4 * np.pi * (3 * kpc) ** 2)   # 3.55e-11 W/m^2
print(f'F_bol threshold (1e4 Lsun @ 3kpc) = {F_THRESH:.3e} W/m^2')

# IRAS band central wavelengths -> frequencies
lam = np.array([12., 25., 60., 100.]) * 1e-6   # m
nu = 2.99792458e8 / lam                         # Hz (descending)

def fbol(f12, f25, f60, f100):
    """Trapezoidal integral of nu*Fnu over log-nu across the 4 IRAS bands.
    Fnu in Jy (1 Jy = 1e-26 W/m^2/Hz). Returns W/m^2 (12-100um only; conservative)."""
    Fnu = np.vstack([f12, f25, f60, f100]).T * 1e-26   # W/m^2/Hz, shape (N,4)
    # integrate Fnu dnu over the band span (trapezoid in nu, ascending order 100->12)
    nua = nu[::-1]                                       # ascending freq
    Fa = Fnu[:, ::-1]
    return np.trapezoid(Fa, nua, axis=1)

# PSC main has only B1950 coords (no galactic cols for constraints). Pull bright-60
# all-sky once, compute galactic b locally, filter |b|>5. Floor 30 Jy keeps < server cap
# and is below the ~50-100 Jy a 1e4@3kpc embedded source needs.
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)
r = v.query_constraints(catalog='II/125/main', Fnu_60='>30')
t = r[0]
print('all-sky Fnu_60>30:', len(t))
# B1950 (FK4) -> Galactic
sc = SkyCoord(t['RA1950'], t['DE1950'], unit=(u.hourangle, u.deg), frame='fk4', equinox='B1950')
gal = sc.galactic
t['GLON'] = gal.l.deg
t['GLAT'] = gal.b.deg
t = t[np.abs(t['GLAT']) > 5]
print('off-plane |b|>5, Fnu_60>30:', len(t))

# numeric arrays (masked -> 0 for missing/upper-limit unreliable)
def col(c):
    a = np.array(t[c], float)
    return np.nan_to_num(a, nan=0.0)
f12, f25, f60, f100 = col('Fnu_12'), col('Fnu_25'), col('Fnu_60'), col('Fnu_100')
# quality flags: 3=good,2=moderate,1=upper limit. Require >=2 on 25&60.
q25 = np.array(t['q_Fnu_25'], float); q60 = np.array(t['q_Fnu_60'], float)
goodq = (np.nan_to_num(q25) >= 2) & (np.nan_to_num(q60) >= 2)

Fb = fbol(f12, f25, f60, f100)
t['F_bol'] = Fb
t['L_at3kpc'] = Fb * 4 * np.pi * (3 * kpc) ** 2 / Lsun   # Lsun if it were at 3 kpc

# Wood-Churchwell red/embedded colors (avoid div0)
with np.errstate(divide='ignore', invalid='ignore'):
    c1 = np.log10(np.where(f12 > 0, f25 / f12, np.nan))    # >=0.57
    c2 = np.log10(np.where(f25 > 0, f60 / f25, np.nan))    # >=0.30
red = (c1 >= 0.57) & (c2 >= 0.30)

passgate = (Fb >= F_THRESH) & goodq
cand = passgate & red
print(f'\nF_bol>=thresh & goodq:        {passgate.sum()}')
print(f'  + Wood-Churchwell red SED:  {cand.sum()}')

out = t[cand]
out.sort('F_bol')
out = out[::-1]
out.write('offplane_iras_candidates.fits', overwrite=True)
print('wrote offplane_iras_candidates.fits  (', len(out), 'candidates )')
# quick look
print('\nTop 25 by F_bol (L_at3kpc = Lsun if placed at 3 kpc):')
for r in out[:25]:
    print(f"  IRAS{r['IRAS']:12s} l={r['GLON']:7.2f} b={r['GLAT']:+6.2f}  "
          f"F60={r['Fnu_60']:8.1f} F100={r['Fnu_100']:8.1f}  L@3kpc={r['L_at3kpc']:.2e}")