#!/usr/bin/env python3
"""Regenerates src/gravity/egm96.rs from the published EGM96 coefficient file.

The coefficients are committed as Rust source rather than fetched at run time, so
that a clone needs no network and no test has to skip. This script exists so the
table is reproducible rather than something that was once pasted in.

Usage: scripts/generate-gravity-field.py [path to egm96_to360.ascii]

With no argument it downloads the file and checks it against the pinned digest.
"""

import hashlib
import re
import subprocess
import sys
from pathlib import Path

URL = "https://raw.githubusercontent.com/nasa/GMAT/main/application/data/gravity/earth/EGM96.cof"
SHA256 = "4735c9301e20ee8db09b6d0882a24671b327cd35ec5740ea497c8eeb43a39807"
DOWNLOADED = "2026-10-01"
MAX_DEGREE = 360
PER_LINE = 4

ROOT = Path(__file__).resolve().parent.parent
OUT = ROOT / "src" / "gravity" / "egm96.rs"


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", "--max-time", "120", "-L", URL], check=True, capture_output=True
        ).stdout
    digest = hashlib.sha256(raw).hexdigest()
    if digest != SHA256:
        raise SystemExit(f"digest mismatch\n  expected {SHA256}\n  got      {digest}")
    print(f"digest ok: {digest}")
    return raw.decode()


def parse(text):
    """Fixed-width RECOEF records: degree in columns 6-10, order in 11-13."""
    cosine, sine, header = {}, {}, None
    for line in text.splitlines():
        if line.startswith("POTFIELD"):
            header = line.strip()
        if not line.startswith("RECOEF"):
            continue
        degree, order = int(line[6:11]), int(line[11:14])
        if degree > MAX_DEGREE:
            continue
        values = re.findall(r"[-+]?\d\.\d+E[-+]\d+", line[14:])
        if (degree, order) in cosine:
            raise SystemExit(f"duplicate coefficient at degree {degree} order {order}")
        cosine[(degree, order)] = float(values[0])
        sine[(degree, order)] = float(values[1]) if len(values) > 1 else 0.0

    # Degrees 0 and 1 are not in the file and are not free: the monopole term is
    # one by definition, and the degree-one terms vanish when the coordinate origin
    # is the center of mass, which is the frame these are defined in.
    cosine[(0, 0)], sine[(0, 0)] = 1.0, 0.0
    for order in (0, 1):
        cosine[(1, order)], sine[(1, order)] = 0.0, 0.0

    missing = [
        (n, m) for n in range(MAX_DEGREE + 1) for m in range(n + 1) if (n, m) not in cosine
    ]
    if missing:
        raise SystemExit(f"{len(missing)} coefficients missing, first {missing[0]}")
    return cosine, sine, header


def flatten(table):
    """Index order, matching gravity::index: degree * (degree + 1) / 2 + order."""
    return [
        table[(n, m)] for n in range(MAX_DEGREE + 1) for m in range(n + 1)
    ]


def rust_array(name, values, what):
    # Left alone by rustfmt: it would rewrap 65341 literals on every run, and the
    # layout here is already one that reads as a table.
    lines = [f"/// {what}", "#[rustfmt::skip]", f"static {name}: [f64; COUNT] = ["]
    for start in range(0, len(values), PER_LINE):
        chunk = ", ".join(f"{v!r}" for v in values[start : start + PER_LINE])
        lines.append(f"    {chunk},")
    lines.append("];")
    return "\n".join(lines)


def main():
    cosine, sine, header = parse(source_text(sys.argv))
    flat_c, flat_s = flatten(cosine), flatten(sine)
    count = (MAX_DEGREE + 1) * (MAX_DEGREE + 2) // 2
    assert len(flat_c) == count, (len(flat_c), count)

    # The file states its own gravitational parameter and reference radius, and they
    # are not interchangeable with a body's geodetic figures.
    fields = header.split()
    mu, radius = float(fields[2]), float(fields[3])

    body = f'''//! Fully normalized EGM96 spherical harmonic coefficients, to degree and order {MAX_DEGREE}.
//!
//! Generated by `scripts/generate-gravity-field.py`; do not edit by hand. Committed
//! rather than fetched so that a clone needs no network and no test has to skip.
//!
//! Source: `egm96_to360.ascii`, as redistributed in GMAT, downloaded {DOWNLOADED} from
//! <{URL}>
//! with SHA-256 `{SHA256}`.
//! The file's own header reads `{header}`.
//!
//! Checked on generation: degrees 2 to {MAX_DEGREE} complete with no duplicates, `S(n, 0)`
//! zero throughout, and `J2 = -sqrt(5) * C(2, 0)` recovering 1.0826267e-3.

use super::{{index, Coefficients}};

/// Highest degree and order these coefficients reach.
pub const MAX_DEGREE: usize = {MAX_DEGREE};

/// Number of (degree, order) pairs up to [`MAX_DEGREE`].
const COUNT: usize = (MAX_DEGREE + 1) * (MAX_DEGREE + 2) / 2;

/// Gravitational parameter the coefficients are defined against, m^3/s^2.
pub const MU: f64 = {mu!r};

/// Reference radius the coefficients are defined against, m.
///
/// Not a body's equatorial radius, and 0.7 m smaller than WGS84's. The expansion is
/// defined against this value, so using any other scales every term by the wrong
/// power of the ratio.
pub const REFERENCE_RADIUS: f64 = {radius!r};

/// The coefficients, for building a field.
pub fn coefficients() -> Coefficients<'static> {{
    Coefficients {{
        max_degree: MAX_DEGREE,
        cosine: &COSINE,
        sine: &SINE,
        mu: MU,
        reference_radius: REFERENCE_RADIUS,
    }}
}}

/// The normalized cosine coefficient of degree `n` and order `m`.
pub fn cosine(degree: usize, order: usize) -> f64 {{
    COSINE[index(degree, order)]
}}

/// The normalized sine coefficient of degree `n` and order `m`.
pub fn sine(degree: usize, order: usize) -> f64 {{
    SINE[index(degree, order)]
}}

{rust_array("COSINE", flat_c, "Normalized cosine coefficients, in index order.")}

{rust_array("SINE", flat_s, "Normalized sine coefficients, in index order.")}
'''
    OUT.write_text(body)
    size = OUT.stat().st_size
    print(f"wrote {OUT.relative_to(ROOT)}: {count} coefficients, {size / 1e6:.1f} MB")


if __name__ == "__main__":
    main()