"""Download and parse the public data used by the impact_window study.

Raw files go to cfg.RAW (outside the shared folder); small derived tables go to data/.

  1. Neumann et al. 2015 (Sci. Adv. 1, e1500852), Table 1: 74 lunar basins >= 200 km from GRAIL
     and LOLA.  Source: Europe PMC full-text XML of PMC4646831.            -> data/neumann2015_basins.csv
  2. Fassett et al. 2012 (JGR 117, E00H06), Table 1: N(20) crater densities superposed on 30
     basins.  Source: NASA NTRS copy of the paper (PDF).  The table was transcribed by hand
     (the PDF text layer drops the +/- glyph) and every transcribed value is cross-checked
     against the downloaded text.                                         -> data/fassett2012_table1.csv
  3. Orgel et al. 2018 (JGR Planets 123, 748): absolute model ages of the same 30 basins with the
     buffered non-sparseness correction.  The journal table is paywalled; the identical basin
     list with model ages is in the open LPSC abstract (LPSC 49, #1395), parsed here.
                                                                          -> data/orgel2018_lpsc_ages.csv
  4. JPL Small-Body Database: every asteroid with a measured diameter >= 10 km.
                                                                          -> data/sbdb_asteroids_d10.csv

If a download fails and the derived CSV already exists, the existing CSV is kept and a warning
is printed, so the analysis can still be rerun offline.
"""
from __future__ import annotations

import html
import json
import os
import re
import subprocess
import urllib.parse
import urllib.request

import pandas as pd

import config as cfg

UA = {"User-Agent": "Mozilla/5.0 (research script; impact_window study)"}


def _get(url, dest, timeout=180):
    if os.path.exists(dest) and os.path.getsize(dest) > 1000:
        return dest
    os.makedirs(os.path.dirname(dest), exist_ok=True)
    req = urllib.request.Request(url, headers=UA)
    with urllib.request.urlopen(req, timeout=timeout) as r, open(dest + ".part", "wb") as f:
        f.write(r.read())
    os.replace(dest + ".part", dest)
    return dest


def _pdftotext(pdf, layout=True):
    txt = pdf[:-4] + (".layout.txt" if layout else ".txt")
    try:
        subprocess.run(["pdftotext"] + (["-layout"] if layout else []) + [pdf, txt], check=True)
        return open(txt, encoding="utf8", errors="replace").read()
    except Exception:
        from pdfminer.high_level import extract_text   # fallback, slower
        return extract_text(pdf)


def _clean(cell):
    cell = re.sub(r"<sup>.*?</sup>", "", cell, flags=re.S)       # footnote markers a, b, c, dagger
    return html.unescape(re.sub(r"<[^>]+>", "", cell)).replace("−", "-").strip()


# ---------------------------------------------------------------------------------------------
def fetch_neumann():
    out = os.path.join(cfg.DATA, "neumann2015_basins.csv")
    xml = open(_get(cfg.URL_NEUMANN, os.path.join(cfg.RAW, "neumann", "neumann2015.xml")),
               encoding="utf8").read()
    table = re.search(r"<table-wrap[^>]*>.*?</table-wrap>", xml, flags=re.S).group(0)
    rows = []
    for tr in re.findall(r"<tr[^>]*>(.*?)</tr>", table, flags=re.S):
        raw = re.findall(r"<td[^>]*?(?:/>|>(.*?)</td>)", tr, flags=re.S)
        cells = [_clean(c) for c in raw]
        if len(cells) != 8 or not re.match(r"^-?\d", cells[1] or ""):
            continue                                              # header rows
        main = cells[3]
        conf = {"(": "probable", "[": "possible", "{": "suggested by scaling"}.get(main[:1], "certain")
        num = lambda s: float(re.search(r"\d+(\.\d+)?", s).group(0)) if re.search(r"\d", s) else float("nan")
        rows.append(dict(name=cells[0].replace("’", "'").replace("–", "-"),
                         lat_N=float(cells[1]), lon_E=float(cells[2]),
                         D_main_km=num(main), main_confidence=conf, D_inner_km=num(cells[4]),
                         notes=cells[5], bouguer_diam_km=cells[6], bouguer_contrast_mgal=cells[7]))
    df = pd.DataFrame(rows)
    assert len(df) == 74, f"expected 74 basins in Neumann et al. Table 1, parsed {len(df)}"
    assert abs(df.loc[df.name == "Imbrium", "D_main_km"].iloc[0] - 1321) < 1
    df.to_csv(out, index=False)
    return df


# ---------------------------------------------------------------------------------------------
# Fassett et al. 2012, Table 1.  Columns: basin, bracketed (uncertain), period, N(20), err, flag L
# ("value may be low"), N(64), err.  N per 10^6 km^2.  Hand transcription of the NTRS PDF.
FASSETT = [
    ("South Pole-Aitken", 0, "PN", 156, 7, 1, 33, 3),
    ("Coulomb-Sarton", 0, "PN", 271, 54, 0, 32, 14),
    ("Dirichlet-Jackson", 0, "PN", 266, 36, 0, 33, 12),
    ("Cruger-Sirsalis", 0, "PN", 262, 46, 0, 39, 15),
    ("Smythii", 0, "PN", 225, 19, 0, 28, 6),
    ("Schiller-Zucchius", 1, "PN", 211, 47, 0, 41, 17),
    ("Amundsen-Ganswindt", 1, "PN", 202, 37, 1, 47, 17),
    ("Nubium", 0, "PN", 195, 18, 0, 33, 7),
    ("Poincare", 1, "PN", 194, 44, 0, 44, 16),
    ("Lorentz", 0, "PN", 179, 31, 0, 32, 12),
    ("Fitzgerald-Jackson", 1, "PN", 175, 34, 1, 43, 15),
    ("Birkhoff", 1, "PN", 170, 33, 1, 46, 17),
    ("Ingenii", 1, "PN", 167, 33, 1, 40, 14),
    ("Serenitatis", 1, "PN?", 298, 60, 0, 28, 20),
    ("Apollo", 0, "PN/N", 151, 23, 0, 12, 6),
    ("Freundlich-Sharonov", 0, "PN/N", 140, 18, 0, 16, 6),
    ("Nectaris", 0, "N", 135, 14, 0, 17, 5),
    ("Korolev", 0, "N/PN", 127, 22, 0, 9, 5),
    ("Mendeleev", 1, "N/PN", 129, 36, 0, 16, 10),
    ("Hertzsprung", 0, "N/PN", 129, 22, 0, 16, 6),
    ("Grimaldi", 1, "N/PN", 126, 28, 1, 22, 10),
    ("Mendel-Rydberg", 0, "N/PN", 125, 17, 0, 12, 5),
    ("Planck", 1, "N/PN", 118, 36, 0, 22, 14),
    ("Moscoviense", 0, "N", 120, 17, 0, 9, 4),
    ("Crisium", 0, "N", 113, 11, 0, 8, 3),
    ("Humorum", 0, "N", 108, 21, 0, 12, 6),
    ("Humboldtianum", 0, "N", 93, 14, 0, 11, 4),
    ("Imbrium", 0, "I", 30, 5, 0, 4, 2),
    ("Schrodinger", 0, "I", 19, 7, 0, float("nan"), float("nan")),
    ("Orientale", 0, "I", 21, 4, 0, 1, 1),
]


def fetch_fassett():
    out = os.path.join(cfg.DATA, "fassett2012_table1.csv")
    df = pd.DataFrame(FASSETT, columns=["name", "bracketed", "period", "N20", "N20_err", "may_be_low",
                                        "N64", "N64_err"])
    checked = 0
    try:
        txt = _pdftotext(_get(cfg.URL_FASSETT, os.path.join(cfg.RAW, "fassett", "fassett2012_ntrs.pdf")))
        i = txt.find("Derived Age Sequence of Lunar Basins")
        block = txt[i:i + 6000].replace("é", "e").replace("ö", "o").replace("ü", "u")
        # In the PDF text layer the +/- glyph is the control character \x06, so N(20) +/- err
        # (and the "L" flag) can be parsed automatically and compared with the transcription.
        pat = re.compile(r"^\[?([A-Za-z\- ]+?)\]?(?: \(no buff\.\))?\s{2,}(?:Beginning of )?"
                         r"(PN\?|PN/N|N/PN|PN|N|I)\s+(\d+)\s*\x06\s*(\d+)(L?)")
        parsed = {}
        for line in block.splitlines():
            m = pat.match(line)
            if m:
                parsed[m.group(1).strip()] = (m.group(2), int(m.group(3)), int(m.group(4)), int(bool(m.group(5))))
        for _, r in df.iterrows():
            got = parsed.get(r["name"])
            want = (r.period, int(r.N20), int(r.N20_err), int(r.may_be_low))
            if got != want:
                raise AssertionError(f"Fassett Table 1 cross-check failed for {r['name']}: "
                                     f"parsed {got}, transcribed {want}")
            checked += 1
        assert len(parsed) == 30, f"parsed {len(parsed)} rows from the downloaded Fassett table"
    except AssertionError:
        raise
    except Exception as e:                                  # network or tool failure only
        print(f"  WARNING: could not download/parse the Fassett et al. PDF ({e}); "
              f"using the transcription without cross-check")
    df["crosschecked_rows"] = checked
    df.to_csv(out, index=False)
    return df


# ---------------------------------------------------------------------------------------------
def fetch_orgel():
    out = os.path.join(cfg.DATA, "orgel2018_lpsc_ages.csv")
    try:
        txt = _pdftotext(_get(cfg.URL_ORGEL, os.path.join(cfg.RAW, "orgel", "orgel_lpsc2018_1395.pdf")))
    except Exception as e:
        if os.path.exists(out):
            print(f"  WARNING: Orgel et al. abstract not downloaded ({e}); using existing {out}")
            return pd.read_csv(out)
        raise
    rows = []
    pat = re.compile(r"^\s{4,}([A-Z][A-Za-zöüé\- ]+?)\s+(?:(PN|N|I)\s+)?(\d\.\d\d)\s*±\s*(0\.\d+),\s*-(0\.\d+)")
    for line in txt.splitlines():
        m = pat.match(line)
        if m:
            name = (m.group(1).strip().replace("ö", "o").replace("ü", "u").replace("é", "e")
                    .replace("Sharanov", "Sharonov"))
            rows.append(dict(name=name, period_marker=m.group(2) or "", age_neukum_Ga=float(m.group(3)),
                             err_plus=float(m.group(4)), err_minus=float(m.group(5))))
    df = pd.DataFrame(rows)
    assert len(df) == 30, f"expected 30 basins in the Orgel et al. LPSC table, parsed {len(df)}"
    assert df.age_neukum_Ga.iloc[0] == 4.31 and df.name.iloc[0] == "South Pole-Aitken"
    df.to_csv(out, index=False)
    return df


# ---------------------------------------------------------------------------------------------
def fetch_sbdb():
    out = os.path.join(cfg.DATA, "sbdb_asteroids_d10.csv")
    dest = os.path.join(cfg.RAW, "sbdb", "sbdb_d10.json")
    try:
        if not (os.path.exists(dest) and os.path.getsize(dest) > 1000):
            q = urllib.parse.urlencode({
                "fields": "spkid,full_name,class,diameter,diameter_sigma,albedo,H,a,e,i",
                "sb-kind": "a", "sb-cdata": json.dumps({"AND": ["diameter|GE|10"]})})
            _get(cfg.URL_SBDB + "?" + q, dest)
        d = json.load(open(dest))
    except Exception as e:
        if os.path.exists(out):
            print(f"  WARNING: SBDB query failed ({e}); using existing {out}")
            return pd.read_csv(out)
        raise
    df = pd.DataFrame(d["data"], columns=d["fields"])
    df["full_name"] = df["full_name"].str.strip()
    for c in ("diameter", "diameter_sigma", "albedo", "H", "a", "e", "i"):
        df[c] = pd.to_numeric(df[c], errors="coerce")
    df = df.sort_values("diameter", ascending=False).reset_index(drop=True)
    df.to_csv(out, index=False)
    return df


def run():
    n = fetch_neumann()
    f = fetch_fassett()
    o = fetch_orgel()
    s = fetch_sbdb()
    mb = s[s["class"].isin(["MBA", "IMB", "OMB"])]
    return dict(neumann_rows=len(n), fassett_rows=len(f), fassett_crosschecked=int(f.crosschecked_rows.iloc[0]),
                orgel_rows=len(o), sbdb_rows=len(s), sbdb_mainbelt_rows=len(mb),
                mainbelt_counts={str(d): int((mb.diameter >= d).sum()) for d in (10, 20, 50, 100, 200, 300, 440, 500)})


if __name__ == "__main__":
    print(json.dumps(run(), indent=1))