#!/usr/bin/env python
"""Make MeerKAT+GLOSTAR+Spitzer+ALMA figures for ALL RMS sources (HII + YSO +
HII/YSO) within SMGPS coverage that have ALMA continuum ON DISK. Coverage is
decided by the footprint index (gather_alma_bands), which sees both
continuum_downloads/ AND the uvdata tree -- so no source with on-disk ALMA is
missed regardless of type, filename, or storage location. One panel 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
SKIP = int(sys.argv[2]) if len(sys.argv) > 2 else 0
OUT = 'meerkat_alma_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)
smgps = (((2 <= l) & (l <= 60)) | ((251 <= l) & (l <= 358))) & (np.abs(b) < 1.5)
# star-formation tracers only (the accretion-timescale population); evolved
# star / PN / OH-IR types are excluded -- they only fall coincidentally near ALMA
typ = np.array([str(x) for x in t['Type']])
sf = np.isin(typ, ['HII region', 'YSO', 'HII/YSO'])
sm = np.where(smgps & sf)[0]

# 60" union-find dedup
i1, i2, _, _ = search_around_sky(cc[sm], cc[sm], DEDUP * u.arcsec)
parent = list(range(len(sm)))
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(sm)):
    groups[find(k)].append(sm[k])
reps = [max(g, key=lambda k: (np.nan_to_num(L[k], nan=-1.0),
                              str(t['Type'][k]) == 'HII region')) for g in groups.values()]
idx = sorted(reps, key=lambda k: l[k])
print(f"{len(sm)} SMGPS entries -> {len(idx)} unique sources (all types)", flush=True)

made = attempted = 0
for n, i in enumerate(idx):
    if attempted >= N:
        break
    if n < SKIP:
        continue
    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]):     # index pre-filter: skip if no on-disk ALMA
        continue
    attempted += 1
    print(f"=== {name}  {t['Type'][i]}  l={l[i]:.3f} b={b[i]:+.3f}", flush=True)
    try:
        r = subprocess.run(['python', 'viz_meerkat_alma.py', 'auto', out, '--field-arcmin', '3',
                            '--center', f"{l[i]},{b[i]}"], timeout=420)
    except subprocess.TimeoutExpired:
        print(f"  TIMEOUT {name} (>420s), skipped", flush=True)
        continue
    if r.returncode == 0 and os.path.exists(out):
        made += 1
print(f"DONE made {made} figures in {OUT}/", flush=True)