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 three decision variables x[0], x[1], x[2]
x = prog.NewContinuousVariables(3, "x")def cost_fun(z):
    cos_z = np.cos(z[0] + z[1])
    sin_z = np.sin(z[0] + z[1])
    return cos_z**2 + cos_z + sin_z
# Add the cost evaluated with x[0] and x[1].
cost1 = prog.AddCost(cost_fun, vars=[x[0], x[1]])
print(cost1)cost2 = prog.AddCost(cost_fun, vars=[x[0], x[2]])
print(cost2)# Add a cost x[1]**2 + x[0] using a lambda function.
cost3 = prog.AddCost(lambda z: z[0]**2 + z[1], vars = [x[1], x[0]])
print(cost3)cost4 = prog.AddCost(lambda z: z[0]**2 + z[1], vars = x[1:])
print(cost4)cost4 = prog.AddQuadraticCost(x[0]**2 + 3 * x[1]**2 + 2*x[0]*x[1] + 2*x[1] * x[0] + 1)
print(cost4)# Add a cost x[0]**2 + 2*x[1]**2 + x[0]*x[1] + 3*x[1] + 1.
cost5 = prog.AddQuadraticCost(
    Q=np.array([[2., 1], [1., 4.]]),
    b=np.array([0., 3.]),
    c=1.,
    vars=x[:2])
print(cost5)cost6 = prog.AddQuadraticErrorCost(
    Q=np.array([[1, 0.5], [0.5, 1]]),
    x_desired=np.array([1., 2.]),
    vars=x[1:])
print(cost6)# Add the L2 norm cost on (A*x[:2] - b).dot(A*x[:2]-b)
cost7 = prog.Add2NormSquaredCost(
    A=np.array([[1., 2.], [2., 3], [3., 4]]),
    b=np.array([2, 3, 1.]),
    vars=x[:2])
print(cost7)## Define a python function to add the constraint x[0]**2 + 2x[1]<=1, -0.5<=sin(x[1])<=0.5
def constraint_evaluator1(z):
    return np.array([z[0]**2+2*z[1], np.sin(z[1])])

constraint1 = prog.AddConstraint(
    constraint_evaluator1,
    lb=np.array([-np.inf, -0.5]),
    ub=np.array([1., 0.5]),
    vars=x[:2])
print(constraint1)

# Add another constraint using lambda function.
constraint2 = prog.AddConstraint(
    lambda z: np.array([z[0]*z[1]]),
    lb=[0.],
    ub=[1.],
    vars=[x[2]])
print(constraint2)## Solve a simple nonlinear 
# min               -x0
# subject to x1 - exp(x0) >= 0
#            x2 - exp(x1) >= 0
#            0 <= x0 <= 100
#            0 <= x1 <= 100
#            0 <= x2 <= 10
prog = MathematicalProgram()
x = prog.NewContinuousVariables(3)
# The cost is a linear function, so we call AddLinearCost
prog.AddLinearCost(-x[0])
# Now add the constraint x1-exp(x0)>=0 and x2-exp(x1)>=0
prog.AddConstraint(
    lambda z: np.array([z[1]-np.exp(z[0]), z[2]-np.exp(z[1])]),
    lb=[0, 0],
    ub=[np.inf, np.inf],
    vars=x)
# Add the bounding box constraint 0<=x0<=100, 0<=x1<=100, 0<=x2<=10
prog.AddBoundingBoxConstraint(0, 100, x[:2])
prog.AddBoundingBoxConstraint(0, 10, x[2])

# Now solve the program with initial guess x=[1, 2, 3]
result = Solve(prog, np.array([1.,2.,3.]))
print(f"Is optimization successful? {result.is_success()}")
print(f"Solution to x: {result.GetSolution(x)}")
print(f"optimal cost: {result.get_optimal_cost()}")import matplotlib.pyplot as plt
prog = MathematicalProgram()
p1 = prog.NewContinuousVariables(2, "p1")
p2 = prog.NewContinuousVariables(2, "p2")

# Add the constraint that p1 is on the unit circle centered at (0, 2)
prog.AddConstraint(
    lambda z: [z[0]**2 + (z[1]-2)**2],
    lb=np.array([1.]),
    ub=np.array([1.]),
    vars=p1)

# Add the constraint that p2 is on the curve y=x*x
prog.AddConstraint(
    lambda z: [z[1] - z[0]**2],
    lb=[0.],
    ub=[0.],
    vars=p2)

# Add the cost on the distance between p1 and p2
prog.AddQuadraticCost((p1-p2).dot(p1-p2))

# Set the value of p1 in initial guess to be [0, 1]
prog.SetInitialGuess(p1, [0., 1.])
# Set the value of p2 in initial guess to be [1, 1]
prog.SetInitialGuess(p2, [1., 1.])

# Now solve the program
result = Solve(prog)
print(f"Is optimization successful? {result.is_success()}")
p1_sol = result.GetSolution(p1)
p2_sol = result.GetSolution(p2)
print(f"solution to p1 {p1_sol}")
print(f"solution to p2 {p2_sol}")
print(f"optimal cost {result.get_optimal_cost()}")

# Plot the solution.
plt.figure()
plt.plot(np.cos(np.linspace(0, 2*np.pi, 100)), 2+np.sin(np.linspace(0, 2*np.pi, 100)))
plt.plot(np.linspace(-2, 2, 100), np.power(np.linspace(-2, 2, 100), 2))
plt.plot(p1_sol[0], p1_sol[1], '*')
plt.plot(p2_sol[0], p2_sol[1], '*')
plt.axis('equal')
plt.show()