"""Step 5 (library, called from s07): same-velocity neighbours from SIMBAD within 10 arcmin.

Why: Cosmicflows-4 line widths come partly from large-beam telescopes (GBT ~9', Effelsberg ~9',
Parkes ~14', Nancay 4'x22' FWHM).  A second gas-rich galaxy inside the beam at the same velocity
broadens the HI profile and makes the target look under-luminous for its "rotation velocity".
The SDSS-based checks only reach 2-2.5' (Arecibo beam) and only inside the SDSS footprint, so we
query SIMBAD (all-sky, heterogeneous completeness) with a TAP upload join.

Returned table: one row per (target, SIMBAD object with a radial velocity within +-1000 km/s and
within 10'); the neighbour definition itself is applied in s07.
"""
import io
import os
import subprocess
import time
import numpy as np
import pandas as pd
from astropy.table import Table
from common import RAW

SIMBAD_TAP = "https://simbad.cds.unistra.fr/simbad/sim-tap/sync"
RADIUS_DEG = 10.0 / 60.0
CHUNK = 1000
# SIMBAD object types that are individual galaxies (groups/clusters/pairs, SNe, HII regions and
# radio/HI source entries are excluded: they are not separate galaxies or may be the target itself)
GALAXY_OTYPES = {"G", "AG?", "GiC", "GiG", "GiP", "EmG", "LSB", "LIN", "Sy1", "Sy2", "SyG", "SBG", "H2G",
                 "rG", "IG", "BiC", "AGN", "BLL", "QSO", "bCG", "Bla", "HzG", "ERO", "LyA", "DLA", "mAL"}

QUERY = """SELECT u.idx, b.main_id, b.ra, b.dec, b.rvz_radvel, b.otype, b.galdim_majaxis,
 DISTANCE(POINT('ICRS',b.ra,b.dec),POINT('ICRS',u.ra,u.dec))*60 AS sep_arcmin
 FROM TAP_UPLOAD.up AS u JOIN basic AS b
 ON 1=CONTAINS(POINT('ICRS',b.ra,b.dec),CIRCLE('ICRS',u.ra,u.dec,%.5f))
 WHERE b.rvz_radvel IS NOT NULL AND b.rvz_radvel BETWEEN u.cz-1000 AND u.cz+1000""" % RADIUS_DEG


def _query_chunk(df, tag):
    vot = os.path.join(RAW, f"simbad_up_{tag}.vot")
    Table.from_pandas(df[["idx", "ra", "dec", "cz"]]).write(vot, format="votable", overwrite=True)
    last = ""
    for k in range(4):
        out = subprocess.run(
            ["curl", "-sS", "--max-time", "280", "-F", "REQUEST=doQuery", "-F", "LANG=ADQL", "-F", "FORMAT=csv",
             "-F", "MAXREC=200000", "-F", "UPLOAD=up,param:up", "-F", f"up=@{vot}", "-F", "QUERY=" + QUERY,
             SIMBAD_TAP], capture_output=True, text=True)
        if out.returncode == 0 and out.stdout.startswith("idx,"):
            os.remove(vot)
            return pd.read_csv(io.StringIO(out.stdout))
        last = (out.stdout[:300] + out.stderr[:300])
        time.sleep(5 * (k + 1))
    raise RuntimeError(f"SIMBAD TAP failed for chunk {tag}: {last}")


def simbad_neighbours(uid, ra, dec, cz, cache="simbad_nb.csv"):
    """All SIMBAD objects with |v - cz| < 1000 km/s within 10' of each target (cached in RAW)."""
    path = os.path.join(RAW, cache)
    targets = pd.DataFrame(dict(uid=np.asarray(uid), ra=np.asarray(ra, float), dec=np.asarray(dec, float),
                                cz=np.asarray(cz, float)))
    done = pd.DataFrame(columns=["uid"])
    if os.path.exists(path) and os.path.exists(path + ".queried"):
        done = pd.read_csv(path + ".queried")
        res = [pd.read_csv(path)]
    else:
        res = []
    todo = targets[~targets.uid.isin(done.uid)].reset_index(drop=True)
    todo["idx"] = np.arange(len(todo))
    for k in range(0, len(todo), CHUNK):
        ch = todo.iloc[k:k + CHUNK]
        t0 = time.time()
        d = _query_chunk(ch, f"{k}")
        d["uid"] = d.idx.map(ch.set_index("idx").uid)
        res.append(d.drop(columns="idx"))
        print(f"    SIMBAD chunk {k // CHUNK + 1}/{(len(todo) - 1) // CHUNK + 1}: {len(d)} rows in {time.time() - t0:.0f}s")
    out = pd.concat(res, ignore_index=True) if res else pd.DataFrame()
    if len(todo):
        out.to_csv(path, index=False)
        pd.concat([done, todo[["uid"]]]).to_csv(path + ".queried", index=False)
    return out[out.uid.isin(targets.uid)].reset_index(drop=True)


if __name__ == "__main__":
    print(__doc__)