#!/usr/bin/env python
"""Off-plane (|b|>1.5) SF RMS sources with ALMA: comparable WISE+VLASS/RACS figures
via viz_offplane.py. One panel-set per unique source (60" dedup). Resumable."""
import os, sys, subprocess
import numpy as np
from astropy.table import Table
from astropy.coordinates import SkyCoord, search_around_sky
import astropy.units as u
from collections import defaultdict
from viz_meerkat_alma import gather_alma_bands

N = int(sys.argv[1]) if len(sys.argv) > 1 else 10**9
OUT = 'offplane_figs'
os.makedirs(OUT, exist_ok=True)
DEDUP = 60.0

t = Table.read('rms_table1_raw.fits')
cc = SkyCoord(t['RAJ2000'], t['DEJ2000'], unit=(u.hourangle, u.deg))
l, b = cc.galactic.l.deg, cc.galactic.b.deg
L = np.array(t['Lbol'], float)
typ = np.array([str(x) for x in t['Type']])
sf = np.isin(typ, ['HII region', 'YSO', 'HII/YSO'])
off = np.where((np.abs(b) > 1.5) & sf)[0]

i1, i2, _, _ = search_around_sky(cc[off], cc[off], DEDUP * u.arcsec)
parent = list(range(len(off)))
def find(x):
    while parent[x] != x:
        parent[x] = parent[parent[x]]; x = parent[x]
    return x
for a_, b_ in zip(i1.tolist(), i2.tolist()):
    ra, rb = find(a_), find(b_)
    if ra != rb:
        parent[ra] = rb
groups = defaultdict(list)
for k in range(len(off)):
    groups[find(k)].append(off[k])
reps = [max(g, key=lambda k: (np.nan_to_num(L[k], nan=-1.0),
                              typ[k] == 'HII region')) for g in groups.values()]
idx = sorted(reps, key=lambda k: l[k])
print(f"{len(off)} off-plane SF entries -> {len(idx)} unique sources", flush=True)

made = attempted = 0
for i in idx:
    if attempted >= N:
        break
    name = str(t['Name'][i])
    out = os.path.join(OUT, f"{name}.png")
    if os.path.exists(out):
        made += 1
        continue
    if not gather_alma_bands(cc[i]):       # only sources with on-disk ALMA
        continue
    attempted += 1
    print(f"=== {name}  {typ[i]}  l={l[i]:.3f} b={b[i]:+.3f}", flush=True)
    try:
        r = subprocess.run(['python', 'viz_offplane.py', out, '--center', f"{l[i]},{b[i]}",
                            '--field-arcmin', '4'], timeout=600)
    except subprocess.TimeoutExpired:
        print(f"  TIMEOUT {name}", flush=True); continue
    if r.returncode == 0 and os.path.exists(out):
        made += 1
print(f"DONE made {made} off-plane figures in {OUT}/", flush=True)