#!/usr/bin/env python3
# SPDX-FileCopyrightText: 2026 VTT Technical Research Centre of Finland Ltd
# SPDX-License-Identifier: AGPL-3.0-or-later
"""Generate one bounded two-field model, never parse arbitrary input code."""
from __future__ import annotations

import argparse
import hashlib
import json
from pathlib import Path

import sympy as sp

VERSION = "1.14.0"


def expressions():
    names = "t x y z source_rate a b c d v_value u_xx u_yy u_zz u_xy u_yz v_xz"
    s = dict(zip(names.split(), sp.symbols(names)))
    source = s["source_rate"] * s["t"] * (s["x"] + s["y"] + s["z"])
    return s, [s["v_value"] + s["a"] * s["u_xy"] + source,
               s["b"] * (s["u_xx"] + s["u_yy"] + s["u_zz"])
               + s["c"] * s["v_xz"] + s["d"] * s["u_yz"] + source]


def cpp(expression):
    # Fixed expressions only contain binary arithmetic; keep repository style
    # deterministic without introducing a formatter dependency.
    code = sp.ccode(expression, standard="c99")
    return " ".join(code.replace("*", " * ").replace("+", " + ").split())


def artifacts():
    if sp.__version__ != VERSION:
        raise RuntimeError(f"require SymPy {VERSION}, got {sp.__version__}")
    symbols, expr = expressions()
    replacements, reduced = sp.cse(expr, symbols=sp.numbered_symbols("cse"), order="canonical")
    bindings = {name: (f"g.{name[0]}.{name[2:]}" if name.startswith(("u_", "v_"))
                      else f"p.{name}" if name in "abcd" else
                      f"position[{ 'xyz'.index(name) }]" if name in "xyz" else name)
                for name in symbols}
    # The generated header carries its own SPDX tags; REUSE must not parse them
    # as this file's license expression.
    # REUSE-IgnoreStart
    lines = ["// SPDX-FileCopyrightText: 2026 VTT Technical Research Centre of Finland Ltd",
             "// SPDX-License-Identifier: AGPL-3.0-or-later", "",
             # REUSE-IgnoreEnd
             "// Generated by generate.py with SymPy 1.14.0. Do not edit.",
             "#pragma once", "#include <openpfc/kernel/data/host_device.hpp>", "",
             "namespace coupled::generated {",
             "template <class Increments, class Position, class Local, class Parameters>",
             "OPENPFC_INLINE_HD Increments rhs(double t, const Position &position, const Local &g,",
             "                                 const Parameters &p,",
             "                                 double source_rate = 0.0) noexcept {"]
    for name in sorted(symbols):
        if bindings[name] != name:
            lines.append(f"  const double {name} = {bindings[name]};")
    for var, value in replacements:
        lines.append(f"  const double {var} = {cpp(value)};")
    lines.extend([f"  return Increments{{{cpp(reduced[0])},",
                  f"                    {cpp(reduced[1])}}};",
                  "}", "} // namespace coupled::generated", ""])
    header = "\n".join(lines).encode()
    manifest = {
        "schema": 1, "sympy": VERSION,
        "header_sha256": hashlib.sha256(header).hexdigest(),
        "fields": {field: sorted({str(symbol)[2:] for e in expr for symbol in e.free_symbols
                                  if str(symbol).startswith(field + "_")})
                   for field in ["u", "v"]},
        "parameters": ["a", "b", "c", "d"], "position": ["x", "y", "z"],
        "time": True, "optional_source": "source_rate*t*(x+y+z), added to both equations; default zero",
        "expressions": [str(e) for e in expr],
        "operations_before_cse": sum(int(sp.count_ops(e)) for e in expr),
        "operations_after_cse": sum(int(sp.count_ops(e)) for _, e in replacements)
                                + sum(int(sp.count_ops(e)) for e in reduced),
        "cse_temporaries": len(replacements),
    }
    return {"generated_rhs.hpp": header,
            "dependency_manifest.json": (json.dumps(manifest, indent=2, sort_keys=True)+"\n").encode()}


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--output-dir", type=Path, default=Path(__file__).resolve().parent)
    parser.add_argument("--check", action="store_true", help="fail if checked-in artifacts differ")
    args = parser.parse_args()
    output = artifacts()
    if args.check:
        for name, content in output.items():
            if not (args.output_dir / name).is_file() or (args.output_dir / name).read_bytes() != content:
                parser.exit(1, f"generated artifact differs: {name}\n")
    else:
        args.output_dir.mkdir(parents=True, exist_ok=True)
        for name, content in output.items():
            (args.output_dir / name).write_bytes(content)


if __name__ == "__main__":
    main()