from pydrake.solvers import MathematicalProgram, Solve
import numpy as np

# Create an empty MathematicalProgram named prog (with no decision variables,
# constraints or costs)
prog = MathematicalProgram()
# Add two decision variables x[0], x[1].
x = prog.NewContinuousVariables(2, "x")# Add a symbolic linear expression as the cost.
cost1 = prog.AddLinearCost(x[0] + 3 * x[1] + 2)
# Print the newly added cost
print(cost1)
# The newly added cost is stored in prog.linear_costs().
print(prog.linear_costs()[0])cost2 = prog.AddLinearCost(2 * x[1] + 3)
print(f"number of linear cost objects: {len(prog.linear_costs())}")# We add a linear cost 3 * x[0] + 4 * x[1] + 5 to prog by specifying the coefficients
# [3., 4] and the constant 5 in AddLinearCost
cost3 = prog.AddLinearCost([3., 4.], 5., x)
print(cost3)print(f"number of linear cost objects before calling AddCost: {len(prog.linear_costs())}")
# Call AddCost to add a linear expression as linear cost. After calling this function,
# len(prog.linear_costs()) will increase by 1.
cost4 = prog.AddCost(x[0] + 3 * x[1] + 5)
print(f"number of linear cost objects after calling AddCost: {len(prog.linear_costs())}")prog = MathematicalProgram()
x = prog.NewContinuousVariables(2, "x")
y = prog.NewContinuousVariables(3, "y")

# Call AddConstraint to add a bounding box constraint x[0] >= 1
bounding_box1 = prog.AddConstraint(x[0] >= 1)
print(f"number of bounding box constraint objects: {len(prog.bounding_box_constraints())}")

# Call AddLinearConstraint to add a bounding box constraint x[1] <= 2
bounding_box2 = prog.AddLinearConstraint(x[1] <= 2)
print(f"number of bounding box constraint objects: {len(prog.bounding_box_constraints())}")

# Call AddConstraint to add a linear equality constraint x[0] + y[1] == 3
linear_eq1 = prog.AddConstraint(x[0] + y[1] == 3.)
print(f"number of linear equality constraint objects: {len(prog.linear_equality_constraints())}")

# Call AddLinearConstraint to add a linear equality constraint x[1] + 2 * y[2] == 1
linear_eq2 = prog.AddLinearConstraint(x[1] + 2 * y[2] == 1)
print(f"number of linear equality constraint objects: {len(prog.linear_equality_constraints())}")

# Call AddConstraint to add a linear inequality constraint x[0] + 3*x[1] + 2*y[2] <= 4
linear_ineq1 = prog.AddConstraint(x[0] + 3*x[1] + 2*y[2] <= 4)
print(f"number of linear inequality constraint objects: {len(prog.linear_constraints())}")

# Call AddLinearConstraint to add a linear inequality constraint x[1] + 4 * y[1] >= 2
linear_ineq2 = prog.AddLinearConstraint(x[1] + 4 * y[1] >= 2)
print(f"number of linear inequality constraint objects: {len(prog.linear_constraints())}")# Add a nonlinear constraint square(x[0]) == 2 by calling AddLinearConstraint. This should
# throw an exception
try:
    prog.AddLinearConstraint(x[0] ** 2 == 2)
except RuntimeError as err:
    print(err.args)# Add a linear constraint 2x[0] + 3x[1] <= 2, 1 <= 4x[1] + 5y[2] <= 3.
# This is equivalent to lower <= A * [x;y[2]] <= upper with
# lower = [-inf, 1], upper = [2, 3], A = [[2, 3, 0], [0, 4, 5]].
linear_constraint = prog.AddLinearConstraint(
    A=[[2., 3., 0], [0., 4., 5.]],
    lb=[-np.inf, 1],
    ub=[2., 3.],
    vars=np.hstack((x, y[2])))
print(linear_constraint)# Add a bounding box constraint -1 <= x[0] <= 2, 3 <= x[1] <= 5
bounding_box3 = prog.AddBoundingBoxConstraint([-1, 3], [2, 5], x)
print(bounding_box3)# Add a bounding box constraint 3 <= y[i] <= 5 for all i.
bounding_box4 = prog.AddBoundingBoxConstraint(3, 5, y)
print(bounding_box4)# Add a linear equality constraint 4 * x[0] + 5 * x[1] == 1
linear_eq3 = prog.AddLinearEqualityConstraint(np.array([[4, 5]]), np.array([1]), x)
print(linear_eq3)# Solve an optimization program
# min -3x[0] - x[1] - 5x[2] -x[3] + 2
# s.t 3x[0] + x[1] + 2x[2] = 30
#     2x[0] + x[1] + 3x[2] + x[3] >= 15
#     2x[1] + 3x[3] <= 25
#     -100 <= x[0] + 2x[2] <= 40
#   x[0], x[1], x[2], x[3] >= 0, x[1] <= 10
prog = MathematicalProgram()
# Declare x as decision variables.
x = prog.NewContinuousVariables(4)
# Add linear costs. To show that calling AddLinearCosts results in the sum of each individual
# cost, we add two costs -3x[0] - x[1] and -5x[2]-x[3]+2
prog.AddLinearCost(-3*x[0] -x[1])
prog.AddLinearCost(-5*x[2] - x[3] + 2)
# Add linear equality constraint 3x[0] + x[1] + 2x[2] == 30
prog.AddLinearConstraint(3*x[0] + x[1] + 2*x[2] == 30)
# Add Linear inequality constraints
prog.AddLinearConstraint(2*x[0] + x[1] + 3*x[2] + x[3] >= 15)
prog.AddLinearConstraint(2*x[1] + 3*x[3] <= 25)
# Add linear inequality constraint -100 <= x[0] + 2x[2] <= 40
prog.AddLinearConstraint(A=[[1., 2.]], lb=[-100], ub=[40], vars=[x[0], x[2]])
prog.AddBoundingBoxConstraint(0, np.inf, x)
prog.AddLinearConstraint(x[1] <= 10)

# Now solve the program.
result = Solve(prog)
print(f"Is solved successfully: {result.is_success()}")
print(f"x optimal value: {result.GetSolution(x)}")
print(f"optimal cost: {result.get_optimal_cost()}")