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 quadratic expression as the quadratic cost.
cost1 = prog.AddQuadraticCost(x[0]**2 + 2*x[0]*x[1] + x[1]**2 + 3*x[0] + 4)
# The newly added cost is returned as cost1
print(cost1)
# The newly added cost is stored inside prog.
print(prog.quadratic_costs()[0])# Add another quadratic cost to prog.
cost2 = prog.AddQuadraticCost(x[1]*x[1] + 3)
print(f"The number of quadratic costs in prog: {len(prog.quadratic_costs())}")# Add the cost x[0]*x[0] + x[0]*x[1] + 1.5*x[1]*x[1] + 2*x[0] + 4*x[1] + 1
cost3 = prog.AddQuadraticCost(Q=[[2, 1], [1, 3]], b=[2, 4], c=1, vars=x)
print(f"cost 3 is {cost3}")# Adds the cost (x - [1;2])' * Q * (x-[1;2])
cost4 = prog.AddQuadraticErrorCost(Q=[[1, 2],[2, 6]], x_desired=[1,2], vars=x)
print(f"cost4 is {cost4}")# Adds the squared norm of (x[0]+2*x[1]-2, x[1] - 3) to the program cost.
cost5 = prog.Add2NormSquaredCost(A=[[1, 2], [0, 1]], b=[2, 3], vars=x)
print(f"cost5 is {cost5}")# Adds a linear cost to the quadratic program
cost6 = prog.AddLinearCost(x[0] + 3 * x[1] + 1)
print(f"cost6 is {cost6}")
print(f"Number of linear costs in prog: {len(prog.linear_costs())}")constraint1 = prog.AddLinearConstraint(x[0] + 3*x[1] <= 5)
# Adds the constraint 1 <= x[0] <= 5 and 1 <= x[1] <= 5
constraint2 = prog.AddBoundingBoxConstraint(1, 5, x)# Call AddQuadraticCost with the is_convex flag unspecified. Drake will check the convexity
# of this cost.
cost7 = prog.AddQuadraticCost(x[0]**2 + 3 * x[1]**2 + x[0] * x[1] + 2 * x[0])
print(f"Is the cost {cost7} convex? {cost7.evaluator().is_convex()}")
cost8 = prog.AddQuadraticCost(x[0]**2 + 3 * x[1]**2 + 4*x[0]*x[1])
print(f"Is the cost {cost8} convex? {cost8.evaluator().is_convex()}")

# Call AddQuadraticCost with specified is_convex flag. Drake won't check the convexity of
# the quadratic cost when this flag is specified.
cost9 = prog.AddQuadraticCost(x[0]**2 + 4 * x[1]**2 + 3 * x[0]*x[1], is_convex=True)
print(f"Is the cost {cost9} convex? {cost9.evaluator().is_convex()}")prog = MathematicalProgram()
x = prog.NewContinuousVariables(3, "x")
prog.AddQuadraticCost(x[0] * x[0] + 2 * x[0] + 3)
# Adds the quadratic cost on the squared norm of the vector
# (x[1] + 3*x[2] - 1, 2*x[1] + 4*x[2] -4)
prog.Add2NormSquaredCost(A = [[1, 3], [2, 4]], b=[1, 4], vars=[x[1], x[2]])

# Adds the linear constraints.
prog.AddLinearEqualityConstraint(x[0] + 2*x[1] == 5)
prog.AddLinearConstraint(x[0] + 4 *x[1] <= 10)
# Sets the bounds for each variable to be within [-1, 10]
prog.AddBoundingBoxConstraint(-1, 10, x)

# Solve the program.
result = Solve(prog)
print(f"optimal solution x: {result.GetSolution(x)}")
print(f"optimal cost: {result.get_optimal_cost()}")