"""Step 1: obtain all inputs.

Small catalogue snapshots are kept in data/ (so that every number in the
report can be regenerated from exactly the tables that were analysed);
bulk astrometry goes to the raw directory.  `--refresh` re-downloads
everything from the public services.

Sources
  * JPL SBDB Query API   https://ssd-api.jpl.nasa.gov/sbdb_query.api
  * JPL radar astrometry https://ssd-api.jpl.nasa.gov/sb_radar.api
  * JPL Horizons API     https://ssd.jpl.nasa.gov/api/horizons.api
  * MPC observations     https://www.minorplanetcenter.net/iau/ECS/MPCAT-OBS/UnnObs.txt.gz
                         https://data.minorplanetcenter.net/api/get-obs
  * MPC observatory codes https://www.minorplanetcenter.net/iau/lists/ObsCodes.html

Revision note: space-based astrometry (NEOWISE, TESS, HST, ...) is now kept together
with the observer position given in the paired 's' record; in the first version these
observations were silently placed at the geocentre.  Roving-observer records are dropped.
"""
import gzip
import io
import json
import os
import re
import subprocess
import sys
import time
import urllib.parse
import urllib.request

import numpy as np
import pandas as pd

from common import DATA, RAW, pack_prov, mjd_from_ymd, H_10M, H_100M

SBDB = "https://ssd-api.jpl.nasa.gov/sbdb_query.api"
FIELDS = ("spkid,full_name,pdes,name,prefix,kind,neo,pha,H,G,diameter,albedo,rot_per,"
          "spec_B,spec_T,orbit_id,epoch_mjd,e,a,q,i,om,w,ma,tp,per,moid,class,source,"
          "soln_date,producer,data_arc,first_obs,last_obs,n_obs_used,n_del_obs_used,"
          "n_dop_obs_used,condition_code,rms,A1,A1_sigma,A2,A2_sigma,A3,A3_sigma,DT,DT_sigma")
UNNOBS_URL = "https://www.minorplanetcenter.net/iau/ECS/MPCAT-OBS/UnnObs.txt.gz"
SEED = 20261012
FREEZE = "2026-05-25"       # UnnObs.txt.gz is dated 2026-06-03; later astrometry comes from the MPC API


def _get(url, data=None, headers=None, timeout=300, retries=4):
    last = None
    for k in range(retries):
        try:
            req = urllib.request.Request(url, data=data, headers=headers or {})
            if data is not None:
                req.get_method = lambda: "GET"       # MPC API wants GET with a JSON body
            with urllib.request.urlopen(req, timeout=timeout) as r:
                return r.read()
        except Exception as exc:          # noqa
            last = exc
            time.sleep(2 + 3 * k)
    raise RuntimeError("download failed: %s (%s)" % (url[:120], last))


def sbdb_table(path, query, refresh):
    if os.path.exists(path) and not refresh:
        return pd.read_csv(path, dtype={"pdes": str, "prefix": str})
    url = SBDB + "?fields=" + FIELDS + "&full-prec=true&" + query
    d = json.loads(_get(url))
    df = pd.DataFrame(d["data"], columns=d["fields"])
    df["full_name"] = df["full_name"].str.strip()
    df.to_csv(path, index=False)
    print("  SBDB ->", os.path.basename(path), len(df), "rows")
    return pd.read_csv(path, dtype={"pdes": str, "prefix": str})


def radar_table(refresh):
    path = os.path.join(DATA, "jpl_radar_astrometry.csv")
    if os.path.exists(path) and not refresh:
        return pd.read_csv(path, dtype={"des": str})
    d = json.loads(_get("https://ssd-api.jpl.nasa.gov/sb_radar.api"))
    df = pd.DataFrame(d["data"], columns=d["fields"])
    df.to_csv(path, index=False)
    print("  radar ->", len(df), "rows")
    return pd.read_csv(path, dtype={"des": str})


def obscodes(refresh):
    path = os.path.join(DATA, "mpc_obscodes.csv")
    if os.path.exists(path) and not refresh:
        return pd.read_csv(path, dtype={"code": str})
    txt = _get("https://www.minorplanetcenter.net/iau/lists/ObsCodes.html").decode("utf-8", "replace")
    rows = []
    for line in txt.splitlines():
        if len(line) < 30 or line.startswith("Code") or line.startswith("<"):
            continue
        code = line[0:3]
        try:
            lon = float(line[4:13]); rc = float(line[13:21]); rs = float(line[21:30])
        except ValueError:
            continue      # space-based or roving: no fixed parallax constants
        rows.append((code, lon, rc, rs))
    df = pd.DataFrame(rows, columns=["code", "lon_deg", "rho_cos", "rho_sin"])
    df.to_csv(path, index=False)
    print("  obscodes ->", len(df))
    return pd.read_csv(path, dtype={"code": str})


# ------------------------------------------------------------ MPC astrometry
def _mjd_of(ln):
    return float(mjd_from_ymd(int(ln[15:19]), int(ln[20:22]), float(ln[23:32])))


def parse_obs80(lines):
    """Parse MPC 80-column optical records.

    Returns a list of (mjd_utc, ra_deg, dec_deg, note2, catalogue, station, ox, oy, oz).
    ox, oy, oz is the geocentric position of a space-based observer [au, equatorial
    J2000], read from the paired 's' record (blank for ground stations, whose
    position is computed from the MPC parallax constants).  Space-based records
    without a usable 's' record are dropped, as are roving-observer ('V'/'v'),
    radar ('R'/'r') and deleted ('X'/'x') records."""
    sat = {}
    for ln in lines:
        if len(ln) >= 80 and ln[14] == "s":
            try:
                unit = 1.0 / 1.495978707e8 if ln[32] == "1" else (1.0 if ln[32] == "2" else None)
                if unit is None:
                    continue
                xyz = []
                for a in (34, 46, 58):
                    v = float(ln[a + 1:a + 11])
                    xyz.append((-v if ln[a] == "-" else v) * unit)
                sat[(ln[0:12], ln[15:32], ln[77:80])] = xyz
            except ValueError:
                continue
    out = []
    for ln in lines:
        if len(ln) < 80:
            continue
        n2 = ln[14]
        if n2 in "sRrvVXx":
            continue
        try:
            mjd = _mjd_of(ln)
            ra = (float(ln[32:34]) + float(ln[35:37]) / 60.0 + float(ln[38:44] or 0) / 3600.0) * 15.0
            sgn = -1.0 if ln[44] == "-" else 1.0
            sec = ln[51:56].strip()
            dec = sgn * (float(ln[45:47]) + float(ln[48:50]) / 60.0 + (float(sec) if sec else 0.0) / 3600.0)
        except ValueError:
            continue
        if n2 == "S":
            xyz = sat.get((ln[0:12], ln[15:32], ln[77:80]))
            if xyz is None:
                continue
            ox, oy, oz = xyz
        else:
            ox = oy = oz = float("nan")
        out.append((mjd, ra, dec, n2 if n2.strip() else "P", ln[71], ln[77:80], ox, oy, oz))
    return out


OBS_FILE = "neo_unn_obs_v2.csv.gz"     # v2: includes the positions of space-based observers


def extract_unnobs(neo, refresh):
    """Stream the MPC unnumbered-object observation file and keep the records
    of unnumbered NEOs.  Output: RAW/neo_unn_obs_v2.csv.gz"""
    out = os.path.join(RAW, OBS_FILE)
    if os.path.exists(out) and not refresh:
        return out
    gz = os.path.join(RAW, "UnnObs.txt.gz")
    if not os.path.exists(gz) or refresh:
        print("  downloading UnnObs.txt.gz (456 MB) ...")
        subprocess.run(["curl", "-sS", "-o", gz, UNNOBS_URL], check=True)
    unn = neo[neo.kind == "au"]
    packed = {}
    for p in unn.pdes:
        k = pack_prov(p)
        if k:
            packed[k] = p
    print("  unnumbered NEOs with packed designation:", len(packed), "of", len(unn))
    n = 0
    n_sat = 0

    def flush(p, group, g):
        nonlocal n, n_sat
        buf = []
        for (mjd, ra, dec, n2, cat, stn, ox, oy, oz) in parse_obs80(group):
            if ox == ox:
                xyz = "%.10f,%.10f,%.10f" % (ox, oy, oz); n_sat += 1
            else:
                xyz = ",,"
            buf.append("%s,%.7f,%.6f,%.6f,%s,%s,%s,%s\n" % (p, mjd, ra, dec, n2, cat.replace(",", " "), stn, xyz))
        g.writelines(buf)
        n += len(buf)

    with gzip.open(gz, "rt", encoding="latin-1") as f, gzip.open(out, "wt") as g:
        g.write("pdes,mjd_utc,ra,dec,note2,cat,stn,ox,oy,oz\n")
        cur, group = None, []
        for ln in f:
            key = ln[5:12]
            if key != cur:
                if group:
                    flush(packed[cur], group, g)
                cur, group = (key, []) if key in packed else (None, [])
            if cur is not None:
                group.append(ln.rstrip("\n").ljust(80))
        if group:
            flush(packed[cur], group, g)
    print("  kept", n, "observations (%d space-based with observer positions) ->" % n_sat, out)
    return out


def mpc_api_obs(pdes, refresh=False):
    """Astrometry of one object from the MPC observations API (cached)."""
    d = os.path.join(RAW, "mpc_api")
    os.makedirs(d, exist_ok=True)
    path = os.path.join(d, pdes.replace(" ", "_").replace("/", "_") + ".txt")
    if not os.path.exists(path) or refresh:
        body = json.dumps({"desigs": [pdes], "output_format": ["OBS80"]}).encode()
        raw = _get("https://data.minorplanetcenter.net/api/get-obs", data=body,
                   headers={"Content-Type": "application/json"}, timeout=180)
        j = json.loads(raw)
        txt = j[0].get("OBS80", "") if isinstance(j, list) and j and isinstance(j[0], dict) else ""
        with open(path, "w") as f:
            f.write(txt or "")
        time.sleep(0.2)
    with open(path) as f:
        return parse_obs80([ln.rstrip("\n").ljust(80) for ln in f])


# ------------------------------------------------------------------ Horizons
_HZ_RE = re.compile(r"\$\$SOE\s*(.*?)\$\$EOE", re.S)


def horizons_state(pdes, jd_tdb):
    """Heliocentric ICRF state [au, au/d] of a small body at jd_tdb from JPL's
    own orbit solution (full N-body)."""
    cmd = "'DES=%s;'" % pdes if not pdes.isdigit() else "'%s;'" % pdes
    q = dict(format="text", COMMAND=cmd, OBJ_DATA="NO", MAKE_EPHEM="YES", EPHEM_TYPE="VECTORS",
             CENTER="'500@10'", TLIST="'%.6f'" % jd_tdb, REF_PLANE="FRAME", VEC_TABLE="2",
             OUT_UNITS="AU-D", CSV_FORMAT="YES")
    txt = _get("https://ssd.jpl.nasa.gov/api/horizons.api?" + urllib.parse.urlencode(q), timeout=120).decode()
    m = _HZ_RE.search(txt)
    if not m:
        return None
    v = [s.strip() for s in m.group(1).strip().split(",")]
    return [float(x) for x in v[2:8]]


def fetch_horizons(plan, refresh=False):
    """plan: DataFrame with pdes, jd_ref.  Cached incrementally in data/."""
    path = os.path.join(DATA, "horizons_states.csv")
    have = {}
    if os.path.exists(path) and not refresh:
        old = pd.read_csv(path, dtype={"pdes": str})
        have = {(r.pdes, round(r.jd_ref, 3)): r for r in old.itertuples()}
    rows = [tuple(r)[1:] for r in have.values()]
    todo = [(r.pdes, r.jd_ref) for r in plan.itertuples() if (r.pdes, round(r.jd_ref, 3)) not in have]
    print("  Horizons states: have %d, to fetch %d" % (len(have), len(todo)))
    t0 = time.time()
    for k, (p, jd) in enumerate(todo):
        try:
            s = horizons_state(p, jd)
        except Exception as exc:      # noqa
            s = None
        if s is None:
            s = [np.nan] * 6
        rows.append((p, jd) + tuple(s))
        if (k + 1) % 100 == 0 or k + 1 == len(todo):
            pd.DataFrame(rows, columns=["pdes", "jd_ref", "x", "y", "z", "vx", "vy", "vz"]).to_csv(path, index=False)
            print("    %d/%d  (%.0f s)" % (k + 1, len(todo), time.time() - t0), flush=True)
    return pd.read_csv(path, dtype={"pdes": str})


# ---------------------------------------------------------------- the plan
def build_plan(neo, ng):
    """Decide, deterministically, which objects are analysed with which model
    and which external look-ups they need."""
    df = neo.copy()
    t0 = pd.Timestamp("1858-11-17")
    df["mjd_first"] = (pd.to_datetime(df.first_obs, errors="coerce") - t0).dt.days
    df["mjd_last"] = (pd.to_datetime(df.last_obs, errors="coerce") - t0).dt.days
    df["arc_d"] = df.mjd_last - df.mjd_first
    freeze = (pd.Timestamp(FREEZE) - t0).days
    in_arc = (df.epoch_mjd >= df.mjd_first - 30) & (df.epoch_mjd <= df.mjd_last + 30)
    # objects whose arc extends past the bulk-file date get their astrometry from the MPC API
    df["post_freeze"] = df.mjd_last > freeze
    df["ocat"] = np.select(
        [df.H.isna() | df.mjd_first.isna(), df.arc_d < 1,
         in_arc & (df.arc_d <= 180), df.arc_d <= 180, df.kind.eq("an")],
        ["no_H_or_date", "arc_lt_1d", "short_inarc", "short_std", "long_numbered"],
        "long_unnumbered")
    df["sbin"] = np.where(df.H > H_10M, "1-10 m", np.where(df.H > H_100M, "10-100 m", ">100 m"))
    has_ng = df.pdes.isin(set(ng.pdes))
    rng = np.random.default_rng(SEED)
    u = pd.Series(rng.random(len(df)), index=df.sort_values("pdes").index).reindex(df.index)
    df["u"] = u
    # sampling fractions (weights = 1/fraction); everything else is analysed in full
    df["frac"] = 1.0
    m = (df.ocat == "short_std") & (df.sbin != "1-10 m") & ~has_ng
    df.loc[m, "frac"] = 0.20
    m2 = (df.ocat == "long_numbered") & ~has_ng
    df.loc[m2, "frac"] = 0.20
    # multi-apparition 10-100 m objects need the (slow) numerical model: 50% sample
    m3 = df.ocat.isin(["long_numbered", "long_unnumbered"]) & (df.sbin == "10-100 m") & df.A1.isna()
    df.loc[m3, "frac"] = 0.50
    df["selected"] = (df.u < df.frac) & ~df.ocat.isin(["no_H_or_date", "arc_lt_1d"])
    df["need_horizons"] = df.selected & (df.ocat == "short_std")
    df["need_mpc_api"] = df.selected & ((df.kind == "an") | df.post_freeze)
    df["jd_ref"] = np.where(df.ocat == "short_std", 0.5 * (df.mjd_first + df.mjd_last) + 2400000.5,
                            df.epoch_mjd + 2400000.5)
    return df


def main(refresh=False, skip_remote=False, only=None):
    print("[1] downloads")
    ng = sbdb_table(os.path.join(DATA, "sbdb_nongrav_all.csv"),
                    "sb-cdata=" + urllib.parse.quote('{"OR":["A1|DF","A2|DF","A3|DF"]}'), refresh)
    neo = sbdb_table(os.path.join(DATA, "sbdb_neo_asteroids.csv"), "sb-group=neo&sb-kind=a", refresh)
    radar_table(refresh)
    obscodes(refresh)
    plan = build_plan(neo, ng)
    keep = ["pdes", "full_name", "kind", "H", "epoch_mjd", "e", "a", "q", "i", "om", "w", "ma", "class",
            "data_arc", "first_obs", "last_obs", "n_obs_used", "n_del_obs_used", "n_dop_obs_used",
            "condition_code", "rms", "A1", "A1_sigma", "A2", "A2_sigma", "A3", "A3_sigma",
            "mjd_first", "mjd_last", "arc_d", "post_freeze", "ocat", "sbin", "u", "frac", "selected",
            "need_horizons", "need_mpc_api", "jd_ref"]
    plan = plan[keep]
    plan.to_csv(os.path.join(DATA, "analysis_plan.csv"), index=False)
    print(pd.crosstab(plan.ocat, plan.sbin, margins=True).to_string())
    print("  selected:", int(plan.selected.sum()), " Horizons:", int(plan.need_horizons.sum()),
          " MPC API:", int(plan.need_mpc_api.sum()))
    extract_unnobs(neo, refresh)
    if skip_remote:
        return
    # validation objects that are not NEO asteroids ('Oumuamua) are handled in s03
    if only in (None, "mpc"):
        api = list(plan[plan.need_mpc_api].pdes)
        t0 = time.time()
        for k, p in enumerate(api):
            mpc_api_obs(p, refresh)
            if (k + 1) % 200 == 0:
                print("    MPC API %d/%d (%.0f s)" % (k + 1, len(api), time.time() - t0), flush=True)
    if only in (None, "horizons"):
        fetch_horizons(plan[plan.need_horizons][["pdes", "jd_ref"]], refresh)


if __name__ == "__main__":
    _only = [a.split("=")[1] for a in sys.argv if a.startswith("--only=")]
    main(refresh="--refresh" in sys.argv, skip_remote="--skip-remote" in sys.argv,
         only=_only[0] if _only else None)