"""Step 1: download the raw public data (cached in RAW; re-runnable).

 (a) Gaia Catalogue of Nearby Stars (GCNS; Gaia Collaboration, Smart et al. 2021),
     VizieR J/A+A/649/A6/table1c, 331,312 rows, selected columns.
 (b) AllWISE Source Catalog rows (IRSA table allwise_p3as_psd) within 12 arcsec of
     every GCNS position propagated to epoch 2010.5 (the WISE cryogenic epoch),
     via IRSA TAP table upload in chunks.
 (c) Pecaut & Mamajek dwarf colour / bolometric-magnitude table (E. Mamajek's
     maintained version) for bolometric fluxes and spectral-class boundaries.
"""
import io
import os
import sys
import time

import numpy as np
import pandas as pd
from astropy.table import Table

from common import RAW, VIZIER_TAP, IRSA_TAP, tap_sync, tap_upload, http_get

GCNS_FILE = os.path.join(RAW, "gcns_table1c.csv")
AW_DIR = os.path.join(RAW, "allwise_chunks")
MAMAJEK = os.path.join(RAW, "EEM_dwarf_UBVIJHK_colors_Teff.txt")
CHUNK = 10000
RADIUS_DEG = 12.0 / 3600.0
CTRL_DIR = os.path.join(RAW, "allwise_control_chunks")
CTRL_RADIUS_DEG = 6.0 / 3600.0
CTRL_OFFSET_DEG = 90.0 / 3600.0

GCNS_QUERY = ('SELECT GaiaEDR3,RA_ICRS,DE_ICRS,Plx,e_Plx,pmRA,pmDE,"Gmag",RFG,BPmag,RFBP,RPmag,RFRP,'
              '"E(BP/RP)",RUWE,IPDfmp,GCNSprob,WDprob,Dist50,"2MASS",Jmag,e_Jmag,Hmag,e_Hmag,Ksmag,e_Ksmag,'
              'WISE,W1mag,e_W1mag,W2mag,e_W2mag,W3mag,e_W3mag,W4mag,e_W4mag FROM "J/A+A/649/A6/table1c"')

AW_COLS = ("designation ra dec glon glat w1mpro w1sigmpro w1snr w1rchi2 w2mpro w2sigmpro w2snr w2rchi2 "
           "w3mpro w3sigmpro w3snr w3rchi2 w4mpro w4sigmpro w4snr w4rchi2 nb na w1sat w2sat w3sat w4sat "
           "cc_flags ext_flg var_flg ph_qual moon_lev w3m w4m w3cov w4cov w3cc_map w4cc_map "
           "w1flux w1sigflux w2flux w2sigflux w3flux w3sigflux w3sky w3sigsk w3conf "
           "w4flux w4sigflux w4sky w4sigsk w4conf w3mag w3sigm w3flg w4mag w4sigm w4flg xscprox "
           "r_2mass n_2mass j_m_2mass j_msig_2mass h_m_2mass h_msig_2mass k_m_2mass k_msig_2mass").split()


def get_gcns():
    if os.path.exists(GCNS_FILE) and os.path.getsize(GCNS_FILE) > 5e7:
        return
    print("downloading GCNS ...", flush=True)
    txt = tap_sync(VIZIER_TAP, GCNS_QUERY, timeout=1800)
    if not txt.startswith("GaiaEDR3"):
        raise RuntimeError("GCNS download failed: " + txt[:500])
    with open(GCNS_FILE, "w") as f:
        f.write(txt)


def get_mamajek():
    if os.path.exists(MAMAJEK) and os.path.getsize(MAMAJEK) > 1e4:
        return
    raw = http_get("https://www.pas.rochester.edu/~emamajek/EEM_dwarf_UBVIJHK_colors_Teff.txt", post=False)
    with open(MAMAJEK, "wb") as f:
        f.write(raw)


def get_allwise(control=False):
    """control=False: AllWISE sources within 12" of each GCNS star (epoch 2010.5).
    control=True : AllWISE sources within 6" of a blank-sky control position 90" north
                   (south for dec > +89.9) of each star, used to measure the surface density
                   of background sources that would mimic an excess (chance alignments)."""
    outdir = CTRL_DIR if control else AW_DIR
    radius = CTRL_RADIUS_DEG if control else RADIUS_DEG
    os.makedirs(outdir, exist_ok=True)
    g = pd.read_csv(GCNS_FILE, usecols=["GaiaEDR3", "RA_ICRS", "DE_ICRS", "pmRA", "pmDE"])
    dt = 2010.5 - 2016.0
    dec = g.DE_ICRS + g.pmDE.fillna(0) * dt / 3.6e6
    ra = (g.RA_ICRS + g.pmRA.fillna(0) * dt / 3.6e6 / np.cos(np.radians(g.DE_ICRS))) % 360.0
    if control:
        dec = dec + np.where(dec > 89.9, -1.0, 1.0) * CTRL_OFFSET_DEG
    sid = g.GaiaEDR3.values.astype("int64")
    n = len(g)
    sel = ", ".join("a." + c for c in AW_COLS)
    q = (f"SELECT u.sid, u.ra AS ra_in, u.dec AS dec_in, {sel} FROM allwise_p3as_psd a, TAP_UPLOAD.mytable u "
         f"WHERE CONTAINS(POINT('ICRS',a.ra,a.dec),CIRCLE('ICRS',u.ra,u.dec,{radius:.6f}))=1")
    nchunk = (n + CHUNK - 1) // CHUNK
    for k in range(nchunk):
        out = os.path.join(outdir, f"aw_{k:03d}.csv")
        if os.path.exists(out) and os.path.getsize(out) > 1000:
            continue
        s = slice(k * CHUNK, min(n, (k + 1) * CHUNK))
        t = Table({"sid": sid[s], "ra": ra.values[s], "dec": dec.values[s]})
        t0 = time.time()
        txt = tap_upload(IRSA_TAP, q, t, timeout=1800)
        if not txt.startswith("sid"):
            raise RuntimeError(f"IRSA chunk {k} failed: {txt[:500]}")
        with open(out + ".tmp", "w") as f:
            f.write(txt)
        os.replace(out + ".tmp", out)
        print(f"chunk {k+1}/{nchunk}: {txt.count(chr(10))-1} rows, {time.time()-t0:.0f}s", flush=True)


if __name__ == "__main__":
    get_gcns()
    get_mamajek()
    if "--control-only" not in sys.argv:
        get_allwise()
    get_allwise(control=True)
    print("download step done")