#!/usr/bin/env python
"""Make MeerKAT+GLOSTAR+Spitzer+ALMA figures for RMS sources that have ALMA data,
within SMGPS coverage. Centered on the RMS position; titled with RMS Name/Type."""
import os, sys, subprocess
import numpy as np
from astropy.table import Table
from astropy.coordinates import SkyCoord
import astropy.units as u
from viz_meerkat_alma import gather_alma_bands

N = int(sys.argv[1]) if len(sys.argv) > 1 else 30
SKIP = int(sys.argv[2]) if len(sys.argv) > 2 else 0     # skip first SKIP matches
OUT = 'meerkat_alma_figs'
os.makedirs(OUT, exist_ok=True)

DEDUP = 60.0   # arcsec: entries within this are the same source (A/B comps etc.)

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)
from astropy.coordinates import search_around_sky
from collections import defaultdict
smgps = (((2 <= l) & (l <= 60)) | ((251 <= l) & (l <= 358))) & (np.abs(b) < 1.5)
sm = np.where(smgps)[0]

# vectorized dedup: union-find over all pairs within DEDUP (C-level search_around_sky)
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])
# keep brightest (HII-preferred) representative per cluster
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 after {DEDUP:.0f}\" dedup",
      flush=True)

import glob


def has_candidate(ld):
    """Cheap filename-only check (no header reads) for any ALMA file near longitude ld."""
    for tk in {f'G{int(round(ld))}', f'G0{int(round(ld))}', f'G{ld:.1f}'}:
        if glob.glob(f'continuum_downloads/*{tk}*pbcor.fits'):
            return True
    return False


made = 0
for i in idx:
    if made >= N:
        break
    name = str(t['Name'][i])
    out = os.path.join(OUT, f"{name}.png")
    if os.path.exists(out):                 # resume
        made += 1
        continue
    if not has_candidate(l[i]):             # cheap gate; viz does real coverage check
        continue
    print(f"=== {name}  {t['Type'][i]}  l={l[i]:.3f} b={b[i]:+.3f}", flush=True)
    try:                                    # all FITS I/O inside the timeout-guarded subprocess
        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)