#!/usr/bin/env python
"""Make MeerKAT+GLOSTAR+Spitzer+ALMA figures for ALL ALMA-observed RMS targets.
Driven directly from alma_census.csv (has_cont=1) -> guarantees every observed
target is attempted, with no lossy filename-longitude gate. Centered on the RMS
position; one panel per unique source (60" dedup of A/B/D components)."""
import os, sys, csv, subprocess
import numpy as np
from astropy.coordinates import SkyCoord, search_around_sky
import astropy.units as u
from collections import defaultdict

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   # arcsec

rows = [r for r in csv.DictReader(open('alma_census.csv')) if r['has_cont'] == '1']
gl = np.array([float(r['GLON']) for r in rows])
gb = np.array([float(r['GLAT']) for r in rows])
cc = SkyCoord(gl, gb, unit='deg', frame='galactic')

# vectorized 60" dedup (union-find) over the ALMA-observed targets
i1, i2, _, _ = search_around_sky(cc, cc, DEDUP * u.arcsec)
parent = list(range(len(rows)))
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(rows)):
    groups[find(k)].append(k)

def nobs(k):
    v = rows[k]['n_obs']
    return int(v) if v not in ('', '0') else 0
# rep per cluster: prefer a real radio class (HII-like) over 'none', then most obs
reps = [max(g, key=lambda k: (rows[k]['radio_class'] != 'none', nobs(k))) for g in groups.values()]
idx = sorted(reps, key=lambda k: gl[k])
print(f"{len(rows)} ALMA-observed targets -> {len(idx)} unique sources after {DEDUP:.0f}\" dedup",
      flush=True)

made = attempted = 0
for n, i in enumerate(idx):
    if attempted >= N:
        break
    if n < SKIP:
        continue
    name = rows[i]['Name'].rstrip('ABD')  # strip component suffix for the figure name
    out = os.path.join(OUT, f"{name}.png")
    if os.path.exists(out):                 # resume
        made += 1
        continue
    attempted += 1
    print(f"=== {name}  {rows[i]['radio_class']}  l={gl[i]:.3f} b={gb[i]:+.3f}", flush=True)
    try:
        r = subprocess.run(['python', 'viz_meerkat_alma.py', 'auto', out, '--field-arcmin', '3',
                            '--center', f"{gl[i]},{gb[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}/ ({len(idx)} unique targets total)", flush=True)