from pydrake.solvers import MathematicalProgram
import numpy as np
import matplotlib.pyplot as plt

# Create an empty MathematicalProgram named prog (with no decision variables, 
# constraints or cost function)
prog = MathematicalProgram()x = prog.NewContinuousVariables(2)
print(x)
print(1 + 2*x[0] + 3*x[1] + 4*x[1])y = prog.NewContinuousVariables(2, "dog")
print(y)
print(y[0] + y[0] + y[1] * y[1] * y[1])var_matrix = prog.NewContinuousVariables(3, 2, "A")
print(var_matrix)# Add the constraint x(0) * x(1) = 1 to prog
prog.AddConstraint(x[0] * x[1] == 1)prog.AddConstraint(x[0] >= 0)
prog.AddConstraint(x[0] - x[1] <= 0)# Add a cost x(0)**2 + 3 to the total cost. Since prog doesn't have a cost before, now the total cost is x(0)**2 + 3
prog.AddCost(x[0] ** 2 + 3)prog.AddCost(x[0] + x[1])from IPython.display import display, Markdown
display(Markdown(prog.ToLatex()))"""
Solves a simple optimization problem
       min x(0)^2 + x(1)^2
subject to x(0) + x(1) = 1
           x(0) <= x(1)
"""
from pydrake.solvers import Solve
# Set up the optimization problem.
prog = MathematicalProgram()
x = prog.NewContinuousVariables(2)
prog.AddConstraint(x[0] + x[1] == 1)
prog.AddConstraint(x[0] <= x[1])
prog.AddCost(x[0] **2 + x[1] ** 2)

# Now solve the optimization problem.
result = Solve(prog)

# print out the result.
print("Success? ", result.is_success())
# Print the solution to the decision variables.
print('x* = ', result.GetSolution(x))
# Print the optimal cost.
print('optimal cost = ', result.get_optimal_cost())
# Print the name of the solver that was called.
print('solver is: ', result.get_solver_id().name())"""
An infeasible optimization problem.
"""
prog = MathematicalProgram()
x = prog.NewContinuousVariables(1)[0]
y = prog.NewContinuousVariables(1)[0]
prog.AddConstraint(x + y >= 1)
prog.AddConstraint(x + y <= 0)
prog.AddCost(x)

result = Solve(prog)
print("Success? ", result.is_success())
print(result.get_solution_result())"""
Demo on manually choosing a solver
Solves the problem
min x(0)
s.t x(0) + x(1) = 1
    0 <= x(1) <= 1
"""
from pydrake.solvers import IpoptSolver
prog = MathematicalProgram()
x = prog.NewContinuousVariables(2)
prog.AddConstraint(x[0] + x[1] == 1)
prog.AddConstraint(0 <= x[1])
prog.AddConstraint(x[1] <= 1)
prog.AddCost(x[0])

# Choose IPOPT as the solver.
# First instantiate an IPOPT solver.

solver = IpoptSolver()
# The initial guess is [1, 1]. The third argument is the options for Ipopt solver,
# and we set no solver options.
result = solver.Solve(prog, np.array([1, 1]), None)

print(result.get_solution_result())
print("x* = ", result.GetSolution(x))
print("Solver is ", result.get_solver_id().name())
print("Ipopt solver status: ", result.get_solver_details().status,
      ", meaning ", result.get_solver_details().ConvertStatusToString())
from pydrake.solvers import MakeSolver
solver = MakeSolver(IpoptSolver().solver_id())
result = solver.Solve(prog)
print(result.get_solution_result())
print("x* = ", result.GetSolution(x))print("Ipopt solver status: ", result.get_solver_details().status,
      ", meaning ", result.get_solver_details().ConvertStatusToString())from pydrake.solvers import IpoptSolver
prog = MathematicalProgram()
x = prog.NewContinuousVariables(2)
prog.AddConstraint(x[0]**2 + x[1]**2 == 100.)
prog.AddCost(x[0]**2-x[1]**2)
solver = IpoptSolver()
# The user doesn't provide an initial guess.
result = solver.Solve(prog, None, None)
print(f"Without a good initial guess, the result is {result.is_success()}")
print(f"solution {result.GetSolution(x)}")
# Pass an initial guess
result = solver.Solve(prog, [-5., 0.], None)
print(f"With a good initial guess, the result is {result.is_success()}")
print(f"solution {result.GetSolution(x)}")# Visualize the solver progress in each iteration through a callback
# Find the closest point on a curve to a desired point.

fig = plt.figure()
curve_x = np.linspace(1, 10, 100)
ax = plt.gca()
ax.plot(curve_x, 9./curve_x)
ax.plot(-curve_x, -9./curve_x)
ax.plot(0, 0, 'o')
x_init = [4., 5.]
ax.plot(x_init[0], x_init[1], 'x', color='red')

def update(x):
    ax.plot(x[0], x[1], 'x', color='red')
    
prog = MathematicalProgram()
x = prog.NewContinuousVariables(2)
prog.AddConstraint(x[0] * x[1] == 9)
prog.AddCost(x[0]**2 + x[1]**2)
prog.AddVisualizationCallback(update, x)
result = Solve(prog, x_init)