check

#!/usr/bin/env python3
import contextlib
import subprocess
import sympy as sp
import tempfile
import re

RATRACER = "./ratracer"

def ratsame(a: str, b: str):
    a = sp.sympify(a)
    b = sp.sympify(b)
    return (a - b).ratsimp() == 0

def strsame(a: str, b: str):
    return a.strip() == b.strip()

@contextlib.contextmanager
def file(content: str = ""):
    with tempfile.NamedTemporaryFile("w") as f:
        f.write(content)
        f.flush()
        yield f.name

def run(*cmd):
    p = subprocess.run([RATRACER, *cmd], stdout=subprocess.PIPE, stderr=subprocess.PIPE, encoding="utf8")
    if p.returncode != 0:
        raise ValueError(f"Ratracer failed with code {p.returncode}\nStandard output:\n{p.stdout}\nStandard error:\n{p.stderr}")
    return p.stdout, p.stderr

def check_output_expr(expr: str, *cmd):
    stdout, stderr = run(*cmd)
    lines = stdout.splitlines()
    if len(lines) < 1:
        print("Expected:")
        print(expr)
        print("Received: nothing")
        raise ValueError(f"Test failed\nStandard error:\n{stderr}")
    result = lines[-1].replace(";", "").strip()
    if not ratsame(result, expr):
        print("Expected:")
        print(expr)
        print("Received:")
        print(result)
        raise ValueError(f"Test failed\nStandard error:\n{stderr}")

def check_output_str(expr: str, *cmd):
    stdout, stderr = run(*cmd)
    if stdout.strip() != expr.strip():
        print("Expected:")
        print(expr)
        print("Received:")
        print(stdout)
        raise ValueError(f"Test failed\nStandard error:\n{stderr}")

def check_log_pattern(pattern: str, *cmd):
    stdout, stderr = run(*cmd)
    if re.search(pattern, stderr, flags=re.MULTILINE) is None:
        print("Expected to find:")
        print(pattern)
        print("Received:")
        print(stdout)
        print("Log:")
        print(stderr)
        raise ValueError(f"Test failed")

def check_trace_output(expr: str, *cmd):
    with file(expr) as fname:
        check_output_expr(expr, "trace-expression", fname, *cmd)

def check_series(expr: str):
    eps = sp.sympify("eps")
    with file(expr) as fn:
        stdout, stderr = run("trace-expression", fn, "to-series", "eps", "0", "reconstruct0")
    series = {}
    order = None
    for line in stdout.splitlines():
        line = line.rstrip()
        if line == "": continue
        if line.startswith("Fire"): continue
        m = re.match("^ORDER.*,eps\\^([0-9-]*)", line)
        if m is not None:
            order = eps**int(m.group(1))
            continue
        m = re.match("^  (.*);", line)
        if m is not None:
            if order in series: raise ValueError(f"duplicate order '{order}'")
            series[order] = line.strip().replace(";", "")
            continue
        raise ValueError(f"Unrecognized output: {line}")
    spseries = sp.series(expr, sp.sympify("eps"), n=25).removeO().as_coefficients_dict()
    for key, value in series.items():
        if key not in spseries and sp.sympify(value) != 0:
            raise ValueError(f"The reported key '{key}' should not appear.\nValue: {value}")
        if not (sp.sympify(value) - spseries[key]).together().is_zero:
            raise ValueError(
                    f"Mismatch in the value at '{key}'.\n"
                    f"Expected:\n{value}\n"
                    f"Got:\n{series[key]}")

check_output_expr("123456", "sh", "echo 123456")
check_trace_output("2*y/(x^2-y^2) + 1/(x+y) + 1/(x-y)", "reconstruct")
check_trace_output("2*y/(x^2-y^2) + 1/(x+y) + 1/(x-y)", "optimize", "reconstruct")
check_trace_output("2*y/(x^2-y^2) + 1/(x+y) + 1/(x-y)", "finalize", "reconstruct")
check_trace_output("2*y/(x^2-y^2) + 1/(x+y) + 1/(x-y)", "optimize", "finalize", "reconstruct")
check_trace_output("1/x^-1", "reconstruct")
check_trace_output("1/x^(-2)", "reconstruct")
check_trace_output(" 1 / x ^ ( -2 ) ", "reconstruct")
check_trace_output("1/x+1/y^2-1/x^-10+2", "reconstruct")
check_trace_output("(x-y)^-2", "reconstruct")
check_trace_output("x+1/2*3/4", "reconstruct")
check_trace_output("x*2147483647 + y*2147483648 + z*2147483649", "optimize", "finalize", "reconstruct")
check_trace_output("x*4294967295 + y*4294967296 + z*4294967297", "optimize", "finalize", "reconstruct")
check_trace_output("x*8589934591 + y*8589934592 + z*8589934593", "optimize", "finalize", "reconstruct")
check_trace_output("a + _a + a_ + C0 + C0_a + C_a0", "finalize", "reconstruct")

with file("1+2") as fn:
    check_output_expr("3", "trace-expression", fn, "finalize", "trace-expression", fn, "reconstruct0")

expr = "2*y/(x^2-y^2) + 1/(x+y) + 1/(x-y)"
with file(expr) as fn1:
    with file() as fn2:
        run("trace-expression", fn1, "save-trace", fn2)
        check_output_expr(expr, "load-trace", fn2, "reconstruct")
    with file() as fn2:
        run("trace-expression", fn1, "optimize", "save-trace", fn2)
        check_output_expr(expr, "load-trace", fn2, "reconstruct")
    with file() as fn2:
        run("trace-expression", fn1, "finalize", "save-trace", fn2)
        check_output_expr(expr, "load-trace", fn2, "reconstruct")
    with file() as fn2:
        run("trace-expression", fn1, "optimize", "finalize", "save-trace", fn2)
        check_output_expr(expr, "load-trace", fn2, "reconstruct")

with file("x+y/x^2") as fn:
    check_output_expr("11+y/121", "set", "x", "11", "trace-expression", fn, "reconstruct")
    check_output_expr("11+y/121", "set", "x", "11", "trace-expression", fn, "optimize", "reconstruct")
    check_output_expr("x+y/x^2", "trace-expression", fn, "set", "x", "11", "reconstruct")

with file("x+y/2+z/3+t/4") as fn:
    check_output_expr("x+5/2+7/3+11/4",
            "set", "y", "5", "set", "z", "7", "set", "t", "11",
            "trace-expression", fn, "reconstruct")

with file("x") as fn:
    check_output_expr("x+1", "set", "x", "x+1", "trace-expression", fn, "reconstruct")

with file("(2*2)+(2*x)+(2*x)+(2*x)+(2*x)") as fn1:
    with file("x") as fn2:
        check_output_expr("x+11", "trace-expression", fn1, "set", "x", "x+11", "optimize", "finalize", "unfinalize", "trace-expression", fn2, "reconstruct")

with file("1-(x-2)/(y+1)+(x-2)^2*(y+1)") as fn1:
    with file() as fn2:
        run("trace-expression", fn1, "optimize", "finalize", "save-trace", fn2)
        check_output_expr("1-(x-2)/(y+1)+(x-2)^2*(y+1)", "load-trace", fn2, "reconstruct")
        check_output_expr("1-2/(y+1)+4*(y+1)", "set", "x", "4", "load-trace", fn2, "reconstruct")

check_series("1/(1-eps)+1")
check_series("1+1/(eps*eps)")
check_series("1+1/eps")
for i in range(10):
    check_series(f"1+1/eps^{i}")
for i in range(10):
    check_series(f"1/eps+eps^{i}")
for i in range(10):
    check_series(f"1/(1-eps)+1/(1-eps)/eps^{i}")
check_series("1/eps^10-1/eps^10+ 1/(1-eps)")
check_series("1/(1-eps)+1")
check_series("1/(1-eps)-1-eps+1/(eps*eps)")
check_series("1/eps^9 + 1/eps")
check_series("1/(1-eps)/eps^7")
check_series("1/(eps+eps*eps+1-1)")
check_series("1/(eps+1)/(eps+1-1)")
check_series("1/eps^6+(2/(3*eps+5)/(7*eps+2-2)+1/eps)/eps")
check_series("1/(1-eps)-1/(1-eps)")
check_series("(1/(1-eps) + eps^2)-1/(1-eps)")
check_series("(1+eps+eps^2+eps^3+eps^4/(1-eps)) - (1+eps+eps^2+eps^4/(1-eps))")
check_series("(1+eps-1-eps)*(1/(1-eps))+1")
check_series("(1+eps-1-eps)*(1/(1-eps))")

system = """\
fam[3]*(-z)
fam[2]*(q)

fam[2]*(-y)
fam[1]*(x)

fam[1]*(p)
fam2@123*(-q)
"""
result = """\
CO[fam2@123,fam[1]] =
  (p)/(q);
CO[fam[3],fam[1]] =
  (q*x)/(z*y);
CO[fam[2],fam[1]] =
  (x)/(y);
"""
with file(system) as fn:
    check_output_str(
        result,
        "load-equations",
        fn,
        "solve-equations",
        "choose-equation-outputs",
        "optimize",
        "finalize",
        "reconstruct"
    )

system = """\
fam[1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19,20]*(x)
fam[2,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19,20]*(-y)
"""
result = """\
CO[fam[2,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19,20],fam[1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19,20]] =
  (x)/(y);
"""
with file(system) as fn:
    check_output_str(
        result,
        "load-equations",
        fn,
        "solve-equations",
        "choose-equation-outputs",
        "optimize",
        "finalize",
        "reconstruct"
    )

system = """\
fam[1]*(x)
fam[2]*(y)
fam[3]*(z)
"""
result = """\
fam[3]*0x345
fam[2]*0x234
fam[1]*0x123
"""
with file(system) as f:
    check_output_str(
        result,
        "set", "x", str(0x123),
        "set", "y", str(0x234),
        "set", "z", str(0x345),
        "load-equations", f,
        "dump-equations")

expr = "x+y"
with file(expr) as fn:
    check_output_str(
        f"{fn} =\n  357;",
        "set", "x", "123",
        "trace-expression",
        fn,
        "evaluate-modular", "--set", "y", "234", "--modulus=1000")

with file("x") as fn, file("0 new\n") as renames:
    check_output_str("0 x\n", "trace-expression", fn, "list-inputs")
    check_output_str(f"0 {fn}\n", "trace-expression", fn, "list-outputs")
    check_output_str("0 new\n",
        "trace-expression", fn,
        "rename-outputs", renames,
        "list-outputs")

with file("x") as fn1, file("x") as fn2:
    check_output_str(
        f"{fn1} =\n  (10)/(1);\n{fn2} =\n  (x)/(1);",
        "set", "x", "10",
        "trace-expression", fn1,
        "unset", "x",
        "trace-expression", fn2,
        "reconstruct")

with file("x") as fn1, \
     file("y") as fn2, \
     file(f"{fn1}\n") as pat1, \
     file(f"{fn2}\n") as pat2:
    check_output_str(
        f"0 {fn1}\n",
        "trace-expression", fn1,
        "trace-expression", fn2,
        "keep-outputs", pat1,
        "list-outputs")
    check_output_str(
        f"0 {fn1}\n",
        "trace-expression", fn1,
        "trace-expression", fn2,
        "drop-outputs", pat2,
        "list-outputs")

system = """\
fam[1]*(24)
fam[2]*(x1+x2)
fam[3]*(x1+x3)
"""
result1 = """\
fam[3]*0x20
fam[2]*0x1e
fam[1]*0x18

CO[fam[3],fam[2]] =
  ((-1)*x2+(-1)*x1)/(x1+x3);
CO[fam[3],fam[1]] =
  (1)/((-1/24)*x1+(-1/24)*x3);
"""
result2 = """\
fam[3]*0x20
fam[1]*0x18

CO[fam[3],fam[1]] =
  (1)/((-1/24)*x1+(-1/24)*x3);
"""
result3 = """\
fam[2]*0x1e
fam[1]*0x18

CO[fam[2],fam[1]] =
  (1)/((-1/24)*x2+(-1/24)*x1);
"""
with file(system) as f:
    check_output_str(
        result1,
        "configure-tracing",
            "--modulus=1013",
            "--set", "x1", "13",
            "--set", "x2", "17",
            "--set", "x3", "19",
        "load-equations", f,
        "dump-equations",
        "solve-equations",
        "choose-equation-outputs",
        "reconstruct")
    check_output_str(
        result2,
        "configure-tracing",
            "--modulus=101",
            "--set", "x1", "13",
            "--set", "x2", "-13",
            "--set", "x3", "19",
        "load-equations", f,
        "dump-equations",
        "solve-equations",
        "choose-equation-outputs",
        "reconstruct")
    check_output_str(
        result3,
        "configure-tracing",
            "--modulus=101",
            "--set", "x1", "13",
            "--set", "x2", "17",
            "--set", "x3", "-13",
        "load-equations", f,
        "dump-equations",
        "solve-equations",
        "choose-equation-outputs",
        "reconstruct")

print("OK")