"""Step 2: census and classification of every asteroid-designated object with a
fitted non-gravitational parameter in the JPL Small-Body Database.

Output: data/census_nongrav.csv, data/census_summary.csv
"""
import os

import numpy as np
import pandas as pd

from common import (DATA, A1_PER_AMR, P_SRP, AUD2_TO_MS2, PV_NOM, PV_LO, PV_HI, diam_m,
                    a2_yark_max, amr_of_A1, amr_sphere)

# ---------------------------------------------------------------------------
# Literature lists (checked against the papers; see report for references)
# ---------------------------------------------------------------------------
# Seligman et al. 2024, PNAS 121, e2406424121, Tables 1-2 (14 "dark comets";
# includes the 7 of Farnocchia+2023, Seligman+2023 and Chesley+2016)
DARK_COMETS_OUTER = ["2001 ME1", "2005 UY6", "1998 FR11", "2012 UR158", "2003 RM"]
DARK_COMETS_INNER = ["2005 VL1", "2010 RF12", "1998 KY26", "2016 NJ33", "2010 VL65",
                     "2013 BA74", "2006 RH120", "2016 GW221", "2013 XY20"]
PRIMARY_TO_PROV = {"139359": "2001 ME1", "452639": "2005 UY6", "152667": "1998 FR11", "523599": "2003 RM"}
# Objects with documented activity / cometary nature
KNOWN_ACTIVE = {
    "3200": "Phaethon: recurrent perihelion activity (Jewitt & Li 2010; Zhang+2023 sodium)",
    "457175": "2008 GO98 = comet 362P (dual-status active object)",
    "2017 U1": "1I/'Oumuamua: interstellar; Micheli+2018 non-gravitational acceleration",
}
# Small objects whose radial acceleration was published as solar radiation pressure on a natural body
PUBLISHED_SRP = {
    "2009 BD": "Micheli+2012 (AMR 2.97e-4 m2/kg); Mommert+2014",
    "2011 MD": "Micheli+2014; Mommert+2014",
    "2012 LA": "Micheli+2013",
    "2015 TC25": "Farnocchia+2017 (in Reddy+2016 context)",
    "2012 TC4": "Reddy+2019 campaign",
    "2020 CD3": "Fedorets+2020; Naidu+2021",
    "2006 RH120": "Kwiatkowski+2009; Chesley+2016 (also listed as dark comet)",
}
# Artificial objects once carried as asteroids (not in the SBDB asteroid tables any more).
# A1 / AMR from JPL Horizons object pages retrieved 2026-10-12.
ARTIFICIAL = [
    dict(name="2020 SO (Surveyor 2 Centaur, 1966-084B)", H=28.44, A1=3.5566674695484e-9,
         src="JPL Horizons solution JPL#90 (D. Farnocchia), A1 with (1au/r)^2"),
    dict(name="J002E3 (Apollo 12 S-IVB)", H=26.43, A1=0.0079 * A1_PER_AMR,
         src="JPL Horizons: AMRAT = 0.0079 m2/kg"),
]

RHO_MIN_YARK = 1000.0   # kg/m3, generous low density for the Yarkovsky bound
RHO_MIN_SRP = 500.0     # kg/m3, lowest plausible bulk density of a natural monolith/rubble pile
CR_MAX = 1.5            # radiation-pressure coefficient upper value (1 = absorption only)
NSIG = 3.0


def classify(df):
    df = df.copy()
    for c in ["H", "A1", "A1_sigma", "A2", "A2_sigma", "A3", "A3_sigma", "diameter", "albedo", "a", "e", "i"]:
        df[c] = pd.to_numeric(df[c], errors="coerce")
    d_meas = df.diameter * 1e3
    df["D_nom"] = np.where(d_meas > 0, d_meas, diam_m(df.H, PV_NOM))
    df["D_min"] = np.where(d_meas > 0, d_meas / 1.3, diam_m(df.H, PV_HI))
    df["D_max"] = np.where(d_meas > 0, d_meas * 1.3, diam_m(df.H, PV_LO))
    for k in ("A1", "A2", "A3"):
        df["snr_" + k] = (df[k] / df[k + "_sigma"]).abs()
    df["a_tot"] = np.sqrt(df.A1.fillna(0) ** 2 + df.A2.fillna(0) ** 2 + df.A3.fillna(0) ** 2)
    # bounds
    df["A2_yark_max"] = a2_yark_max(df.D_min, RHO_MIN_YARK)
    df["A2_yark_max_nom"] = a2_yark_max(df.D_nom, 2000.0)
    df["A1_srp_max"] = CR_MAX * amr_sphere(df.D_min, RHO_MIN_SRP) * A1_PER_AMR
    df["A1_srp_nom"] = amr_sphere(df.D_nom, 2000.0) * A1_PER_AMR
    df["amr_equiv"] = amr_of_A1(df.A1)                    # m2/kg for C_R = 1 (signed)
    df["amr_equiv_sigma"] = amr_of_A1(df.A1_sigma)
    df["amr_total_equiv"] = amr_of_A1(df.a_tot)           # |a| expressed as area-to-mass
    df["amr_rock_nom"] = amr_sphere(df.D_nom, 2000.0)
    df["rho_implied"] = 1.5 * A1_PER_AMR / (df.A1 * df.D_nom)   # density if A1 is pure SRP, C_R=1
    df["rho_implied_hi"] = CR_MAX * 1.5 * A1_PER_AMR / ((df.A1 - df.A1_sigma) * df.D_min)

    sig1 = df.snr_A1 >= NSIG
    sig2 = df.snr_A2 >= NSIG
    sig3 = df.snr_A3 >= NSIG
    df["significant"] = sig1 | sig2 | sig3
    a1_neg = sig1 & (df.A1 < 0)
    a1_excess = sig1 & (df.A1 > 0) & ((df.A1 - NSIG * df.A1_sigma) > df.A1_srp_max)
    a2_excess = sig2 & ((df.A2.abs() - NSIG * df.A2_sigma) > df.A2_yark_max)
    df["flag_A1_sunward"] = a1_neg
    df["flag_A1_excess"] = a1_excess
    df["flag_A2_excess"] = a2_excess
    df["flag_A3"] = sig3
    nonrad = a1_neg | a1_excess | a2_excess | sig3
    key = df.pdes.map(lambda p: PRIMARY_TO_PROV.get(p, p))
    is_dc = key.isin(DARK_COMETS_OUTER + DARK_COMETS_INNER)
    is_act = df.pdes.isin(KNOWN_ACTIVE)
    cls = np.full(len(df), "not significant (<3 sigma)", dtype=object)
    cls[(sig2 & ~nonrad).values] = "Yarkovsky-consistent"
    cls[(sig1 & (df.A1 > 0) & ~nonrad).values] = "radiation pressure on natural body"
    cls[(sig1 & (df.A1 > 0) & sig2 & ~nonrad).values] = "radiation pressure + Yarkovsky"
    # non-radiative term that is not in the published dark-comet / active lists (the label is
    # refined after the literature check in s04_limits.refine_census)
    cls[(nonrad & df.significant).values] = "non-radiative, not in published lists"
    cls[(is_dc & df.significant).values] = "published dark comet"
    cls[(is_act & df.significant).values] = "known active / interstellar"
    df["class_ng"] = cls
    # [revision] 'Yarkovsky-consistent' only means 'not excluded by the generous bound' (1000 kg/m3,
    # albedo 0.45).  Flag the objects that exceed the *nominal* maximum (2000 kg/m3, albedo 0.14).
    df["A2_over_nominal"] = df.A2.abs() / df.A2_yark_max_nom
    df["flag_A2_over_nominal"] = sig2 & (df.A2.abs() > df.A2_yark_max_nom)
    df["flag_A2_over_nominal_3sig"] = sig2 & ((df.A2.abs() - NSIG * df.A2_sigma) > df.A2_yark_max_nom)
    # Tisserand parameter with respect to Jupiter
    aj = 5.2044
    df["T_jup"] = aj / df.a + 2.0 * np.cos(np.radians(df.i)) * np.sqrt(df.a / aj * (1.0 - df.e ** 2))
    df["lit_note"] = [KNOWN_ACTIVE.get(p, PUBLISHED_SRP.get(p, "")) for p in df.pdes]
    return df


def main():
    print("[2] census")
    ng = pd.read_csv(os.path.join(DATA, "sbdb_nongrav_all.csv"), dtype={"pdes": str}, low_memory=False)
    ast = ng[ng.kind.str.startswith("a")].copy()
    cen = classify(ast)
    cols = ["full_name", "pdes", "kind", "neo", "class", "H", "D_nom", "D_min", "D_max", "q", "e", "a", "i",
            "data_arc", "n_obs_used", "n_del_obs_used", "condition_code", "producer", "soln_date", "rms", "T_jup",
            "A1", "A1_sigma", "snr_A1", "A2", "A2_sigma", "snr_A2", "A3", "A3_sigma", "snr_A3",
            "amr_equiv", "amr_equiv_sigma", "amr_total_equiv", "amr_rock_nom", "rho_implied",
            "A1_srp_max", "A2_yark_max", "A2_yark_max_nom",
            "flag_A1_sunward", "flag_A1_excess", "flag_A2_excess", "flag_A3", "significant",
            "A2_over_nominal", "flag_A2_over_nominal", "flag_A2_over_nominal_3sig",
            "class_ng", "lit_note"]
    cen[cols].to_csv(os.path.join(DATA, "census_nongrav.csv"), index=False)
    summ = cen.class_ng.value_counts().rename_axis("class").reset_index(name="n")
    summ.to_csv(os.path.join(DATA, "census_summary.csv"), index=False)
    print("  asteroid-designated objects with fitted non-grav parameters:", len(cen),
          "(A1: %d, A2: %d, A3: %d)" % (cen.A1.notna().sum(), cen.A2.notna().sum(), cen.A3.notna().sum()))
    print(summ.to_string(index=False))
    return cen


if __name__ == "__main__":
    cen = main()
    pd.set_option("display.width", 250); pd.set_option("display.max_rows", 200)
    x = cen[~cen.class_ng.isin(["Yarkovsky-consistent", "not significant (<3 sigma)"])]
    print(x[["full_name", "H", "D_nom", "A1", "snr_A1", "A2", "snr_A2", "A3", "snr_A3", "amr_equiv",
             "rho_implied", "flag_A1_sunward", "flag_A1_excess", "flag_A2_excess", "flag_A3", "class_ng"]]
          .sort_values("class_ng").to_string())