"""Daily geocentric ephemerides (JPL Horizons) 1998-2026 for every catalogued
near-Earth asteroid on an Earth-like or co-orbital orbit:
   set A: |a - 1 au| < 0.006            (co-orbital zone)
   set B: 0.9 < a < 1.1, e < 0.1, i < 5 (Earth-like; the orbits of spent rocket stages)
Used (s05) to test the survey simulator against the real discovery
circumstances of these objects.

For each object we fetch the daily r, Delta, phase angle, solar elongation,
sky rate, ecliptic latitude and declination (full arrays kept in the raw
directory) and store the derived daily g = V - H (with and without trailing loss),
solar elongation, declination flag and sky band in data/earthlike_geom.npz, with
data/earthlike_objects.csv.  Slow (about 2 s per object); cached.  Run with --refresh
to re-query.  (Revision: the survey weight is no longer stored, because the elongation
weight is now measured in s02; s05 computes it from the stored elongation.  If the raw
ephemeris cache exists the derived file is rebuilt from it without new queries.)
"""
import os
import re
import sys
import time
import urllib.parse

import numpy as np
import pandas as pd

import simlib as S
from common import DATA, RAW, fetch_text, phase_HG, trailing_loss
from s01_neo_completeness import load_neo

HOR = "https://ssd.jpl.nasa.gov/api/horizons.api"
START, STOP = "1998-01-01", "2026-10-01"

HALF_MONTH = "ABCDEFGHJKLMNOPQRSTUVWXY"


def designation_window(pdes):
    """(year_fraction_start, year_fraction_end) of the half-month encoded in a provisional
    designation such as '2016 HO3'; None if not a provisional designation."""
    m = re.match(r"^(\d{4}) ([A-HJ-Y])[A-Z]\d*$", pdes.strip())
    if not m:
        return None
    y = int(m.group(1))
    k = HALF_MONTH.index(m.group(2))
    return (y + k / 24.0, y + (k + 1) / 24.0)


def prov_from_fullname(full_name):
    m = re.search(r"\((\d{4} [A-Z]{2}\d*)\)", full_name)
    return m.group(1) if m else None


def query(pdes):
    cmd = f"'{pdes};'" if pdes.strip().isdigit() else f"'DES={pdes};'"
    q = {"format": "text", "COMMAND": cmd, "OBJ_DATA": "NO", "MAKE_EPHEM": "YES", "EPHEM_TYPE": "OBSERVER",
         "CENTER": "'500@399'", "START_TIME": f"'{START}'", "STOP_TIME": f"'{STOP}'", "STEP_SIZE": "'1 d'",
         "QUANTITIES": "'1,3,19,20,23,24,31'", "CSV_FORMAT": "YES", "ANG_FORMAT": "DEG"}
    txt = fetch_text(HOR + "?" + urllib.parse.urlencode(q), timeout=180, retries=3)
    if "$$SOE" not in txt:
        return None
    body = txt.split("$$SOE")[1].split("$$EOE")[0].strip().split("\n")
    out = np.empty((len(body), 7), np.float32)
    for k, line in enumerate(body):
        f = [x.strip() for x in line.split(",")]
        # 0 date, 1,2 flags, 3 RA, 4 DEC, 5 dRA*cosD, 6 dDEC, 7 r, 8 rdot, 9 delta, 10 deldot, 11 S-O-T, 12 /r, 13 S-T-O, 14 EcLon, 15 EcLat
        try:
            rate = np.hypot(float(f[5]), float(f[6])) * 24.0 / 3600.0      # arcsec/h -> deg/day
            out[k] = (float(f[7]), float(f[9]), float(f[13]), float(f[11]), rate, float(f[15]), float(f[4]))
        except Exception:
            out[k] = np.nan
    return out


def derive(eph, out_npz):
    """Compact derived geometry used by the simulator tests.  eph columns: r, delta, phase, elong, rate, eclat, dec."""
    r, d, ph, el, rate, eclat, dec = [eph[:, :, k].astype(float) for k in range(7)]
    g = 5 * np.log10(r * d) + phase_HG(np.radians(ph)) + trailing_loss(rate)
    g_notrail = 5 * np.log10(r * d) + phase_HG(np.radians(ph))
    band = (np.abs(eclat) > S.BETA_SPLIT)
    np.savez_compressed(out_npz, g_centimag=np.round(np.clip(g, -300, 300) * 100).astype(np.int16),
                        g_notrail_centimag=np.round(np.clip(g_notrail, -300, 300) * 100).astype(np.int16),
                        dec_ok=(dec > S.DEC_MIN).astype(np.uint8), band=band.astype(np.uint8),
                        elong_decideg=np.round(np.nan_to_num(el, nan=0.0) * 10).astype(np.int16), start=np.array(START))


def main(refresh=False):
    out_npz = os.path.join(DATA, "earthlike_geom.npz")
    out_csv = os.path.join(DATA, "earthlike_objects.csv")
    raw_npz = os.path.join(RAW, "earthlike_ephem_full.npz")
    if os.path.exists(out_npz) and os.path.exists(out_csv) and not refresh:
        if "dec_ok" in np.load(out_npz).files:
            print("[ephem] cached:", out_npz)
            return
        if os.path.exists(raw_npz):
            eph = np.load(raw_npz)["eph"]
            if eph.shape[0] == len(pd.read_csv(out_csv)):
                derive(eph, out_npz)
                print("[ephem] rebuilt", out_npz, "from the raw ephemeris cache", eph.shape)
                return
    neo = load_neo()
    A = (neo.a > 0.994) & (neo.a < 1.006)
    B = (neo.a > 0.9) & (neo.a < 1.1) & (neo.e < 0.1) & (neo.i < 5)
    sel = neo[(A | B) & neo.H.notna()].copy()
    sel["setA"] = A[sel.index]
    sel["setB"] = B[sel.index]
    sel["prov"] = [prov_from_fullname(n) for n in sel.full_name]
    print(f"[ephem] {len(sel)} objects (A: {sel.setA.sum()}, B: {sel.setB.sum()})")
    cache_dir = os.path.join(RAW, "horizons")
    os.makedirs(cache_dir, exist_ok=True)
    arrs, keep = [], []
    t0 = time.time()
    for n, (_, r) in enumerate(sel.iterrows()):
        key = r.pdes.strip().replace(" ", "_")
        fn = os.path.join(cache_dir, key + ".npy")
        if os.path.exists(fn) and not refresh:
            a = np.load(fn)
        else:
            try:
                a = query(r.pdes)
            except Exception as ex:
                print("[ephem] failed", r.pdes, ex)
                a = None
            if a is None:
                continue
            np.save(fn, a)
            time.sleep(0.3)
        arrs.append(a)
        keep.append(r)
        if (n + 1) % 25 == 0:
            print(f"[ephem] {n + 1}/{len(sel)} {time.time() - t0:.0f}s", flush=True)
    nmin = min(len(a) for a in arrs)
    eph = np.stack([a[:nmin] for a in arrs])                      # columns: r, delta, phase, elong, rate, eclat, dec
    np.savez_compressed(raw_npz, eph=eph)
    derive(eph, out_npz)
    kd = pd.DataFrame(keep)[["pdes", "full_name", "prov", "H", "a", "e", "i", "moid", "first_obs", "last_obs",
                              "data_arc", "condition_code", "disc_year", "setA", "setB"]]
    kd.to_csv(out_csv, index=False)
    print(f"[ephem] wrote {out_npz} {eph.shape}")


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