#!/usr/bin/env python3
"""Download every raw input used by the fp_kiii study into RAW (re-runnable; cached).

Sources
  * 6dFGSv Fundamental Plane data, Campbell et al. 2014 (VizieR J/MNRAS/443/1231)
  * SDSS PV catalogue, Howlett et al. 2022 (Zenodo record 6824749)
  * SDSS DR17 SkyServer SQL: our own, un-clipped early-type parent sample
"""
import os
import sys
import time
import urllib.parse
import urllib.request

RAW = os.environ.get("FP_KIII_RAW", "fermi_raw/fp_kiii")
VIZ = "https://tapvizier.cds.unistra.fr/TAPVizieR/tap/sync"
SKY = "https://skyserver.sdss.org/dr17/SkyServerWS/SearchTools/SqlSearch"


def fetch(url, dest, data=None, timeout=900, tries=8, min_bytes=50):
    if os.path.exists(dest) and os.path.getsize(dest) >= min_bytes:
        return dest
    os.makedirs(os.path.dirname(dest), exist_ok=True)
    last = None
    for k in range(tries):
        try:
            req = urllib.request.Request(url, data=data)
            with urllib.request.urlopen(req, timeout=timeout) as r:
                blob = r.read()
            if len(blob) < min_bytes or blob[:1] == b"{" or b"<html" in blob[:300].lower():
                raise RuntimeError("bad payload: %r" % blob[:200])
            with open(dest + ".part", "wb") as f:
                f.write(blob)
            os.replace(dest + ".part", dest)
            return dest
        except Exception as e:  # noqa
            last = e
            time.sleep(15 * (k + 1))
    raise RuntimeError("failed to fetch %s: %s" % (url[:200], last))


def vizier(table, dest):
    q = {"REQUEST": "doQuery", "LANG": "ADQL", "FORMAT": "csv", "MAXREC": "500000",
         "QUERY": 'SELECT * FROM "%s"' % table}
    return fetch(VIZ, dest, data=urllib.parse.urlencode(q).encode())


SDSS_SQL = """
SELECT s.specObjID, s.bestObjID, s.plate, s.mjd, s.fiberID, s.ra, s.dec, s.z, s.zErr,
 s.velDisp, s.velDispErr, s.snMedian, s.legacy_target1,
 p.petroMag_r, p.petroMagErr_r, p.petroR50_r, p.petroR90_r, p.petroR50_i, p.petroR90_i,
 p.deVMag_g, p.deVMag_r, p.deVMag_i, p.deVMagErr_r, p.deVRad_r, p.deVRadErr_r, p.deVAB_r, p.deVAB_i, p.deVRad_i,
 p.modelMag_g, p.modelMag_r, p.modelMag_i, p.cModelMag_r, p.fiberMag_r,
 p.extinction_g, p.extinction_r, p.extinction_i, p.fracDeV_r,
 p.lnLDeV_r, p.lnLExp_r, p.lnLDeV_i, p.lnLExp_i, p.clean, p.flags_r, p.nChild, p.parentID,
 l.h_alpha_eqw, l.h_alpha_eqw_err, x.d4000_n,
 e.sigmaStars, e.sigmaStarsErr, e.EW_Ha_6562, e.bpt
FROM SpecObj s
JOIN PhotoObjAll p ON s.bestObjID = p.objID
LEFT JOIN galSpecLine l ON l.specObjID = s.specObjID
LEFT JOIN galSpecIndx x ON x.specObjID = s.specObjID
LEFT JOIN emissionLinesPort e ON e.specObjID = s.specObjID
WHERE s.class='GALAXY' AND s.zWarning=0 AND s.sciencePrimary=1 AND s.survey='sdss'
 AND s.z BETWEEN 0.0033 AND 0.12 AND s.velDisp > 50 AND p.fracDeV_r >= 0.5
 AND s.ra >= {lo} AND s.ra < {hi}
"""


def sdss_chunks(step=4):
    out = []
    for lo in range(0, 360, step):
        dest = os.path.join(RAW, "sdss_own", "sdss_etg_ra%03d_%03d.csv" % (lo, lo + step))
        sql = " ".join(SDSS_SQL.format(lo=lo, hi=lo + step).split())
        url = SKY + "?" + urllib.parse.urlencode({"cmd": sql, "format": "csv"})
        # an empty RA slice returns only the header (~900 bytes) which is fine
        fetch(url, dest, timeout=900, min_bytes=100)
        out.append(dest)
        print("  sdss chunk", lo, os.path.getsize(dest), flush=True)
    return out


def main():
    for t in ("FPsample", "table2", "table9"):
        vizier("J/MNRAS/443/1231/" + t, os.path.join(RAW, "6df_%s.csv" % t))
    print("6dFGSv tables ok", flush=True)
    for f in ("SDSS_PV_public.dat", "data_description.pdf"):
        fetch("https://zenodo.org/api/records/6824749/files/%s/content" % f,
              os.path.join(RAW, "sdss_pv", f))
    print("SDSS PV (Howlett 2022) ok", flush=True)
    sdss_chunks()
    print("SDSS DR17 own sample ok", flush=True)


if __name__ == "__main__":
    sys.exit(main())