"""Empirical completeness limits of the minor-planet catalogue from the main belt outward.

Reads the full MPC orbit catalogue (MPCORB.DAT.gz), splits it into dynamical
regions by semimajor axis and, for each region, estimates the absolute
magnitude to which the catalogue is complete from the shape of the differential
H distribution: a power law is fitted to the bright side (1.0 to 3.0 mag brighter
than the mode) and H_90 / H_50 are the magnitudes at which the observed counts
fall to 90% / 50% of the extrapolated power law (in the spirit of Hendler &
Malhotra 2020).  The limits are converted to diameters for albedo 0.2, 0.05 and
0.01, and compared with a purely photometric estimate (V = 21.5 at opposition).

Outputs: data/belt_completeness.csv, data/belt_H_histograms.csv
"""
import gzip
import os

import numpy as np
import pandas as pd

from common import DATA, RAW, D_from_H

REGIONS = [  # name, a_min, a_max, representative heliocentric distance for the photometric estimate
    ("Hungaria (1.78-2.0 au)", 1.78, 2.0, 1.9),
    ("Inner main belt (2.0-2.5 au)", 2.0, 2.5, 2.3),
    ("Middle main belt (2.5-2.82 au)", 2.5, 2.82, 2.65),
    ("Outer main belt (2.82-3.27 au)", 2.82, 3.27, 3.1),
    ("Cybele (3.27-3.7 au)", 3.27, 3.7, 3.45),
    ("Hilda (3.7-4.2 au)", 3.7, 4.2, 3.95),
    ("Jupiter Trojans (4.8-5.5 au)", 4.8, 5.5, 5.2),
    ("Centaurs / giant-planet region (5.5-30 au)", 5.5, 30.0, 15.0),
    ("Trans-Neptunian (30-100 au)", 30.0, 100.0, 42.0),
]


def load_mpcorb():
    cache = os.path.join(RAW, "mpcorb_parsed.npz")
    if os.path.exists(cache):
        z = np.load(cache)
        return z["H"], z["a"], z["e"], z["i"], z["year"]
    H, a, e, inc, yr = [], [], [], [], []
    started = False
    with gzip.open(os.path.join(RAW, "MPCORB.DAT.gz"), "rt", encoding="latin-1") as f:
        for line in f:
            if not started:
                if line.startswith("-----"):
                    started = True
                continue
            if len(line) < 160:
                continue
            try:
                h = float(line[8:13]) if line[8:13].strip() else np.nan
                aa = float(line[92:103])
                ee = float(line[70:79])
                ii = float(line[59:68])
            except ValueError:
                continue
            # discovery year from the packed provisional designation of unnumbered objects (e.g. K25A01B)
            des = line[0:7]
            y = 0
            if des[0] in "IJK" and des[1:3].isdigit() and not des[:5].isdigit():
                y = {"I": 1800, "J": 1900, "K": 2000}[des[0]] + int(des[1:3])
            H.append(h)
            a.append(aa)
            e.append(ee)
            inc.append(ii)
            yr.append(y)
    H, a, e, inc, yr = map(np.array, (H, a, e, inc, yr))
    np.savez_compressed(cache, H=H, a=a, e=e, i=inc, year=yr)
    return H, a, e, inc, yr


def limits(Hs, step=0.2):
    edges = np.arange(np.floor(np.nanmin(Hs)) - 0.2, np.ceil(np.nanmax(Hs)) + 0.4, step)
    cnt, _ = np.histogram(Hs, edges)
    cen = 0.5 * (edges[1:] + edges[:-1])
    # smooth mode
    sm = np.convolve(cnt, np.ones(3) / 3, mode="same")
    ipk = int(np.argmax(sm))
    Hpk = cen[ipk]
    fitsel = (cen >= Hpk - 3.0) & (cen <= Hpk - 1.0) & (cnt > 0)
    if fitsel.sum() < 4:
        return dict(H_mode=Hpk, slope=np.nan, H_90=np.nan, H_50=np.nan), cen, cnt, np.full(len(cen), np.nan)
    slope, icpt = np.polyfit(cen[fitsel], np.log10(cnt[fitsel]), 1)
    model = 10 ** (icpt + slope * cen)
    ratio = cnt / model
    def crossing(level):
        for k in range(np.where(cen >= Hpk - 1.0)[0][0], len(cen)):
            if ratio[k] < level:
                if k == 0:
                    return cen[k]
                r0, r1 = ratio[k - 1], ratio[k]
                return cen[k - 1] + (r0 - level) / (r0 - r1 + 1e-12) * step if r0 > level else cen[k]
        return np.nan
    return dict(H_mode=Hpk, slope=slope, H_90=crossing(0.9), H_50=crossing(0.5)), cen, cnt, ratio


def main():
    H, a, e, inc, yr = load_mpcorb()
    q = a * (1 - e)
    print(f"[belts] MPCORB objects parsed: {len(H)}; with H: {np.isfinite(H).sum()}; "
          f"unnumbered with 2025 designations: {(yr == 2025).sum()}, 2026: {(yr == 2026).sum()}")
    rows, hrows = [], []
    for name, a0, a1, rhel in REGIONS:
        sel = (a >= a0) & (a < a1) & np.isfinite(H) & (q > 1.3)
        if name.startswith("Hungaria"):
            sel &= (e < 0.18) & (inc > 16) & (inc < 34)
        if name.startswith("Jupiter"):
            sel &= (e < 0.3)
        Hs = H[sel]
        res, cen, cnt, ratio = limits(Hs)
        # photometric estimate: V=21.5 at opposition (zero phase), distance rhel, rhel-1
        H_phot = 21.5 - 5 * np.log10(rhel * (rhel - 1.0))
        row = dict(region=name, n_known=int(sel.sum()), H_mode=res["H_mode"], bright_slope=res["slope"],
                   H_90=res["H_90"], H_50=res["H_50"], H_phot_V21p5=H_phot)
        for tag, Hx in [("H90", res["H_90"]), ("H50", res["H_50"])]:
            for p in (0.2, 0.05, 0.01):
                row[f"D_km_{tag}_p{p}"] = float(D_from_H(Hx, p)) / 1e3 if np.isfinite(Hx) else np.nan
        rows.append(row)
        for c, n, r in zip(cen, cnt, ratio):
            hrows.append(dict(region=name, H=c, count=int(n), ratio_to_powerlaw=r))
    out = pd.DataFrame(rows)
    out.to_csv(os.path.join(DATA, "belt_completeness.csv"), index=False)
    pd.DataFrame(hrows).to_csv(os.path.join(DATA, "belt_H_histograms.csv"), index=False)
    pd.set_option("display.width", 250)
    print(out.round(2).to_string(index=False))
    return out


if __name__ == "__main__":
    main()