import os
import symforce

symforce.set_symbolic_api("symengine")
symforce.set_log_level("warning")
symforce.set_epsilon_to_symbol()

import symforce.symbolic as se
from symforce import codegen
from symforce.codegen.backends.pytorch.pytorch_config import PyTorchConfig
from symforce.values import Values

def gen_py(inputs, outputs, name):
    gen = codegen.Codegen(
        inputs=inputs,
        outputs=outputs,
        config=PyTorchConfig(),
        name=name,
    )
    data = gen.generate_function()

    # Print what we generated
    print("Files generated in {}:\n".format(data.output_dir))
    for f in data.generated_files:
        print("  |- {}".format(os.path.relpath(f, data.output_dir)))

def main():
    x = se.Symbol('x')
    z = se.Symbol('z')
    y = se.Max(se.Scalar(0.0), (x**2 - z**2))
    w = se.sqrt(y)

    dwdx = [se.diff(w, x), se.diff(w, z)]

    inputs = Values(x=se.Matrix([x, z]))
    outputs = Values(dwdx=se.Matrix(dwdx))
    gen_py(inputs=inputs, outputs=outputs, name='dwdx')

if __name__=='__main__':
    main()