#!/usr/bin/env python
"""SED figure per nearby-sample source from the CDS VizieR photometry service.
nuFnu vs wavelength, all catalog photometry within RADIUS; annotate adopted d and Lbol.
Output: sed_figs/<Name>.png + sed_figs/sed_status.csv"""
import os
import csv
import time
import numpy as np
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
from astropy.table import Table
from astropy.coordinates import SkyCoord
import astropy.units as u
import requests
from io import BytesIO
from extinction_profiles import load_targets

RADIUS_AS = 12.0
OUT = 'sed_figs'
os.makedirs(OUT, exist_ok=True)


def fetch_sed(ra, dec, rs=RADIUS_AS):
    url = f'https://vizier.cfa.harvard.edu/viz-bin/sed?-c={ra},{dec}&-c.rs={rs}'
    r = requests.get(url, timeout=60)
    r.raise_for_status()
    return Table.read(BytesIO(r.content), format='votable')


def main(only=None):
    targs = load_targets()
    # restrict: <2kpc members + all curated
    sel = [t for t in targs if t['sample'] == 'curated' or t['d_best'] < 2.0]
    status = []
    for t in sel:
        if only and t['name'] not in only:
            continue
        fn = os.path.join(OUT, f"{t['name']}.png")
        c = SkyCoord(t['glon'] * u.deg, t['glat'] * u.deg, frame='galactic').icrs
        try:
            tab = fetch_sed(c.ra.deg, c.dec.deg)
        except Exception as e:
            print(f"{t['name']}: SED fetch failed: {str(e)[:60]}")
            status.append((t['name'], 'fetch-fail'))
            continue
        if len(tab) == 0:
            status.append((t['name'], 'empty'))
            continue
        freq = np.array(tab['sed_freq'], float)   # GHz
        flux = np.array(tab['sed_flux'], float)   # Jy
        efl = np.array(tab['sed_eflux'], float)
        lam = 2.998e5 / freq                       # um
        nufnu = freq * 1e9 * flux * 1e-26          # W/m2
        ok = np.isfinite(lam) & np.isfinite(nufnu) & (nufnu > 0)
        fig, ax = plt.subplots(figsize=(6.5, 4.5))
        ax.scatter(lam[ok], nufnu[ok], s=14, c=np.log10(lam[ok]), cmap='plasma',
                   alpha=0.8, edgecolors='none')
        ax.set_xscale('log')
        ax.set_yscale('log')
        ax.set_xlabel(r'$\lambda$ [$\mu$m]')
        ax.set_ylabel(r'$\nu F_\nu$ [W m$^{-2}$]')
        d = t['d_best']
        L = t.get('L', None)
        # Lbol -> expected integrated nuFnu scale line: Fbol = L/(4 pi d^2)
        ax.set_title(f"{t['name']}  VizieR photometry r<{RADIUS_AS:.0f}\"  d={d:.2f} kpc",
                     fontsize=10)
        ax.grid(alpha=0.2, which='both')
        fig.tight_layout()
        fig.savefig(fn, dpi=110)
        plt.close(fig)
        status.append((t['name'], f'ok:{ok.sum()}pts'))
        print(f"{t['name']}: {ok.sum()} photometry points")
        time.sleep(0.4)
    with open(os.path.join(OUT, 'sed_status.csv'), 'w', newline='') as fh:
        w = csv.writer(fh)
        w.writerow(['name', 'status'])
        w.writerows(status)


if __name__ == '__main__':
    main()