#!/usr/bin/env python3
"""Generate and verify the explicit operator formulas.

The generated block is distributed both as operators.tex and inline in the
standalone manuscript source between fixed marker comments.
"""
from __future__ import annotations

import argparse
from pathlib import Path

from sympy import Add, Matrix, Rational, expand, factor, latex, symbols

x, y, z = symbols("x y z")
u = 1 + x * y
F = Matrix([
    u**3 * z + y**2 * u * (4 + 3 * x * y),
    y + 3 * x * u**2 * z + 3 * x * y**2 * (4 + 3 * x * y),
    2 * x - 3 * x**2 * y - x**3 * z,
])
J = F.jacobian(Matrix([x, y, z]))
A = J.adjugate().applyfunc(expand) * Rational(-1, 2)


def tex(expr) -> str:
    return latex(factor(expr))


def render() -> str:
    lines: list[str] = []
    lines.append(r"\begin{align}")
    for i, expr in enumerate(F, 1):
        ending = r" \\" if i < 3 else ""
        lines.append(rf"\varphi(x_{i}) &= {tex(expr)}{ending}")
    lines.append(r"\end{align}")
    lines.append("")
    lines.append(r"\begingroup\small")

    order = [(0, 0), (1, 0), (2, 0), (0, 1), (1, 1), (2, 1), (0, 2), (1, 2)]
    for row, col in order:
        lines.append(r"\begin{align}")
        lines.append(rf"a_{{{row+1}{col+1}}} &= {tex(A[row, col])}\\[-2pt]\nonumber")
        lines.append(r"\end{align}")

    # The last coefficient is split across lines only for typesetting.
    doubled = expand(2 * A[2, 2])
    terms = list(Add.make_args(doubled))
    # Sort by SymPy's conventional polynomial display order through the LaTeX string.
    full = latex(doubled)
    pieces = full.split(" + ")
    if len(pieces) != 14:
        raise RuntimeError(f"Unexpected a33 term count: {len(pieces)}")
    groups = [pieces[:4], pieces[4:9], pieces[9:]]
    lines.append(r"\begin{align}")
    lines.append(r"a_{33} &= \tfrac{1}{2}\big( " + " +".join(groups[0]) + r" \\")
    lines.append(r"&\qquad +" + " +".join(groups[1]) + r" \\")
    lines.append(r"&\qquad +" + " +".join(groups[2]) + r" \big)\nonumber")
    lines.append(r"\end{align}")
    lines.append(r"\endgroup")
    lines.append("")
    return "\n".join(lines)


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--output", type=Path, default=Path(__file__).resolve().parent / "operators.tex")
    parser.add_argument("--check", action="store_true")
    parser.add_argument("--check-source", type=Path)
    args = parser.parse_args()
    generated = render()
    if args.check:
        if not args.output.exists():
            raise SystemExit(f"missing {args.output}")
        current = args.output.read_text()
        if current != generated:
            raise SystemExit(f"{args.output.name} does not match generated formulas")
        print(f"PASSED: {args.output.name} matches generated formulas")

    if args.check_source is not None:
        if not args.check_source.exists():
            raise SystemExit(f"missing {args.check_source}")
        source = args.check_source.read_text()
        begin = "% BEGIN GENERATED OPERATORS\n"
        end = "% END GENERATED OPERATORS"
        if begin not in source or end not in source:
            raise SystemExit(f"generated-operator markers missing from {args.check_source.name}")
        inline = source.split(begin, 1)[1].split(end, 1)[0]
        if inline != generated:
            raise SystemExit(f"inline formulas in {args.check_source.name} do not match generated formulas")
        print(f"PASSED: inline formulas in {args.check_source.name} match generated formulas")

    if args.check or args.check_source is not None:
        return 0

    args.output.write_text(generated)
    print(f"Wrote {args.output}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())