"""Near-Earth-object survey completeness as of the download date.

Known near-Earth asteroids (JPL SBDB, cross-checked against MPC NEA.txt) per
absolute-magnitude bin, divided by two published debiased population models
(NEOMOD2 = Nesvorny et al. 2024; HC21 = Harris & Chodas 2021, as tabulated in
NEOMOD2 Table 3).  Also the same split by discovery epoch (used to calibrate the
survey simulator in s02) and the Earth-like-orbit subset.

Outputs: data/neo_completeness.csv, data/neo_known_by_epoch.csv
"""
import json
import os
import re

import numpy as np
import pandas as pd

from common import DATA, RAW, D_from_H

EPOCH_ENDS = [2005.0, 2011.0, 2016.0, 2021.0, 2027.0]


def load_neo():
    d = json.load(open(os.path.join(RAW, "sbdb_neo.json")))
    df = pd.DataFrame(d["data"], columns=d["fields"])
    for c in ["H", "a", "e", "i", "q", "ad", "moid", "diameter", "albedo", "data_arc", "om", "w", "ma"]:
        df[c] = pd.to_numeric(df[c], errors="coerce")
    df = df[df.kind.str.startswith("a")].copy()        # asteroids only (no comets)

    def disc_year(name, first_obs):
        m = re.search(r"\((\d{4}) [A-Z]{2}\d*\)", name)
        if m:
            return float(m.group(1))
        m = re.search(r"\(A(\d{3}) [A-Z]{2}\)", name)
        if m:
            return 1000.0 + float(m.group(1))
        if "P-L" in name:
            return 1960.0
        m = re.search(r"T-(\d)", name)
        if m:
            return 1971.0
        try:
            return float(str(first_obs)[:4])
        except Exception:
            return np.nan
    df["disc_year"] = [disc_year(n, f) for n, f in zip(df.full_name, df.first_obs)]
    return df


def load_models():
    return pd.read_csv(os.path.join(DATA, "neo_population_models.csv"), comment="#")


def main():
    neo = load_neo()
    n_mpc = sum(1 for _ in open(os.path.join(RAW, "NEA.txt")))
    print(f"[neo] JPL SBDB near-Earth asteroids: {len(neo)} (with H: {neo.H.notna().sum()}); MPC NEA.txt rows: {n_mpc}")
    mod = load_models()
    mod["dN_hc21"] = np.diff(np.concatenate([[mod.N_hc21.iloc[0] - mod.dN_neomod2.iloc[0] * mod.N_hc21.iloc[0] / mod.N_neomod2.iloc[0]],
                                             mod.N_hc21.values]))
    rows = []
    for _, r in mod.iterrows():
        sel = (neo.H >= r.H1) & (neo.H < r.H2)
        k = int(sel.sum())
        kc = int((neo.H < r.H2).sum())
        row = dict(H1=r.H1, H2=r.H2, D_m_p020=float(D_from_H((r.H1 + r.H2) / 2, 0.20)),
                   D_m_p005=float(D_from_H((r.H1 + r.H2) / 2, 0.05)),
                   known_in_bin=k, dN_neomod2=r.dN_neomod2, dN_hc21=r.dN_hc21,
                   compl_diff_neomod2=k / r.dN_neomod2, compl_diff_hc21=k / r.dN_hc21,
                   known_cum=kc, N_neomod2=r.N_neomod2, N_hc21=r.N_hc21,
                   compl_cum_neomod2=kc / r.N_neomod2, compl_cum_hc21=kc / r.N_hc21,
                   known_cum_2022=r.N_mpc_2022)
        for ye in EPOCH_ENDS:
            row[f"known_in_bin_lt{int(ye)}"] = int((sel & (neo.disc_year < ye)).sum())
        rows.append(row)
    out = pd.DataFrame(rows)
    out.to_csv(os.path.join(DATA, "neo_completeness.csv"), index=False)
    pd.set_option("display.width", 200)
    print(out[["H1", "H2", "D_m_p020", "known_in_bin", "dN_neomod2", "compl_diff_neomod2", "compl_diff_hc21",
               "known_cum", "compl_cum_neomod2", "compl_cum_hc21"]].round(4).to_string(index=False))
    faint = neo[neo.H >= 27.75]
    print(f"[neo] known with H>=27.75: {len(faint)}; H>=30: {(neo.H >= 30).sum()}; H>=32: {(neo.H >= 32).sum()}; faintest H = {neo.H.max():.2f}")
    # Earth-like orbits (the region a parked artefact or a human rocket stage would occupy)
    el = neo[(neo.a > 0.9) & (neo.a < 1.1) & (neo.e < 0.1) & (neo.i < 5)]
    co = neo[(neo.a > 0.994) & (neo.a < 1.006)]
    print(f"[neo] Earth-like (0.9<a<1.1, e<0.1, i<5 deg): {len(el)} known; median H {el.H.median():.1f}; "
          f"brightest H {el.H.min():.1f}")
    print(f"[neo] co-orbital zone |a-1|<0.006: {len(co)} known; with H<22: {(co.H < 22).sum()}; H<25: {(co.H < 25).sum()}")
    el.sort_values("H")[["full_name", "H", "a", "e", "i", "moid", "first_obs", "data_arc"]].head(40).to_csv(
        os.path.join(DATA, "neo_earthlike_brightest.csv"), index=False)
    return out


if __name__ == "__main__":
    main()