"""Import the nominal graph symbolically; CasADi performs AD."""
import casadi as ca

x = ca.MX.sym("x", 2)
u = ca.MX.sym("u")
F = ca.Function("F", [x, u], [ca.vertcat(x[0]+x[1]+u/2, x[1]+u)])

vibr = ca.GraphBuilder("vibr.onnx").create("vibr", {
    "symbolic": True, "is_diff_in": [True, False]})

opti = ca.Opti()
τ = opti.parameter()
opti.set_value(τ, 0)

N = 3
X = opti.variable(2, N+1)
U = opti.variable(1, N)
opti.minimize(ca.sumsqr(U))

for k in range(N):
    opti.subject_to(       X[:, k+1] == F(X[:, k], U[:, k]))
    opti.subject_to(-1 <= (U[:, k]   <= 1))
    opti.subject_to( 0 <= (X[1, k]   <= 1.5))
    opti.subject_to(vibr(X[:, k], τ) <= .48)

opti.subject_to(X[:, 0] == [0, 0])
opti.subject_to(X[:, N] == [1.8, 0])

opti.solver("ipopt")

sol = opti.solve()
print("Wetness 0, objective:", sol.value(opti.f))
print("First acceleration:", sol.value(U[0, 0]))

opti.set_value(τ, 1)
sol = opti.solve()
print("Wetness 1, objective:", sol.value(opti.f))
print("First acceleration:", sol.value(U[0, 0]))