#!/usr/bin/env python
"""Common-beam spectral index for the ALMA-IMF HII regions: cut out B3+B6 around
each source, convolve both to a common beam, then measure peak alpha. Removes the
peak-alpha beam bias (publication-grade)."""
import os, glob, numpy as np
from astropy.table import Table
from astropy.io import fits
from astropy.wcs import WCS
from astropy.coordinates import SkyCoord
from astropy.nddata import Cutout2D
import astropy.units as u
from astropy.convolution import convolve_fft, Gaussian2DKernel
try:
    from radio_beam import Beam
    HAVE_RB=True
except Exception:
    HAVE_RB=False
LOG=open('convolve_alpha.log','w')
def p(*a):
    s=' '.join(str(x) for x in a);print(s);print(s,file=LOG);LOG.flush()
p('radio_beam available:',HAVE_RB)

IMGROOT='/orange/adamginsburg/ALMAIMF_Images'
NU={'B3':100.0,'B6':230.0}
m=Table.read('master_phase1_classified.fits')
hii=np.array([str(x) for x in m['simple_type']])=='HII region'
m=m[hii]
mc=SkyCoord(np.array(m['GLON'],float)*u.deg,np.array(m['GLAT'],float)*u.deg,frame='galactic')

def find_cont(field,band):
    base=os.path.join(IMGROOT,field,band)
    for sub in ['cleanest','bsens']:
        d=os.path.join(base,sub)
        if os.path.isdir(d):
            c=glob.glob(os.path.join(d,'*continuum_merged*image.tt0*.fits'))
            def score(f):
                b=os.path.basename(f).lower()
                if any(x in b for x in ['mask','residual','.model','.psf','.pb.','.alpha','.weight','dirty','preselfcal','finalmodel']):return -100
                s=0
                if 'finaliter' in b:s+=5
                if 'pbcor' in b:s+=2
                return s
            c=[f for f in c if score(f)>0]
            if c: return sorted(c,key=score)[-1]
    return None

# the 10 measured sources (field, name)
TARGS=[('W51-IRS2','G049.4885-00.3799'),('W51-IRS2','G049.4903-00.3694'),
       ('W51-E','G049.4885-00.3799'),('G333.60','G333.6032-00.2184'),
       ('G010.62','G010.6235-00.3834'),('G010.62','G010.6260-00.3744'),
       ('G010.62','G010.6311-00.3864'),('G327.29','G327.2852-00.5735'),
       ('G328.25','G328.2523-00.5320C'),('G338.93','G338.9237+00.5618')]
name2i={str(m['Name'][i]):i for i in range(len(m))}

def beam_from_hdr(h):
    return Beam(h['BMAJ']*u.deg,h['BMIN']*u.deg,h['BPA']*u.deg)

cache={}
def load(field,band):
    k=(field,band)
    if k in cache: return cache[k]
    f=find_cont(field,band)
    if not f: cache[k]=None; return None
    h=fits.open(f)[0]; d=np.squeeze(h.data); w=WCS(h.header).celestial
    pix=abs(h.header['CDELT2'])*3600  # arcsec/pix
    cache[k]=(d,w,h.header,pix,beam_from_hdr(h.header)); return cache[k]

results=[]
for field,nm in TARGS:
    if nm not in name2i: p(f'{nm} not in table'); continue
    i=name2i[nm]; sc=mc[i].icrs
    L3=load(field,'B3'); L6=load(field,'B6')
    if not (L3 and L6): p(f'{field}/{nm}: missing band'); continue
    d3,w3,h3,px3,b3=L3; d6,w6,h6,px6,b6=L6
    try:
        cb=b3.commonbeam_with(b6) if hasattr(b3,'commonbeam_with') else Beam(max(b3.major,b6.major),max(b3.minor,b6.minor),0*u.deg)
    except Exception:
        cb=Beam(max(b3.major,b6.major),max(b3.minor,b6.minor),0*u.deg)
    out={}
    for tag,(d,w,hh,px,bm) in [('B3',L3),('B6',L6)]:
        try:
            cut=Cutout2D(d,sc,30*u.arcsec,wcs=w)
            arr=cut.data.astype(float)
            # kernel to go from bm -> cb
            try:
                kbeam=cb.deconvolve(bm)
                kern=kbeam.as_kernel(px*u.arcsec)
                conv=convolve_fft(np.nan_to_num(arr),kern,preserve_nan=False,allow_huge=True)
                # rescale Jy/beam: multiply by cb_area/bm_area
                conv*= (cb.sr/bm.sr).value
            except Exception as e:
                conv=arr;
            cy,cx=np.array(conv.shape)//2
            sub=conv[cy-3:cy+4,cx-3:cx+4]
            out[tag]=np.nanmax(sub)
        except Exception as e:
            out[tag]=np.nan; p(f'  {nm} {tag} ERR {repr(e)[:50]}')
    s3,s6=out.get('B3',np.nan),out.get('B6',np.nan)
    if np.isfinite(s3) and np.isfinite(s6) and s3>0 and s6>0:
        a=np.log(s6/s3)/np.log(NU['B6']/NU['B3'])
        kind='DUST' if a>2.5 else ('free-free' if a<1.5 else 'mixed')
        results.append((nm,field,float(s3),float(s6),float(a),kind))
        p(f'{nm:22s} cb={cb.major.to(u.arcsec).value:.2f}\"  Scb_B3={s3*1e3:.2f} Scb_B6={s6*1e3:.2f} mJy/bm  alpha_cb={a:+.2f}  {kind}')
    else:
        p(f'{nm:22s} no measure s3={s3} s6={s6}')

if results:
    T=Table(rows=results,names=['Name','field','S_B3_cb','S_B6_cb','alpha_cb','kind'])
    T.write('almaimf_alpha_commonbeam.fits',overwrite=True)
    al=np.array([r[4] for r in results])
    p('\n=== COMMON-BEAM SUMMARY ===')
    p(f'  N={len(al)}  alpha range {al.min():+.2f}..{al.max():+.2f}')
    p(f'  free-free (<1.5): {int((al<1.5).sum())}  mixed: {int(((al>=1.5)&(al<=2.5)).sum())}  dust(>2.5): {int((al>2.5).sum())}')
LOG.close()