"""Step 6 [revision]: independent data from ESA's NEO Coordination Centre (NEOCC).

(a) NEOCC orbit catalogue (neo_km.cat) and the orbit files (.eq0) of every NEO for which
    ESA's orbit determination solves for a non-gravitational parameter: an independent
    census of fitted area-to-mass ratios and Yarkovsky parameters -> data/neocc_nongrav.csv
(b) NEOCC residual files (.rwo): the post-fit astrometric residuals of ESA's own orbit
    solutions, computed with a complete dynamical model (planetary ephemerides, relativity,
    perturbing asteroids).  A1 is measured by fitting 6 state corrections + A1 to these
    residuals, with partial derivatives from this study's models (numerical model below
    100 m, two-body model above).  Because the residuals no longer depend on this study's
    simplified dynamics, the measurement becomes usable on multi-apparition arcs, where the
    residuals about my own integration are dominated by dynamical-model error.  Two fits are
    stored: A1 alone (nc_*) and A1 with the transverse term A2 solved as a nuisance parameter
    (nc2_*; the primary measurement on multi-apparition arcs), plus A2-only and A3-only fits
    used for validation against JPL's fitted A2.
    -> data/neocc_measurement.csv

Sources: https://neo.ssa.esa.int/PSDB-portlet/download?file=<name>.rwo | .eq0 | neo_km.cat
Raw files are cached in RAW/neocc/.
"""
import os
import re
import sys
import time
from concurrent.futures import ThreadPoolExecutor

import numpy as np
import pandas as pd

import orbfisher as of
import s01_download as dl
import s03_fisher as s3
from common import DATA, RAW, A1_PER_AMR, mjd_from_ymd

BASE = "https://neo.ssa.esa.int/PSDB-portlet/download?file="
CACHE = os.path.join(RAW, "neocc")
N_KEPLER = 900        # random multi-apparition objects > 100 m (two-body partials)
N_SINGLE = 300        # single-apparition cross-check sample
_DATE = re.compile(r"\s(\d{4}) (\d{2}) (\d{2}\.\d+)\s")


def _name(pdes):
    return pdes.replace(" ", "")


def fetch(name, refresh=False):
    """Download (or read from cache) one NEOCC file; returns its text ('' if unavailable)."""
    os.makedirs(CACHE, exist_ok=True)
    path = os.path.join(CACHE, name)
    if os.path.exists(path) and not refresh:
        with open(path, encoding="latin-1") as f:
            return f.read()
    try:
        txt = dl._get(BASE + name, timeout=120, retries=3).decode("latin-1")
    except Exception:      # noqa
        txt = ""
    if "END_OF_HEADER" not in txt:
        txt = ""
    with open(path, "w", encoding="latin-1") as f:
        f.write(txt)
    return txt


def fetch_many(names, refresh=False, workers=3):
    todo = [n for n in names if refresh or not os.path.exists(os.path.join(CACHE, n))]
    if todo:
        print("  NEOCC: downloading %d files ..." % len(todo), flush=True)
        t0 = time.time()
        with ThreadPoolExecutor(max_workers=workers) as ex:
            for k, _ in enumerate(ex.map(lambda n: fetch(n, refresh), todo)):
                if (k + 1) % 500 == 0:
                    print("    %d/%d (%.0f s)" % (k + 1, len(todo), time.time() - t0), flush=True)


# ------------------------------------------------------------ (a) non-gravitational census
def neocc_nongrav(refresh=False):
    cat = fetch("neo_km.cat", refresh)
    rows = []
    for ln in cat.splitlines():
        p = ln.split()
        if len(p) < 11 or ln.startswith("!") or "=" in ln:
            continue
        try:
            rows.append((p[0], float(p[1]), float(p[2]), float(p[3]), float(p[8]), int(p[-1])))
        except ValueError:
            continue
    c = pd.DataFrame(rows, columns=["name", "epoch_mjd", "a", "e", "H", "n_nongrav"])
    ng = c[c.n_nongrav > 0].copy()
    fetch_many([n + ".eq0" for n in ng.name], refresh)
    out = []
    for r in ng.itertuples():
        txt = fetch(r.name + ".eq0")
        lsp = ngr = rms = None
        for ln in txt.splitlines():
            if ln.startswith(" LSP"):
                lsp = ln.split()[1:]
            elif ln.startswith(" NGR"):
                ngr = [float(x) for x in ln.split()[1:]]
            elif ln.startswith("! RMS"):
                rms = [float(x) for x in ln.split()[2:]]
        d = dict(name=r.name, H=r.H, a=r.a, e=r.e, n_nongrav=r.n_nongrav, amr_m2_per_kg=np.nan, amr_sigma=np.nan,
                 yarkovsky_1e10_aud2=np.nan, yarkovsky_sigma=np.nan, solved="")
        if lsp and ngr:
            npar = int(lsp[1]); solved = [int(x) for x in lsp[3:]]
            d["solved"] = ",".join(str(x) for x in solved)
            sig = (rms or [])[6:6 + len(solved)]
            for j, k in enumerate(solved):
                if k == 1 and len(ngr) >= 1:          # area-to-mass ratio, m^2/t
                    d["amr_m2_per_kg"] = ngr[0] * 1e-3
                    d["amr_sigma"] = sig[j] * 1e-3 if j < len(sig) else np.nan
                elif k == 2 and len(ngr) >= npar:     # Yarkovsky parameter, 1e-10 au/d^2
                    d["yarkovsky_1e10_aud2"] = ngr[1]
                    d["yarkovsky_sigma"] = sig[j] if j < len(sig) else np.nan
        out.append(d)
    t = pd.DataFrame(out)
    t.to_csv(os.path.join(DATA, "neocc_nongrav.csv"), index=False)
    print("  NEOCC catalogue: %d NEOs, %d with non-gravitational parameters (%d with a fitted area-to-mass ratio)" %
          (len(c), len(t), int(t.amr_m2_per_kg.notna().sum())))
    return c, t


# ------------------------------------------------------------ (b) residual files
def parse_rwo(txt):
    """Optical astrometry of a NEOCC .rwo file -> DataFrame (ground-based records only) and
    header rms.  Satellite ('S'), radar and roving records are skipped."""
    rows = []
    rms_hdr = np.nan
    n_bad = n_sat = 0
    for ln in txt.splitlines():
        if ln.startswith("RMSast"):
            try:
                rms_hdr = float(ln.split("=")[1])
            except ValueError:
                pass
            continue
        if ln.startswith("!"):
            if "Radar" in ln:
                break
            continue
        m = _DATE.search(ln)
        if not m:
            continue
        pre = ln[:m.start()].split()
        if len(pre) < 3:
            continue
        if pre[1] != "O":
            n_sat += pre[1] == "S" and pre[2] == "S"
            continue
        p = ln[m.end():].split()
        try:
            if p[6] not in ("F", "T") or p[14] not in ("F", "T"):
                raise ValueError
            mjd = float(mjd_from_ymd(int(m.group(1)), int(m.group(2)), float(m.group(3))))
            ra = (float(p[1]) + float(p[2]) / 60.0 + float(p[3]) / 3600.0) * 15.0
            sgn = -1.0 if p[9].startswith("-") else 1.0
            dec = sgn * (abs(float(p[9])) + float(p[10]) / 60.0 + float(p[11]) / 3600.0)
            rows.append((mjd, ra, dec, float(p[5]), float(p[8]), float(p[13]), float(p[16]), p[-4], float(p[-3]),
                         int(p[-2]), pre[2]))
        except (ValueError, IndexError):
            n_bad += 1
    df = pd.DataFrame(rows, columns=["mjd_utc", "ra", "dec", "sig_ra", "res_ra", "sig_dec", "res_dec", "stn", "chi",
                                     "sel", "type"])
    return df, rms_hdr, n_bad, int(n_sat)


def targets(plan, fr):
    m = plan.merge(fr[["pdes", "model", "n_obs", "sig_A1_opt"]], on="pdes", how="left")
    ok = m.selected & m.model.isin(["numeric", "kepler"])
    multi = ok & (m.arc_d > 180)
    a = m[multi & (m.model == "numeric")].assign(why="multi-apparition, numerical partials")
    kp = m[multi & (m.model == "kepler")].sort_values("u")
    u_max = float(kp.u.iloc[min(N_KEPLER, len(kp)) - 1])
    b = kp[kp.u <= u_max].assign(why="multi-apparition, two-body partials (random)")
    c = m[ok & ((m.rms >= 1.4) | (m.n_obs_used < 0.8 * m.n_obs))].assign(why="flagged by fit quality (rms or rejections)")
    s1 = m[ok & (m.arc_d <= 180) & (m.model == "numeric") & (m.sig_A1_opt / A1_PER_AMR <= 1.0 / 6.0) & m.A1.isna()]
    d = s1.sort_values("u").head(N_SINGLE).assign(why="single-apparition cross-check")
    t = pd.concat([a, b, c, d]).drop_duplicates("pdes", keep="first")
    return t, u_max


def measure(row, store, hz, nongrav_names):
    name = _name(row.pdes)
    txt = fetch(name + ".rwo")
    res = dict(pdes=row.pdes, why=row.why, nc_status="ok")
    if not txt:
        res["nc_status"] = "no_file"
        return res
    df, rms_hdr, n_bad, n_sat = parse_rwo(txt)
    res.update(nc_rms_header=rms_hdr, nc_n_optical=len(df), nc_n_unparsed=n_bad, nc_n_satellite=n_sat,
               nc_n_rejected_by_neocc=int((df.sel == 0).sum()) if len(df) else 0,
               nc_nongrav_in_solution=name in nongrav_names)
    known = df.stn.isin(store.oc)
    df = df[known].sort_values("mjd_utc").reset_index(drop=True)
    if len(df) < 5:
        res["nc_status"] = "too_few"
        return res
    par = np.array([store.oc[s_] for s_ in df.stn])
    night = np.floor(df.mjd_utc.values + par[:, 0] / 360.0 - 0.5).astype(np.int64)
    _, grp = np.unique(np.array(["%s%d" % (a, b) for a, b in zip(df.stn, night)]), return_inverse=True)
    obs = dict(mjd_utc=df.mjd_utc.values, ra=df.ra.values, dec=df.dec.values, sig=df.sig_ra.values,
               off=of.observer_offset(df.mjd_utc.values, par[:, 0], par[:, 1], par[:, 2]))
    ng = [0.0 if not np.isfinite(x) else float(x) for x in (row.A1, row.A2, row.A3)]
    try:
        if row.model == "numeric":
            if row.ocat == "short_inarc":
                state0 = of.elements_to_state(row.a, row.e, row.i, row.om, row.w, row.ma); t_ref = float(row.epoch_mjd)
            else:
                state0, jd0 = s3.hz_state(hz, row.pdes, float(np.median(obs["mjd_utc"])) + 2400000.5)
                if state0 is None:
                    res["nc_status"] = "no_state"
                    return res
                t_ref = jd0 - 2400000.5
            out = of.numeric_fisher(state0, t_ref, obs, None, ng=ng)
        else:
            state0 = of.elements_to_state(row.a, row.e, row.i, row.om, row.w, row.ma)
            out = of.kepler_fisher(state0, float(row.epoch_mjd), obs, None)
    except Exception as exc:     # noqa
        res["nc_status"] = "failed:" + type(exc).__name__
        return res
    res["nc_partials"] = row.model
    res["nc_own_res_med"] = float(np.median(np.hypot(out["res_ra"], out["res_dec"])))
    args = (out["Ha"], out["Hd"], df.res_ra.values, df.res_dec.values, df.sig_ra.values)
    m = of.fit_a1(*args, groups=grp, sig_d=df.sig_dec.values)                      # 6 state + A1
    res.update({"nc_" + k[2:]: v for k, v in m.items()})
    # 6 state + A2 + A1: the transverse (Yarkovsky) term solved as a nuisance parameter.  On
    # multi-year arcs an unmodelled Yarkovsky drift otherwise leaks into A1 (see report).
    m2 = of.fit_a1(*args, groups=grp, sig_d=df.sig_dec.values, cols=[0, 1, 2, 3, 4, 5, 7, 6])
    res.update({"nc2_" + k[2:]: v for k, v in m2.items()})
    # single extra parameter A2 (comparison with JPL's fitted A2) and, numerical partials only, A3
    for tag, col in (("a2", 7), ("a3", 8)):
        if col >= out["Ha"].shape[1] or not np.any(out["Ha"][:, col]):
            continue
        mm = of.fit_a1(*args, sig_d=df.sig_dec.values, cols=[0, 1, 2, 3, 4, 5, col])
        if "m_A1" in mm:
            res.update({"nc_%s_val" % tag: mm["m_A1"], "nc_%s_sig" % tag: mm["m_sig"], "nc_%s_rms7" % tag: mm["m_rms7"],
                        "nc_%s_dchi2" % tag: mm["m_dchi2"]})
    # Night-compressed significance: every station-night counts as ONE observation (errors fully
    # correlated within a night), only observations ESA itself kept, no further clipping.  This is
    # a pessimistic bound on the significance, used to judge candidates (not to define samples).
    keep = (df.sel.values == 1)
    if keep.sum() >= 5:
        _, g2, cnt = np.unique(grp[keep], return_inverse=True, return_counts=True)
        f = np.sqrt(np.minimum(cnt[g2], 4.0))          # ESA's weights already carry sqrt(n/4) for n > 4
        a2 = (out["Ha"][keep], out["Hd"][keep], df.res_ra.values[keep], df.res_dec.values[keep], df.sig_ra.values[keep] * f)
        for tag, cols in (("nc", [0, 1, 2, 3, 4, 5, 6]), ("nc2", [0, 1, 2, 3, 4, 5, 7, 6]), ("nc_a3", [0, 1, 2, 3, 4, 5, 8])):
            if max(cols) >= out["Ha"].shape[1] or not np.any(out["Ha"][:, cols[-1]]):
                continue
            mm = of.fit_a1(*a2, sig_d=df.sig_dec.values[keep] * f, cols=cols, clip=1e9)
            if "m_A1" not in mm:
                continue
            dof = 2 * cnt.size - len(cols)
            if dof < 1:
                continue
            chi = mm["m_rms7"] ** 2 * max(2 * int(keep.sum()) - len(cols), 1)
            sg = mm["m_sig"] * max(1.0, np.sqrt(chi / dof))
            res.update({tag + "_night_val": mm["m_A1"], tag + "_night_sig": sg, tag + "_night_z": mm["m_A1"] / sg})
        res["nc_n_nights_kept"] = int(cnt.size)
    res["nc_n_ground"] = len(df)
    return res


def main(refresh=False):
    print("[6] ESA NEOCC: non-gravitational census and residual-based measurement of A1")
    cat, ngt = neocc_nongrav(refresh)
    plan = s3.load_plan()
    fr = pd.read_csv(os.path.join(DATA, "fisher_results.csv"), dtype={"pdes": str})
    tg, u_max = targets(plan, fr)
    print("  targets:", tg.why.value_counts().to_dict(), " (u_max of the random two-body sample: %.4f)" % u_max)
    fetch_many([_name(p) + ".rwo" for p in tg.pdes], refresh)
    store = s3.ObsStore()
    hz = s3.load_hz()
    names_ng = set(ngt.name)
    rows = []
    t0 = time.time()
    for k, row in enumerate(tg.itertuples(index=False)):
        rows.append(measure(row, store, hz, names_ng))
        if (k + 1) % 500 == 0:
            print("    %d/%d (%.0f s)" % (k + 1, len(tg), time.time() - t0), flush=True)
    out = pd.DataFrame(rows)
    out["nc_u_max"] = u_max
    out.to_csv(os.path.join(DATA, "neocc_measurement.csv"), index=False)
    print("  wrote neocc_measurement.csv:", out.nc_status.value_counts().to_dict(), "(%.0f s)" % (time.time() - t0))
    return out


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