"""Download every public table used by the study (VizieR TAP) into RAW_DIR.

Tables (all from VizieR, https://tapvizier.cds.unistra.fr/TAPVizieR/tap):
  Botelho et al. 2019, MNRAS 482, 1690   J/MNRAS/482/1690  table1 (parameters, ages, Si, Nd, Eu), table3 (Th)
  Spina et al. 2018,  MNRAS 474, 2580    J/MNRAS/474/2580  table2 (parameters), table3 (n-capture incl. Eu), table4 (ages)
  Bedell et al. 2018, ApJ 865, 68        J/ApJ/865/68      table2 (light elements incl. Mg, Si)
  Delgado Mena et al. 2017, A&A 606, A94 J/A+A/606/A94     table1 (parameters), table3 (Eu, Mg, Si)
  Delgado Mena et al. 2019, A&A 624, A78 J/A+A/624/A78     table2 (masses, ages)
  Battistini & Bensby 2016, A&A 586, A49  J/A+A/586/A49  abundances (Eu for F and G dwarfs)
  Bensby et al. 2014, A&A 562, A71        J/A+A/562/A71  tablec3 (ages, Mg, Si for the same stars)
  Mishenina et al. 2022, MNRAS 516, 3786  J/MNRAS/516/3786 tablea1 (Th for 276 disc dwarfs; no ages; check only)
  Mowlavi et al. 2012, A&A 541, A41       J/A+A/541/A41  grids (Geneva tracks 0.5-1.0 Msun, Z = 0.014)
  GALAH DR3, Buder et al. 2021, MNRAS 506, 150   J/MNRAS/506/150  stars + ages (VAC)
  GALAH DR4, Buder et al. 2025, PASA 42, e051    J/other/PASA/42.51  allstar

Usage:  python fetch_data.py [--refresh]
Files already present are kept unless --refresh is given.
"""
import os
import sys
import time
import urllib.parse
import urllib.request

from common import RAW_DIR

TAP = "https://tapvizier.cds.unistra.fr/TAPVizieR/tap/sync"

# Quality cuts for the GALAH samples (Sun-like dwarfs, turn-off stars and subgiants):
#   flag_sp = 0 and flag_fe_h = 0 (recommended by the GALAH release papers), 5300 < Teff < 6400 K,
#   3.6 < log g < 4.7, S/N per pixel in CCD 2 > 40, and at least one of [Eu/Fe], [K/Fe] unflagged.
# (The last condition makes these downloads a SUBSET of all Sun-like dwarfs; the *_count.csv queries
# below give the size of the full set, which is the correct denominator for detection fractions.)
GALAH3 = (
    'SELECT s.GALAH, s.Teff, s.logg, s."[Fe/H]", s."e_[Fe/H]", s.snrc2iraf, s.snrc3iraf, '
    's."[Mg/Fe]", s."e_[Mg/Fe]", s."f_[Mg/Fe]", s."[Si/Fe]", s."e_[Si/Fe]", s."f_[Si/Fe]", '
    's."[K/Fe]", s."e_[K/Fe]", s."f_[K/Fe]", s."[Eu/Fe]", s."e_[Eu/Fe]", s."f_[Eu/Fe]", '
    'a.Age, a.e_Age, a.Age16, a.Age50, a.Age84, a.Mact '
    'FROM "J/MNRAS/506/150/stars" AS s JOIN "J/MNRAS/506/150/ages" AS a ON s.GALAH = a.GALAH '
    'WHERE s.Flagsp=0 AND s."q_[Fe/H]"=0 AND s.Teff BETWEEN 5300 AND 6400 AND s.logg BETWEEN 3.6 AND 4.7 '
    'AND s.snrc2iraf>40 AND (s."f_[Eu/Fe]"=0 OR s."f_[K/Fe]"=0)'
)
GALAH4 = (
    'SELECT s.GALAH, s.Teff, s.logg, s."[Fe/H]", s."e_[Fe/H]", s.SNRpxCCD2, s.SNRpxCCD3, '
    's."[Mg/Fe]", s."e_[Mg/Fe]", s."f_[Mg/Fe]", s."[Si/Fe]", s."e_[Si/Fe]", s."f_[Si/Fe]", '
    's."[K/Fe]", s."e_[K/Fe]", s."f_[K/Fe]", s."[Eu/Fe]", s."e_[Eu/Fe]", s."f_[Eu/Fe]", s.Mass, s.Age '
    'FROM "J/other/PASA/42.51/allstar" AS s '
    'WHERE s.flagSp=0 AND s."f_[Fe/H]"=0 AND s.Teff BETWEEN 5300 AND 6400 AND s.logg BETWEEN 3.6 AND 4.7 '
    'AND s.SNRpxCCD2>40 AND (s."f_[Eu/Fe]"=0 OR s."f_[K/Fe]"=0)'
)

# Denominators for the GALAH europium detection fraction: all Sun-like dwarfs passing the same cuts,
# WITHOUT the requirement that Eu or K be unflagged (the downloads above are pre-filtered on that).
GALAH3_COUNT = ('SELECT COUNT(*) AS n FROM "J/MNRAS/506/150/stars" AS s JOIN "J/MNRAS/506/150/ages" AS a ON s.GALAH = a.GALAH '
                'WHERE s.Flagsp=0 AND s."q_[Fe/H]"=0 AND s.Teff BETWEEN 5300 AND 6400 AND s.logg BETWEEN 3.6 AND 4.7 '
                'AND s.snrc2iraf>40')
GALAH4_COUNT = ('SELECT COUNT(*) AS n FROM "J/other/PASA/42.51/allstar" AS s '
                'WHERE s.flagSp=0 AND s."f_[Fe/H]"=0 AND s.Teff BETWEEN 5300 AND 6400 AND s.logg BETWEEN 3.6 AND 4.7 '
                'AND s.SNRpxCCD2>40')

QUERIES = {
    "botelho2019_table1.csv": 'SELECT * FROM "J/MNRAS/482/1690/table1"',
    "botelho2019_table3.csv": 'SELECT * FROM "J/MNRAS/482/1690/table3"',
    "spina2018_table2.csv": 'SELECT * FROM "J/MNRAS/474/2580/table2"',
    "spina2018_table3.csv": 'SELECT * FROM "J/MNRAS/474/2580/table3"',
    "spina2018_table4.csv": 'SELECT * FROM "J/MNRAS/474/2580/table4"',
    "bedell2018_table2.csv": 'SELECT * FROM "J/ApJ/865/68/table2"',
    "delgadomena2017_table1.csv": 'SELECT * FROM "J/A+A/606/A94/table1"',
    "delgadomena2017_table3.csv": 'SELECT * FROM "J/A+A/606/A94/table3"',
    "delgadomena2019_table2.csv": 'SELECT * FROM "J/A+A/624/A78/table2"',
    "battistini2016_abundances.csv": 'SELECT recno, HIP, Teff, logg, "[Fe/H]", "[Eu/H]", "[Nd/H]", "[Sm/H]" '
                                     'FROM "J/A+A/586/A49/abundances"',
    "bensby2014_tablec3.csv": 'SELECT * FROM "J/A+A/562/A71/tablec3"',
    "mishenina2022_tablea1.csv": 'SELECT * FROM "J/MNRAS/516/3786/tablea1"',
    "geneva_tracks_lowmass.csv": 'SELECT Phase, Minit, Zinit, Age, logL, logTe, "Xc(H)" FROM "J/A+A/541/A41/grids" '
                                 'WHERE Minit < 1.05 AND Zinit BETWEEN 0.0139 AND 0.0141',
    "galah_dr3_sunlike.csv": GALAH3,
    "galah_dr4_sunlike.csv": GALAH4,
    "galah_dr3_sunlike_count.csv": GALAH3_COUNT,
    "galah_dr4_sunlike_count.csv": GALAH4_COUNT,
}


def tap_query(adql, timeout=600, tries=4):
    data = urllib.parse.urlencode(
        {"REQUEST": "doQuery", "LANG": "ADQL", "FORMAT": "csv", "MAXREC": "2000000", "QUERY": adql}).encode()
    last = None
    for i in range(tries):
        try:
            with urllib.request.urlopen(urllib.request.Request(TAP, data=data), timeout=timeout) as r:
                txt = r.read().decode("utf-8", errors="replace")
            if txt.lstrip().startswith("<?xml"):
                raise RuntimeError("TAP error: " + txt[:600])
            return txt
        except Exception as exc:  # noqa: BLE001  (network errors are retried)
            last = exc
            time.sleep(3 + 5 * i)
    raise RuntimeError(f"TAP query failed after {tries} tries: {last}")


def fetch_all(refresh=False, verbose=True):
    paths = {}
    for name, adql in QUERIES.items():
        path = os.path.join(RAW_DIR, name)
        if refresh or not os.path.exists(path) or os.path.getsize(path) < (3 if name.endswith("_count.csv") else 200):
            if verbose:
                print(f"[fetch] {name} ...", flush=True)
            txt = tap_query(adql)
            with open(path, "w") as f:
                f.write(txt)
            time.sleep(1.0)
        if verbose:
            with open(path) as f:
                n = sum(1 for _ in f) - 1
            print(f"[fetch] {name}: {n} rows", flush=True)
        paths[name] = path
    return paths


if __name__ == "__main__":
    fetch_all(refresh="--refresh" in sys.argv)