"""Real astronomical objects with sourced values, for naming things in generated questions.

Values are rounded published figures. Planet/moon data: NASA Planetary Fact Sheet; exoplanets:
NASA Exoplanet Archive (minimum masses where only radial-velocity data exist); stars: SIMBAD and
the cited literature summaries; black holes: published mass estimates (EHT 2019/2022 for M87* and
Sgr A*); galaxies and quasars: NED.
"""

from __future__ import annotations

import random
from dataclasses import dataclass, field

from zenith.engines.constants import AU

SOLAR = "NASA Planetary Fact Sheet"
EXO = "NASA Exoplanet Archive"
STAR = "SIMBAD / published stellar parameters"
BH = "published mass estimates"
NED = "NASA/IPAC Extragalactic Database (NED)"
R_EARTH_KM = 6371.0
R_JUP_KM = 69911.0
M_JUP_MEARTH = 317.8


@dataclass(frozen=True)
class Obj:
    id: str
    name: str
    aliases: tuple[str, ...]
    kind: str
    props: dict[str, float] = field(hash=False)
    source: str


# Which engine parameter each property can fill, per kind: prop -> ((sim, op, param), ...).
_PLANETARY = {
    "a_au": (("orbit", "kepler", "a_au"),),
    "e": (("orbit", "kepler", "e"),),
    "mass_mearth": (("orbit", "kepler", "planet_mass_mearth"), ("gravity", "escape", "mass_mearth")),
    "radius_km": (("gravity", "escape", "radius_km"),),
    "host_mass_msun": (("orbit", "kepler", "star_mass_msun"),),
}
KIND_PARAMS: dict[str, dict[str, tuple[tuple[str, str, str], ...]]] = {
    "planet": _PLANETARY,
    "dwarf_planet": _PLANETARY,
    "exoplanet": _PLANETARY,
    "moon": {"mass_mearth": (("gravity", "escape", "mass_mearth"),), "radius_km": (("gravity", "escape", "radius_km"),)},
    "star": {
        "mass_msun": (("stellar", "main_sequence", "mass_msun"), ("orbit", "kepler", "star_mass_msun"),
                      ("gravity", "black_hole", "mass_msun")),
        "temp_k": (("stellar", "blackbody", "temp_k"),),
        "radius_rsun": (("stellar", "blackbody", "radius_rsun"),),
    },
    "black_hole": {"mass_msun": (("gravity", "black_hole", "mass_msun"),)},
    "galaxy": {"z": (("cosmology", "hubble", "z"),), "distance_mpc": (("cosmology", "hubble", "distance_mpc"),)},
    "quasar": {"z": (("cosmology", "hubble", "z"),)},
    "satellite": {"altitude_km": (("orbit", "satellite", "altitude_km"),)},
    "country": {"lat": (("darksky", "darkest", "lat"),), "lon": (("darksky", "darkest", "lon"),),
                "radius_km": (("darksky", "darkest", "radius_km"),)},
    "place": {"lat": (("darksky", "site", "lat"),), "lon": (("darksky", "site", "lon"),)},
}


def _planetary(kind, source, rows, host=1.0):
    out = []
    for row in rows:
        id_, name, a, e, mass, radius, *rest = row
        h = rest[0] if rest else host
        props = {"a_au": a, "e": e, "host_mass_msun": h}
        if mass is not None:
            props["mass_mearth"] = mass
        if radius is not None:
            props["radius_km"] = radius
        out.append(Obj(id_, name, (), kind, props, source))
    return out


# id, name, a [AU], e, mass [M⊕], radius [km]
_PLANETS = [
    ("mercury", "Mercury", 0.387, 0.2056, 0.0553, 2439.7),
    ("venus", "Venus", 0.723, 0.0068, 0.815, 6051.8),
    ("earth", "Earth", 1.000, 0.0167, 1.0, 6371.0),
    ("mars", "Mars", 1.524, 0.0934, 0.107, 3389.5),
    ("jupiter", "Jupiter", 5.203, 0.0484, 317.8, 69911.0),
    ("saturn", "Saturn", 9.537, 0.0539, 95.2, 58232.0),
    ("uranus", "Uranus", 19.19, 0.0473, 14.5, 25362.0),
    ("neptune", "Neptune", 30.07, 0.0086, 17.1, 24622.0),
]
_DWARFS = [
    ("pluto", "Pluto", 39.48, 0.2488, 0.00218, 1188.3),
    ("eris", "Eris", 67.86, 0.4361, 0.0028, 1163.0),
    ("ceres", "Ceres", 2.767, 0.0758, 0.000157, 469.7),
    ("haumea", "Haumea", 43.13, 0.1913, 0.000671, 780.0),
    ("makemake", "Makemake", 45.79, 0.159, 0.00052, 715.0),
]
# id, name, mass [M⊕], radius [km]
_MOONS = [
    ("moon", "the Moon", 0.0123, 1737.4), ("io", "Io", 0.0150, 1821.6), ("europa", "Europa", 0.0080, 1560.8),
    ("ganymede", "Ganymede", 0.0248, 2634.1), ("callisto", "Callisto", 0.0180, 2410.3),
    ("titan", "Titan", 0.0225, 2574.7), ("enceladus", "Enceladus", 1.8e-5, 252.1),
    ("triton", "Triton", 0.00359, 1353.4), ("titania", "Titania", 5.9e-4, 788.4),
    ("oberon", "Oberon", 5.1e-4, 761.4), ("rhea", "Rhea", 3.86e-4, 763.8), ("iapetus", "Iapetus", 3.03e-4, 734.5),
    ("dione", "Dione", 1.84e-4, 561.4), ("tethys", "Tethys", 1.03e-4, 531.1), ("charon", "Charon", 2.65e-4, 606.0),
    ("ariel", "Ariel", 2.1e-4, 578.9), ("umbriel", "Umbriel", 2.0e-4, 584.7), ("miranda", "Miranda", 1.1e-5, 235.8),
]
# id, name, a [AU], e, mass [M⊕] or None, radius [km] or None, host mass [M☉]
_EXOPLANETS = [
    ("proxima_b", "Proxima Centauri b", 0.0485, 0.02, 1.07, None, 0.122),
    ("trappist_1b", "TRAPPIST-1b", 0.01154, 0.006, 1.374, 1.116 * R_EARTH_KM, 0.0898),
    ("trappist_1c", "TRAPPIST-1c", 0.01580, 0.007, 1.308, 1.097 * R_EARTH_KM, 0.0898),
    ("trappist_1d", "TRAPPIST-1d", 0.02227, 0.008, 0.388, 0.788 * R_EARTH_KM, 0.0898),
    ("trappist_1e", "TRAPPIST-1e", 0.02925, 0.005, 0.692, 0.920 * R_EARTH_KM, 0.0898),
    ("trappist_1f", "TRAPPIST-1f", 0.03849, 0.010, 1.039, 1.045 * R_EARTH_KM, 0.0898),
    ("trappist_1g", "TRAPPIST-1g", 0.04683, 0.002, 1.321, 1.129 * R_EARTH_KM, 0.0898),
    ("trappist_1h", "TRAPPIST-1h", 0.06189, 0.006, 0.326, 0.755 * R_EARTH_KM, 0.0898),
    ("51_peg_b", "51 Pegasi b", 0.0527, 0.013, 146.0, None, 1.11),
    ("hd_209458_b", "HD 209458 b", 0.04747, 0.0, 219.0, 1.38 * R_JUP_KM, 1.15),
    ("kepler_22b", "Kepler-22b", 0.849, 0.0, None, 2.4 * R_EARTH_KM, 0.97),
    ("kepler_452b", "Kepler-452b", 1.046, 0.0, None, 1.63 * R_EARTH_KM, 1.04),
    ("kepler_186f", "Kepler-186f", 0.432, 0.04, None, 1.17 * R_EARTH_KM, 0.54),
    ("toi_700_d", "TOI-700 d", 0.163, 0.032, None, 1.19 * R_EARTH_KM, 0.415),
    ("lhs_1140_b", "LHS 1140 b", 0.0957, 0.0, 5.6, 1.73 * R_EARTH_KM, 0.18),
    ("55_cnc_e", "55 Cancri e", 0.01544, 0.05, 7.99, 1.875 * R_EARTH_KM, 0.905),
    ("gj_1214_b", "GJ 1214 b", 0.01490, 0.0, 8.17, 2.74 * R_EARTH_KM, 0.181),
    ("hd_189733_b", "HD 189733 b", 0.0313, 0.0, 359.0, 1.13 * R_JUP_KM, 0.81),
    ("wasp_12b", "WASP-12b", 0.0234, 0.0, 467.0, 1.9 * R_JUP_KM, 1.43),
    ("kepler_10b", "Kepler-10b", 0.01685, 0.0, 3.26, 1.47 * R_EARTH_KM, 0.91),
    ("gliese_581_c", "Gliese 581 c", 0.0721, 0.0, 5.5, None, 0.31),
    ("gliese_667_cc", "Gliese 667 Cc", 0.125, 0.02, 3.8, None, 0.33),
    ("ross_128_b", "Ross 128 b", 0.0496, 0.116, 1.40, None, 0.168),
    ("teegarden_b", "Teegarden's Star b", 0.0252, 0.0, 1.05, None, 0.097),
    ("barnard_b", "Barnard's Star b", 0.0229, 0.0, 0.37, None, 0.162),
    ("tau_ceti_e", "Tau Ceti e", 0.538, 0.18, 3.93, None, 0.783),
    ("eps_eri_b", "Epsilon Eridani b", 3.48, 0.07, 0.66 * M_JUP_MEARTH, None, 0.82),
    ("hr_8799_b", "HR 8799 b", 68.0, 0.0, 7 * M_JUP_MEARTH, None, 1.47),
    ("beta_pic_b", "Beta Pictoris b", 10.26, 0.10, 9.3 * M_JUP_MEARTH, None, 1.75),
    ("kepler_90h", "Kepler-90h", 1.01, 0.01, None, 1.01 * R_JUP_KM, 1.2),
    ("k2_18_b", "K2-18 b", 0.1429, 0.09, 8.63, 2.61 * R_EARTH_KM, 0.495),
    ("wasp_121b", "WASP-121b", 0.02544, 0.0, 369.0, 1.75 * R_JUP_KM, 1.36),
    ("kelt_9b", "KELT-9b", 0.03462, 0.0, 915.0, 1.89 * R_JUP_KM, 2.52),
    ("hd_80606_b", "HD 80606 b", 0.449, 0.933, 1303.0, None, 0.97),
    ("kepler_1649c", "Kepler-1649c", 0.0649, 0.0, None, 1.06 * R_EARTH_KM, 0.198),
    ("ltt_1445_ab", "LTT 1445 Ab", 0.0381, 0.0, 2.87, 1.30 * R_EARTH_KM, 0.257),
    ("corot_7b", "CoRoT-7b", 0.0172, 0.0, 5.74, 1.58 * R_EARTH_KM, 0.91),
]
# id, name, mass [M☉], Teff [K], radius [R☉], distance [ly]
_STARS = [
    ("sun", "the Sun", 1.0, 5772, 1.0, AU / 9.4607304725808e15),
    ("proxima_centauri", "Proxima Centauri", 0.122, 3042, 0.154, 4.246),
    ("alpha_cen_a", "Alpha Centauri A", 1.079, 5790, 1.2175, 4.37),
    ("alpha_cen_b", "Alpha Centauri B", 0.909, 5260, 0.8591, 4.37),
    ("barnards_star", "Barnard's Star", 0.162, 3223, 0.187, 5.96),
    ("sirius_a", "Sirius A", 2.063, 9940, 1.711, 8.6),
    ("sirius_b", "Sirius B", 1.018, 25000, 0.0084, 8.6),
    ("vega", "Vega", 2.135, 9602, 2.36, 25.0),
    ("betelgeuse", "Betelgeuse", 16.5, 3600, 764.0, 548.0),
    ("rigel", "Rigel", 21.0, 12100, 78.9, 860.0),
    ("polaris", "Polaris", 5.4, 6015, 37.5, 433.0),
    ("arcturus", "Arcturus", 1.08, 4286, 25.4, 36.7),
    ("aldebaran", "Aldebaran", 1.16, 3910, 44.2, 65.3),
    ("antares", "Antares", 12.0, 3660, 680.0, 550.0),
    ("capella", "Capella Aa", 2.57, 4970, 11.98, 42.9),
    ("procyon", "Procyon A", 1.499, 6530, 2.048, 11.46),
    ("altair", "Altair", 1.86, 7700, 1.79, 16.73),
    ("deneb", "Deneb", 19.0, 8525, 203.0, 2600.0),
    ("spica", "Spica", 11.43, 22300, 7.47, 250.0),
    ("canopus", "Canopus", 8.0, 7400, 71.0, 310.0),
    ("regulus", "Regulus", 3.8, 12460, 4.35, 79.3),
    ("fomalhaut", "Fomalhaut", 1.92, 8590, 1.84, 25.1),
    ("pollux", "Pollux", 1.91, 4586, 9.06, 33.8),
    ("castor", "Castor A", 2.37, 10286, 2.4, 51.0),
    ("achernar", "Achernar", 6.0, 15000, 7.3, 139.0),
    ("bellatrix", "Bellatrix", 7.7, 21800, 5.75, 250.0),
    ("mira", "Mira", 1.18, 3000, 332.0, 300.0),
    ("tau_ceti", "Tau Ceti", 0.783, 5344, 0.793, 11.9),
    ("epsilon_eridani", "Epsilon Eridani", 0.82, 5084, 0.735, 10.5),
    ("wolf_359", "Wolf 359", 0.11, 2749, 0.144, 7.86),
    ("luytens_star", "Luyten's Star", 0.29, 3382, 0.315, 12.35),
    ("trappist_1", "TRAPPIST-1", 0.0898, 2566, 0.1192, 40.7),
    ("kepler_186", "Kepler-186", 0.544, 3755, 0.523, 579.0),
    ("ross_128", "Ross 128", 0.168, 3192, 0.197, 11.0),
    ("teegardens_star", "Teegarden's Star", 0.097, 2904, 0.107, 12.5),
    ("51_pegasi", "51 Pegasi", 1.11, 5768, 1.152, 50.6),
    ("hd_209458", "HD 209458", 1.15, 6065, 1.20, 157.0),
    ("eta_carinae", "Eta Carinae A", 100.0, 35000, 240.0, 7500.0),
    ("vy_cma", "VY Canis Majoris", 17.0, 3490, 1420.0, 3820.0),
    ("gliese_581", "Gliese 581", 0.31, 3480, 0.30, 20.5),
    ("kepler_452", "Kepler-452", 1.04, 5757, 1.11, 1800.0),
    ("kepler_22", "Kepler-22", 0.97, 5518, 0.98, 640.0),
    ("hr_8799", "HR 8799", 1.47, 7430, 1.34, 133.0),
    ("beta_pictoris", "Beta Pictoris", 1.75, 8052, 1.5, 63.4),
    ("lhs_1140", "LHS 1140", 0.18, 3216, 0.21, 48.9),
    ("k2_18", "K2-18", 0.495, 3457, 0.469, 124.0),
    ("toi_700", "TOI-700", 0.415, 3480, 0.42, 101.4),
    ("eps_indi_a", "Epsilon Indi A", 0.76, 4649, 0.71, 11.87),
    ("61_cygni_a", "61 Cygni A", 0.70, 4526, 0.665, 11.4),
    ("alnitak", "Alnitak", 33.0, 29000, 20.0, 1260.0),
]
# id, name, mass [M☉], distance [ly]
_BLACK_HOLES = [
    ("sgr_a_star", "Sagittarius A*", 4.3e6, 26670.0),
    ("m87_star", "M87*", 6.5e9, 5.35e7),
    ("cygnus_x1", "Cygnus X-1", 21.2, 7240.0),
    ("ton_618", "the black hole in TON 618", 6.6e10, 1.8e10),
    ("gaia_bh1", "Gaia BH1", 9.62, 1560.0),
    ("gaia_bh3", "Gaia BH3", 33.0, 1926.0),
    ("v404_cygni", "V404 Cygni", 9.0, 7800.0),
    ("gw150914", "the GW150914 merger remnant", 62.0, 1.3e9),
    ("a0620_00", "A0620-00", 6.6, 3300.0),
    ("grs_1915", "GRS 1915+105", 12.4, 28000.0),
    ("m31_star", "Andromeda's central black hole", 1.4e8, 2.5e6),
    ("oj_287_bh", "the primary black hole of OJ 287", 1.8e10, 3.5e9),
]
# id, name, {"z": ...} or {"distance_mpc": ...}
_GALAXIES = [
    ("andromeda", "the Andromeda Galaxy", {"distance_mpc": 0.765}),
    ("triangulum", "the Triangulum Galaxy", {"distance_mpc": 0.86}),
    ("lmc", "the Large Magellanic Cloud", {"distance_mpc": 0.050}),
    ("smc", "the Small Magellanic Cloud", {"distance_mpc": 0.062}),
    ("whirlpool", "the Whirlpool Galaxy", {"distance_mpc": 8.58}),
    ("sombrero", "the Sombrero Galaxy", {"distance_mpc": 9.55}),
    ("pinwheel", "the Pinwheel Galaxy", {"distance_mpc": 6.4}),
    ("centaurus_a", "Centaurus A", {"distance_mpc": 3.8}),
    ("m87", "Messier 87", {"distance_mpc": 16.4}),
    ("m82", "the Cigar Galaxy", {"distance_mpc": 3.5}),
    ("m81", "Bode's Galaxy", {"distance_mpc": 3.6}),
    ("ngc_1300", "NGC 1300", {"distance_mpc": 18.7}),
    ("ngc_4993", "NGC 4993", {"distance_mpc": 40.0}),
    ("ngc_6822", "Barnard's Galaxy", {"distance_mpc": 0.50}),
    ("cartwheel", "the Cartwheel Galaxy", {"z": 0.030}),
    ("hoags_object", "Hoag's Object", {"z": 0.0426}),
    ("gn_z11", "GN-z11", {"z": 10.6}),
    ("jades_gs_z14_0", "JADES-GS-z14-0", {"z": 14.32}),
    ("macs0647_jd", "MACS0647-JD", {"z": 10.17}),
    ("egs_zs8_1", "EGS-zs8-1", {"z": 7.73}),
]
# id, name, z
_QUASARS = [
    ("3c_273", "3C 273", 0.158), ("3c_48", "3C 48", 0.367), ("ton_618_qso", "TON 618", 2.219),
    ("ulas_j1342", "ULAS J1342+0928", 7.54), ("j0313_1806", "J0313-1806", 7.64),
    ("apm_08279", "APM 08279+5255", 3.91), ("twin_quasar", "the Twin Quasar (QSO 0957+561)", 1.41),
    ("3c_279", "3C 279", 0.536), ("oj_287", "OJ 287", 0.306), ("cloverleaf", "the Cloverleaf Quasar", 2.56),
]


# id, name, typical altitude [km] (published orbital elements, rounded)
_SATELLITES = [
    ("iss", "the International Space Station", 420.0), ("hubble", "the Hubble Space Telescope", 535.0),
    ("tiangong", "the Tiangong space station", 390.0), ("starlink", "a Starlink satellite", 550.0),
    ("gps", "a GPS satellite", 20200.0), ("geostationary", "a geostationary satellite", 35786.0),
    ("galileo", "a Galileo navigation satellite", 23222.0), ("noaa", "a NOAA weather satellite", 850.0),
]
# id, name, reference point (lat, lon) near the country's centre, radius [km] covering most of it
_COUNTRIES = [
    ("zimbabwe", "Zimbabwe", -19.0, 29.15, 350.0), ("botswana", "Botswana", -22.3, 24.7, 450.0),
    ("zambia", "Zambia", -13.5, 27.8, 500.0), ("mozambique", "Mozambique", -18.0, 35.0, 500.0),
    ("south_africa", "South Africa", -29.0, 25.0, 500.0), ("namibia", "Namibia", -22.6, 17.1, 500.0),
    ("malawi", "Malawi", -13.25, 34.3, 300.0), ("lesotho", "Lesotho", -29.6, 28.2, 120.0),
    ("eswatini", "Eswatini", -26.5, 31.5, 80.0),
]
SAT = "published orbital elements"
GEO = "approximate geographic centre"


def _build() -> tuple[Obj, ...]:
    objs = (_planetary("planet", SOLAR, _PLANETS) + _planetary("dwarf_planet", SOLAR, _DWARFS)
            + _planetary("exoplanet", EXO, _EXOPLANETS))
    objs += [Obj(i, n, (), "moon", {"mass_mearth": m, "radius_km": r}, SOLAR) for i, n, m, r in _MOONS]
    objs += [Obj(i, n, (), "star", {"mass_msun": m, "temp_k": float(t), "radius_rsun": r, "distance_ly": d}, STAR)
             for i, n, m, t, r, d in _STARS]
    objs += [Obj(i, n, (), "black_hole", {"mass_msun": m, "distance_ly": d}, BH) for i, n, m, d in _BLACK_HOLES]
    objs += [Obj(i, n, (), "galaxy", props, NED) for i, n, props in _GALAXIES]
    objs += [Obj(i, n, (), "quasar", {"z": z}, NED) for i, n, z in _QUASARS]
    objs += [Obj(i, n, (), "satellite", {"altitude_km": h}, SAT) for i, n, h in _SATELLITES]
    objs += [Obj(i, n, (), "country", {"lat": la, "lon": lo, "radius_km": r}, GEO) for i, n, la, lo, r in _COUNTRIES]
    return tuple(objs)


OBJECTS: tuple[Obj, ...] = _build()


def all_objects() -> list[Obj]:
    """The fixed table plus the dark-sky map's places (read from the installed place list)."""
    from zenith.sft.places import place_objects

    return [*OBJECTS, *place_objects()]


def objects_by_kind(kind: str) -> list[Obj]:
    return [o for o in all_objects() if o.kind == kind]


ALWAYS_TRAIN = {"zimbabwe"}  # the headline "around Zimbabwe" question must be trained, never only evaluated


def split_objects(seed: int = 0, held_out: float = 0.15) -> tuple[list[Obj], list[Obj]]:
    """Deterministic (train, held-out) split with ~15% of each kind held out (at least one)."""
    objs = all_objects()
    held: set[str] = set()
    for kind in sorted({o.kind for o in objs}):
        ids = sorted(o.id for o in objs if o.kind == kind and o.id not in ALWAYS_TRAIN)
        random.Random(f"{seed}:{kind}").shuffle(ids)
        held.update(ids[: max(1, round(held_out * len(ids)))])
    return [o for o in objs if o.id not in held], [o for o in objs if o.id in held]