"""Download every raw input (skips files that already exist in RAW).

Raw files go to fermi_raw/seta_haystack/ (override with $SETA_RAW).
"""
import json
import os
import re
import time
import urllib.parse

from common import RAW, fetch, fetch_text

SBDB = "https://ssd-api.jpl.nasa.gov/sbdb_query.api"

FILES = {
    # Minor Planet Center orbit catalogues
    "NEA.txt": "https://www.minorplanetcenter.net/iau/MPCORB/NEA.txt",
    "MPCORB.DAT.gz": "https://www.minorplanetcenter.net/iau/MPCORB/MPCORB.DAT.gz",
    # Granvik et al. 2018 debiased NEO model realisation (802,000 synthetic NEOs, 17<H<25)
    "granvik2018.dat.gz": "https://www.mv.helsinki.fi/home/mgranvik/data/Granvik+_2018_Icarus/Granvik+_2018_Icarus.dat.gz",
    # J. McDowell's General Catalog of Artificial Space Objects: heliocentric and deep-space catalogues
    "hcocat.tsv": "https://planet4589.org/space/gcat/tsv/cat/hcocat.tsv",
    "deepcat.tsv": "https://planet4589.org/space/gcat/tsv/cat/deepcat.tsv",
    "landercat.tsv": "https://planet4589.org/space/gcat/tsv/cat/landercat.tsv",
    # Project Pluto (B. Gray) index of artificial objects picked up by asteroid surveys
    "pseudo.htm": "https://www.projectpluto.com/pluto/mpecs/pseudo.htm",
    # mirror of the MPC editorial notices index (used to list deleted designations)
    "MPCEdNot.htm": "https://birtwhistle.org.uk/MPCEdNot.htm",
    # HiRISE RDR cumulative index (image footprints)
    "RDRCUMINDEX.TAB": "https://hirise-pds.lpl.arizona.edu/PDS/INDEX/RDRCUMINDEX.TAB",
    "RDRCUMINDEX.LBL": "https://hirise-pds.lpl.arizona.edu/PDS/INDEX/RDRCUMINDEX.LBL",
    # NASA Orbital Debris Program Office FAQ (source of the "objects > 10 cm are tracked" statement)
    "odpo_faq.html": "https://orbitaldebris.jsc.nasa.gov/faq/",
    # JPL SBDB: all near-Earth objects with orbit and physical fields
    "sbdb_neo.json": SBDB + "?fields=spkid,pdes,full_name,kind,prefix,class,H,G,diameter,albedo,a,e,i,om,w,ma,q,ad,moid,"
                            "first_obs,last_obs,data_arc,n_obs_used,condition_code&sb-group=neo",
    "sbdb_neo_phys.json": SBDB + "?fields=spkid,pdes,H,A1,A2,A3,rot_per,spec_B,spec_T,diameter,albedo,BV,UB&sb-group=neo&sb-kind=a",
    # (revision) the same near-Earth asteroids with the osculating epoch of each orbit: used in s01b to compute
    # the sky geometry on the night of first observation
    "sbdb_neo_epoch.json": SBDB + "?fields=pdes,full_name,kind,H,epoch,a,e,i,om,w,ma,tp,q,first_obs,last_obs,data_arc,producer"
                                  "&sb-group=neo&sb-kind=a",
    # (revision) B. Gray's list of high-orbit artificial objects for which he maintains orbits fitted to
    # astrometry reported by asteroid surveys (used in s05 as a check of the high-Earth-orbit class)
    "gray_tle_list.txt": "https://raw.githubusercontent.com/Bill-Gray/tles/master/tle_list.txt",
}

# MPECs announcing deletion of minor-planet designations (read and classified in s08)
MPECS = ["2026-J119", "2025-U223", "2025-P132", "2025-A49", "2024-N07", "2024-J297", "2023-P34",
         "2022-U79", "2021-H90", "2021-D62", "2020-X56", "2020-G97", "2015-H125", "2014-X55",
         "2013-Q25", "2012-D99", "2012-B08", "2010-N23"]


def mpec_packed(m):
    y, rest = m.split("-")
    h, n = rest[0], int(rest[1:])
    c = {"18": "I", "19": "J", "20": "K"}[y[:2]] + y[2:]
    if n < 100:
        p = "%02d" % n
    elif n < 360:
        p = chr(ord("A") + (n - 100) // 10) + str(n % 10)
    else:
        p = chr(ord("a") + (n - 360) // 10) + str(n % 10)
    return c, c + h + p


def sbdb_count(cdata=None, **kw):
    q = {"fields": "spkid", "limit": "1", "sb-kind": "a"}
    q.update(kw)
    if cdata:
        q["sb-cdata"] = json.dumps(cdata)
    d = json.loads(fetch_text(SBDB + "?" + urllib.parse.urlencode(q)))
    return int(d["count"])


def main():
    for name, url in FILES.items():
        dest = os.path.join(RAW, name)
        had = os.path.exists(dest)
        fetch(url, dest)
        print(f"[download] {name:22s} {'cached' if had else 'fetched'} {os.path.getsize(dest):>12,d} bytes")
    os.makedirs(os.path.join(RAW, "mpec"), exist_ok=True)
    for m in MPECS:
        d, p = mpec_packed(m)
        dest = os.path.join(RAW, "mpec", p + ".html")
        if not os.path.exists(dest):
            fetch(f"https://www.minorplanetcenter.net/mpec/{d}/{p}.html", dest, min_bytes=10)
            time.sleep(1)
    print(f"[download] {len(MPECS)} MPEC editorial notices cached")
    # counts of physically characterised asteroids in SBDB
    dest = os.path.join(RAW, "sbdb_charact_counts.json")
    if not os.path.exists(dest):
        res = {"all": sbdb_count()}
        for f in ["albedo", "rot_per", "spec_B", "spec_T", "A1", "A2", "diameter"]:
            res[f] = sbdb_count({"AND": [f + "|DF"]})
            time.sleep(0.5)
        for cls in ["IMB", "MBA", "OMB", "TJN", "CEN", "TNO", "MCA"]:
            res["class_" + cls] = sbdb_count(**{"sb-class": cls})
            res[f"class_{cls}_albedo"] = sbdb_count({"AND": ["albedo|DF"]}, **{"sb-class": cls})
            res[f"class_{cls}_rot"] = sbdb_count({"AND": ["rot_per|DF"]}, **{"sb-class": cls})
            time.sleep(0.5)
        json.dump(res, open(dest, "w"), indent=1)
    print("[download] SBDB characterisation counts cached")
    # (revision) the same counts as a function of absolute magnitude, so that the recognition model can use the
    # characterisation rate of objects of the right size instead of a catalogue average
    dest = os.path.join(RAW, "sbdb_charact_by_H.json")
    if not os.path.exists(dest):
        edges = [-5, 8, 10, 12, 14, 15, 16, 17, 18, 20, 30]
        groups = {"main_belt": "IMB,MBA,OMB", "trojans": "TJN", "tno_centaur": "TNO,CEN"}
        anyphys = {"OR": ["albedo|DF", "rot_per|DF", "spec_B|DF", "spec_T|DF"]}
        res = {"H_edges": edges, "groups": groups, "counts": {}}
        for g, cls in groups.items():
            rows = []
            for a, b in zip(edges[:-1], edges[1:]):
                rg = f"H|RG|{a}|{b - 0.001:.3f}"
                rows.append(dict(H_lo=a, H_hi=b, n=sbdb_count({"AND": [rg]}, **{"sb-class": cls}),
                                 n_albedo=sbdb_count({"AND": [rg, "albedo|DF"]}, **{"sb-class": cls}),
                                 n_any_physical=sbdb_count({"AND": [rg, anyphys]}, **{"sb-class": cls})))
                time.sleep(0.3)
            res["counts"][g] = rows
        json.dump(res, open(dest, "w"), indent=1)
    print("[download] SBDB characterisation counts by absolute magnitude cached")
    # interstellar-object check: every SBDB object with e > 1.05
    dest = os.path.join(RAW, "sbdb_hyperbolic.json")
    if not os.path.exists(dest):
        q = {"fields": "full_name,kind,prefix,class,H,e,q,i,first_obs,last_obs,tp",
             "sb-cdata": json.dumps({"AND": ["e|GT|1.05"]})}
        open(dest, "w").write(fetch_text(SBDB + "?" + urllib.parse.urlencode(q)))
    print("[download] SBDB hyperbolic-object list cached")


if __name__ == "__main__":
    main()