wadamesh/scripts/wmm/gen_lua.py

#!/usr/bin/env python3
"""Generate the WMM declination block that lives inside a Lua Store app.

The block is generated rather than hand-pasted for two reasons. It is four
kilobytes of coefficients that nobody can eyeball, and the wrapper around it
matters: the coefficient tables are named G/H/GD/HD, and an app that happens
to use a global named H (gpscompass uses it for the screen height) has that
global silently eaten by a module-scope `local H`. That actually happened.
Owning the `do ... end` here means it cannot come unwrapped again.

    gen_lua.py --emit                      print the block
    gen_lua.py --update <app.lua>          rewrite the block in place
    gen_lua.py --check  <app.lua>          fail if the app has drifted

--check is the one to wire into CI: it regenerates from WMM.COF and compares,
so a coefficient file that gets updated without regenerating, or a block that
gets edited by hand, is caught rather than shipped.

Verify the maths itself with verify.py (NOAA's own 100 test values), and the
generated Lua with the `declination` scenario in scripts/lua-harness.
"""
import argparse
import os
import sys

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from wmm import load_cof

HERE = os.path.dirname(os.path.abspath(__file__))
COF = os.path.join(HERE, "WMM.COF")

BEGIN = "-- WMM-GEN BEGIN -- generated by scripts/wmm/gen_lua.py; do not edit by hand"
END = "end -- WMM-GEN END"


def fmt(v):
    """Shortest round-trip-safe spelling: the coefficients are one decimal."""
    s = f"{v:.1f}"
    if s.endswith(".0"): s = s[:-2]
    if s == "-0": s = "0"
    if s.startswith("0."): s = s[1:]        # 0.2 -> .2
    if s.startswith("-0."): s = "-" + s[2:]  # -0.2 -> -.2
    return s


def _arr(name, vals):
    body, line = [], ""
    for v in vals:
        if len(line) + len(v) + 1 > 96:
            body.append(line)
            line = ""
        line += v + ","
    if line:
        body.append(line)
    return f"local {name}={{" + "\n".join(body) + "}\n"


def gen_block(cof=COF):
    """The complete Lua block, markers included, ready to sit in an app."""
    epoch, N, g, h, gd, hd = load_cof(cof)
    G, H, GD, HD = [], [], [], []
    for n in range(1, N + 1):
        for m in range(0, n + 1):
            G.append(fmt(g[(n, m)]))
            GD.append(fmt(gd[(n, m)]))
            if m >= 1:
                H.append(fmt(h[(n, m)]))
                HD.append(fmt(hd[(n, m)]))

    with open(os.path.join(HERE, "wmm_block.lua"), encoding="utf-8") as f:
        maths = f.read()

    head = (
        f"{BEGIN}\n"
        f"-- WMM{int(epoch)} magnetic declination, degree {N}. East positive:\n"
        f"--     true bearing = magnetic bearing + declination\n"
        f"-- Data: NOAA/NCEI WMM.COF epoch {epoch}, valid {int(epoch)}.0-{int(epoch) + 5}.0.\n"
        f"-- Source: https://www.ncei.noaa.gov/sites/default/files/2024-12/WMM2025COF.zip\n"
        f"-- Everything below is scoped: the tables are named G/H/GD/HD and would\n"
        f"-- otherwise shadow an app's own globals of those names.\n"
        f"local declination\n"
        f"do\n"
        f"local EPOCH,NMAX={epoch},{N}\n"
    )
    arrays = _arr("G", G) + _arr("H", H) + _arr("GD", GD) + _arr("HD", HD)
    return head + arrays + maths + END + "\n"


def _split(path):
    """(before, block, after) around the markers in `path`."""
    src = open(path, encoding="utf-8").read()
    a = src.find(BEGIN)
    if a < 0:
        sys.exit(f"{path}: no '{BEGIN}' marker")
    b = src.find(END, a)
    if b < 0:
        sys.exit(f"{path}: BEGIN marker has no matching '{END}'")
    return src[:a], src[a:b + len(END) + 1], src[b + len(END) + 1:]


def main():
    ap = argparse.ArgumentParser(description=__doc__,
                                 formatter_class=argparse.RawDescriptionHelpFormatter)
    g = ap.add_mutually_exclusive_group(required=True)
    g.add_argument("--emit", action="store_true", help="print the block to stdout")
    g.add_argument("--update", metavar="APP", help="rewrite the block inside APP")
    g.add_argument("--check", metavar="APP", help="exit 1 if APP has drifted")
    ap.add_argument("--cof", default=COF, help="coefficient file (default: WMM.COF)")
    args = ap.parse_args()

    block = gen_block(args.cof)

    if args.emit:
        sys.stdout.write(block)
        return 0

    path = args.update or args.check
    before, current, after = _split(path)

    if args.check:
        if current == block:
            print(f"{os.path.basename(path)}: up to date ({len(block)} bytes)")
            return 0
        print(f"{os.path.basename(path)}: the generated block has DRIFTED from "
              f"{os.path.basename(args.cof)}", file=sys.stderr)
        import difflib
        diff = difflib.unified_diff(current.splitlines(), block.splitlines(),
                                    "in app", "generated", lineterm="", n=1)
        for i, line in enumerate(diff):
            if i > 40:
                print("  ...", file=sys.stderr)
                break
            print("  " + line, file=sys.stderr)
        print("  run: scripts/wmm/gen_lua.py --update " + path, file=sys.stderr)
        return 1

    if current == block:
        print(f"{os.path.basename(path)}: already up to date")
        return 0
    open(path, "w", encoding="utf-8").write(before + block + after)
    print(f"{os.path.basename(path)}: block updated ({len(block)} bytes)")
    return 0


if __name__ == "__main__":
    sys.exit(main())