#!/usr/bin/env python3
"""Regenerates data/gfz-drivers.bin, the solar and geomagnetic record the model reads.

Usage: scripts/fetch-drivers.py [path to a local copy of the source file]

The record is committed as a compact binary blob and embedded with include_bytes!, so a
caller needs no network, no data directory and nothing to parse at startup. The blob is
numbers only; its provenance is written alongside it in data/gfz-drivers.txt, which the
crate compiles in as a string so that a program can print where its weather came from.

Source: GFZ Potsdam's combined Kp/ap/Ap/SN/F10.7 record, under CC BY 4.0. Chosen over
Celestrak's SW-All.csv for two reasons: it states a license where Celestrak states none,
which matters for a file committed into an MIT crate, and Celestrak's runs decades past
today on its own predictions, which are not measurements.

The sunspot column is dropped, being the one part of the GFZ file under CC BY-NC 4.0, and
so is the flux adjusted to one astronomical unit: NRLMSISE-00 asks for the flux at the
Earth's actual distance from the Sun, which is the observed column.

Unlike the coefficient tables, this record is meant to change — it gains a row a day — so
there is no digest pinned in the crate. What is recorded instead is the digest of what was
downloaded, with the date and the coverage, in the provenance file.
"""

import hashlib
import struct
import subprocess
import sys
from datetime import date, timedelta
from pathlib import Path

URL = "https://kp.gfz-potsdam.de/app/files/Kp_ap_Ap_SN_F107_since_1932.txt"
ATTRIBUTION = (
    "GFZ German Research Centre for Geosciences and Geomagnetic Observatory Niemegk, "
    "CC BY 4.0"
)
MISSING = -1.0

# Geomagnetic indices reach back to 1932 and the solar flux to February 1947, but the early
# flux record is sporadic: 105 missing days in 1947 and one hole of 44. From 1957 it is all
# but complete, which is where this record starts and where the other published driver
# files start too.
START = date(1957, 1, 1)

# A run of unmeasured flux days this long is refused rather than interpolated.
LONGEST_FILL = 7

ROOT = Path(__file__).resolve().parent.parent
OUT = ROOT / "data" / "gfz-drivers.bin"
NOTE = ROOT / "data" / "gfz-drivers.txt"

# Modified Julian Day zero. The record is indexed by MJD because a UT day is not always
# 86400 seconds and MJD still increments by exactly one across a leap second.
MJD_EPOCH = date(1858, 11, 17)

MAGIC = b"NRLMSISEDRV"
VERSION = 1

# Field positions in the source, which is whitespace-aligned with 28 columns.
YEAR, MONTH, DAY = 0, 1, 2
AP_FIRST, AP_LAST = 15, 22
AP_DAILY = 23
F107_OBSERVED = 25
DEFINITIVE = 27


def source_text(argv):
    if len(argv) > 1:
        raw = Path(argv[1]).read_bytes()
        print(f"reading {argv[1]}")
    else:
        print(f"downloading {URL}")
        raw = subprocess.run(
            ["curl", "-sS", "--fail", "--max-time", "180", "-L", URL],
            check=True,
            capture_output=True,
        ).stdout
    return raw.decode(), hashlib.sha256(raw).hexdigest()


def parse(text):
    """Daily rows from START onward, with the flux left as None where unmeasured."""
    rows = []
    for line in text.splitlines():
        if line.startswith("#") or not line.strip():
            continue
        fields = line.split()
        if len(fields) != 28:
            raise SystemExit(f"expected 28 fields, found {len(fields)}: {line[:60]}")

        when = date(int(fields[YEAR]), int(fields[MONTH]), int(fields[DAY]))
        if when < START:
            continue

        observed = float(fields[F107_OBSERVED])
        rows.append(
            {
                "date": when,
                "ap": [int(float(f)) for f in fields[AP_FIRST : AP_LAST + 1]],
                "ap_daily": int(float(fields[AP_DAILY])),
                "flux": None if observed <= MISSING else observed,
                "definitive": int(fields[DEFINITIVE]) == 1,
                "filled": False,
            }
        )

    if not rows:
        raise SystemExit(f"no rows at or after {START}")

    # The record is addressed by day, so a missing date would shift everything after it.
    for earlier, later in zip(rows, rows[1:]):
        if later["date"] - earlier["date"] != timedelta(days=1):
            raise SystemExit(f"missing day between {earlier['date']} and {later['date']}")

    return rows


def trim(rows):
    """Drops unmeasured days at either end.

    START is a round date rather than a measurement and the flux happens not to have been
    taken on it. The far end matters more: an unmeasured day there would have to be
    extrapolated, and refusing to evaluate past the record is precisely so that nothing is.
    """
    measured = [index for index, row in enumerate(rows) if row["flux"] is not None]
    if not measured:
        raise SystemExit("no measured solar flux at all")
    return rows[measured[0] : measured[-1] + 1]


def fill(rows):
    """Interpolates the days the flux was not measured, marks them, counts them.

    The gaps are isolated — one or two days, five at the very worst — and the flux moves a
    few percent over that, so a straight line between the measurements either side is a
    fair reconstruction. A long gap is not, and is refused.
    """
    measured = [index for index, row in enumerate(rows) if row["flux"] is not None]
    filled = 0
    for before, after in zip(measured, measured[1:]):
        span = after - before
        if span == 1:
            continue
        if span - 1 >= LONGEST_FILL:
            raise SystemExit(
                f"{span - 1} unmeasured days from {rows[before + 1]['date']}, too many "
                f"to interpolate"
            )
        for offset in range(1, span):
            weight = offset / span
            rows[before + offset]["flux"] = (
                rows[before]["flux"] * (1.0 - weight) + rows[after]["flux"] * weight
            )
            rows[before + offset]["filled"] = True
            filled += 1
    return filled


def mjd(when):
    return when.toordinal() - MJD_EPOCH.toordinal()


def encode(rows, generated):
    """The blob: a fixed header, then twenty-one bytes a day, little-endian throughout.

    The flux is stored in tenths because that is the precision the source states it to, so
    nothing is lost and a float is not needed. ap and Ap are whole numbers on the Kp scale,
    which reaches 400.
    """
    blob = bytearray()
    blob += MAGIC
    blob += struct.pack("<BiI", VERSION, mjd(rows[0]["date"]), len(rows))
    blob += struct.pack("<I", int(generated.strftime("%Y%m%d")))

    for row in rows:
        tenths = round(row["flux"] * 10.0)
        if not 0 <= tenths <= 0xFFFF:
            raise SystemExit(f"{row['date']}: flux {row['flux']} does not fit")
        values = [row["ap_daily"], *row["ap"]]
        for value in values:
            if not 0 <= value <= 0xFFFF:
                raise SystemExit(f"{row['date']}: ap {value} does not fit")
        flags = (1 if row["definitive"] else 0) | (0 if row["filled"] else 2)
        blob += struct.pack("<H", tenths)
        blob += struct.pack("<9H", *values)
        blob += struct.pack("<B", flags)
    return bytes(blob)


def note(rows, digest, generated, filled):
    """The provenance, which the crate compiles in as a string."""
    definitive = sum(1 for row in rows if row["definitive"])
    highest = max(rows, key=lambda row: row["flux"])
    lowest = min(rows, key=lambda row: row["flux"])
    dropouts = [row for row in rows if row["flux"] < 60.0 and not row["filled"]]
    peaks = [row for row in rows if row["flux"] > 400.0 and not row["filled"]]
    return f"""The solar and geomagnetic record NRLMSISE-00 reads.

Source      {URL}
Attribution {ATTRIBUTION}
Downloaded  {generated.isoformat()}
Digest      sha256 {digest}
            of the file as downloaded, not of the blob next to this note.

Coverage    {rows[0]['date'].isoformat()} to {rows[-1]['date'].isoformat()}, \
{len(rows)} days
            Modified Julian Day {mjd(rows[0]['date'])} to {mjd(rows[-1]['date'])}
            The last day a model can be evaluated on is forty days earlier than the last
            day held, because the 81-day flux average the model wants is centered.

Flux        The observed 10.7 cm radio flux, at the Earth's actual distance from the Sun,
            which is what the model asks for. The flux adjusted to one astronomical unit
            is not kept. {filled} of {len(rows)} days were not measured and are
            interpolated from the days either side; they are marked in the blob and
            reported by Day::flux_measured.

Geomagnetic The daily mean Ap and the eight three-hour ap indices. {definitive} of
            {len(rows)} days are definitive rather than preliminary, which Day::definitive
            reports.

Extremes    The flux runs from {lowest['flux']} on {lowest['date'].isoformat()} to \
{highest['flux']} on {highest['date'].isoformat()}.
            Both ends are in the source and are passed through unaltered. The high end is
            real: the daily flux includes radio bursts, so a large flare puts it several
            times above the solar-maximum background, and there are {len(peaks)} measured
            days above 400. The low end looks like instrument trouble rather than the Sun —
            {len(dropouts)} measured days fall below 60, each an isolated day between
            neighbors around 70, and both are in early March. Whether to smooth them is a
            decision for whoever is using the record, not for the record.

Dropped     The sunspot number, which is the one part of the source under CC BY-NC 4.0,
            and the adjusted flux.

Regenerate  scripts/fetch-drivers.py
"""


def main():
    text, digest = source_text(sys.argv)
    rows = trim(parse(text))
    filled = fill(rows)
    generated = date.today()

    OUT.parent.mkdir(parents=True, exist_ok=True)
    OUT.write_bytes(encode(rows, generated))
    NOTE.write_text(note(rows, digest, generated, filled))

    print(f"digest {digest}")
    print(f"interpolated {filled} unmeasured days of {len(rows)}")
    print(
        f"wrote {OUT.relative_to(ROOT)}: {len(rows)} days, "
        f"{rows[0]['date']} to {rows[-1]['date']}, {OUT.stat().st_size / 1e3:.0f} KB"
    )


if __name__ == "__main__":
    main()