#!/usr/bin/env python3
"""Regenerates src/atmosphere/space_weather.rs, the drivers a density model reads.

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

The record is emitted as Rust rather than as data to parse, the same way the gravity
coefficients are. Nothing at runtime reads a file or validates a field.

Unlike the gravity coefficients and the kernels, this record is *meant* to change: it gains
a row a day. So there is no digest to pin it to. What this records instead is the digest of
what was downloaded, with the date and the coverage, in the output's own header —
provenance for a file that moves rather than a lock on one that does not.

Source: GFZ Potsdam's combined Kp/ap/Ap/SN/F10.7 record. Chosen over Celestrak's
SW-All.csv for two reasons. It carries an explicit license where Celestrak states none, and
that matters for a file committed into an MIT crate. And Celestrak's runs decades past
today on its own predictions, which are not measurements; a propagation that silently used
them would not be reproducible.

The sunspot column is dropped. It is the one part of the GFZ file under CC BY-NC 4.0 rather
than CC BY 4.0, and no density model here wants it.
"""

import hashlib
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: 39 gaps in seventy years, none longer than five days, none longer than two
# since 2020. So that is where the retained record starts, which is also where the other
# published driver files start.
START = date(1957, 1, 1)

# A run of unmeasured days this long is refused rather than interpolated. Five is the worst
# the source holds from 1957 on, so this leaves room for it to change without quietly
# inventing a fortnight of weather.
LONGEST_FILL = 7

ROOT = Path(__file__).resolve().parent.parent
OUT = ROOT / "src" / "atmosphere" / "space_weather.rs"

# 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)

# 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, F107_ADJUSTED = 25, 26
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

        def flux(at):
            value = float(fields[at])
            return None if value <= MISSING else value

        rows.append(
            {
                "date": when,
                "ap": [int(float(f)) for f in fields[AP_FIRST : AP_LAST + 1]],
                "ap_daily": int(float(fields[AP_DAILY])),
                "observed": flux(F107_OBSERVED),
                "adjusted": flux(F107_ADJUSTED),
                "definitive": int(fields[DEFINITIVE]),
                "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 any 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 rather than interpolated, and the whole point of refusing to propagate
    past the record is not to do that.
    """
    measured = [index for index, row in enumerate(rows) if row["observed"] 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["observed"] 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
            for key in ("observed", "adjusted"):
                rows[before + offset][key] = (
                    rows[before][key] * (1.0 - weight) + rows[after][key] * weight
                )
            rows[before + offset]["filled"] = True
            filled += 1
    return filled


def main():
    text, digest = source_text(sys.argv)
    rows = trim(parse(text))
    filled = fill(rows)
    print(f"digest {digest}")
    print(f"interpolated {filled} unmeasured days of {len(rows)}")

    OUT.parent.mkdir(parents=True, exist_ok=True)
    OUT.write_text(rust_source(rows, digest, date.today().isoformat(), filled))
    print(
        f"wrote {OUT.relative_to(ROOT)}: {len(rows)} days, "
        f"{rows[0]['date']} to {rows[-1]['date']}, {OUT.stat().st_size / 1e6:.1f} MB"
    )


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


def wrapped(tokens, width=96, indent="    "):
    """The tokens as comma-separated lines, packed to `width`."""
    lines, line = [], indent
    for token in tokens:
        piece = token + ","
        if len(line) + len(piece) > width and line != indent:
            lines.append(line)
            line = indent
        line += piece
    if line != indent:
        lines.append(line)
    return "\n".join(lines)


def array(name, kind, doc, tokens, count):
    return (
        f"{doc}\n#[rustfmt::skip]\npub(super) static {name}: [{kind}; {count}] = [\n"
        f"{wrapped(tokens)}\n];\n"
    )


def rust_source(rows, digest, downloaded, filled):
    """The generated module: provenance in the docs, the record in six flat arrays."""
    first, last = rows[0]["date"], rows[-1]["date"]
    doc = f"""//! Solar and geomagnetic drivers for atmospheric density, one row per UT day.
//!
//! Generated by `scripts/fetch-space-weather.py`; do not edit by hand. Committed as Rust
//! rather than as data so that a clone needs no network, no data directory and no parser.
//!
//! Source: <{URL}>
//! downloaded {downloaded}, SHA-256 `{digest}`.
//!
//! Attribution: {ATTRIBUTION}.
//!
//! The digest records what was downloaded rather than pinning it: the source gains a row a
//! day, so there is nothing fixed to pin. Its sunspot column is dropped, being the one part
//! of it under a non-commercial license.
//!
//! Coverage: {first} to {last}, {len(rows)} days. Geomagnetic indices reach back to 1932
//! and the flux to 1947, but the early flux record is too sporadic to use: 105 missing days
//! in 1947 and one hole of 44. Of the days here, {filled} had no flux measurement and were
//! interpolated from their neighbors, never more than five in a row.
"""
    parts = [
        doc,
        f"/// Modified Julian Day of the first row, {first}.\n"
        f"pub(super) const FIRST_DAY: i64 = {mjd(first)};\n",
        f"/// How many days the record holds.\npub(super) const DAYS: usize = {len(rows)};\n",
        array(
            "AP",
            "[f64; 8]",
            "/// The eight three-hour geomagnetic indices of each day, earliest first.",
            ["[" + ",".join(f"{v}." for v in row["ap"]) + "]" for row in rows],
            "DAYS",
        ),
        array(
            "AP_DAILY",
            "f64",
            "/// The daily mean geomagnetic index.",
            [f"{row['ap_daily']}." for row in rows],
            "DAYS",
        ),
        array(
            "FLUX_OBSERVED",
            "f64",
            "/// Solar flux as measured, in solar flux units.",
            [f"{row['observed']:.1f}" for row in rows],
            "DAYS",
        ),
        array(
            "FLUX_ADJUSTED",
            "f64",
            "/// Solar flux scaled to one astronomical unit.",
            [f"{row['adjusted']:.1f}" for row in rows],
            "DAYS",
        ),
        array(
            "DEFINITIVE",
            "bool",
            "/// Whether the geomagnetic indices are settled rather than preliminary.",
            ["true" if row["definitive"] else "false" for row in rows],
            "DAYS",
        ),
        array(
            "FLUX_MEASURED",
            "bool",
            "/// Whether the flux was measured that day rather than interpolated across a gap.",
            ["false" if row["filled"] else "true" for row in rows],
            "DAYS",
        ),
    ]
    return "\n".join(parts)


if __name__ == "__main__":
    main()