from IPython.display import Markdown, display
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
from matplotlib import cm
from pydrake.solvers import MathematicalProgram
from pydrake.solvers import IpoptSolver
import numpy as np


def constraint(x):
    return [np.cos(x[0]) + 2 * np.cos(x[0] - x[1])]

# Find a solution satisfying
# 1 <= cos(x[0]) + 2 * cos(x[0] - x[1]) <= 2
# This problem has infinitely many solutions, for example, x = [0, pi/2]

# To visualize the constraint, I draw the landscape of cos(x[0]) + 2 * cos(x[0] - x[1])), and
# also highlight the point x = [0, 0]. You could see that x = [0, 0] is at a
# peak of the landscape.
def draw_constraint_landscape():
    fig = plt.figure()
    ax = fig.add_subplot(111, projection='3d')
    x_mesh, y_mesh = np.meshgrid(np.linspace(-np.pi, np.pi, 31), np.linspace(-np.pi, np.pi, 31))
    constraint_val = np.cos(x_mesh) + 2 * np.cos(x_mesh - y_mesh)
    surf = ax.plot_surface(x_mesh, y_mesh, constraint_val, cmap=cm.coolwarm, alpha=0.8)
    ax.plot([0], [0], [3], marker='.', color='g', markersize=20)
    ax.set_xlabel("x[0]")
    ax.set_ylabel("x[1]")
    ax.set_zlabel("cos(x[0]) + 2 * cos(x[0]-x[1])")
    fig.show()
    

draw_constraint_landscape()

prog = MathematicalProgram()
x = prog.NewContinuousVariables(2, "x")

prog.AddConstraint(constraint, [1], [2], x)

solver = IpoptSolver()
# With the initial guess being (0, 0), the solver cannot find a solution.
prog.SetInitialGuess(x, [0, 0])
result = solver.Solve(prog)
print(f"Starting from x=[0, 0], the solver result is {result.get_solution_result()}")
print(f"The solver gets stuck at x={result.GetSolution(x)}")

# With a different initial guess, the solver can find the solution
prog.SetInitialGuess(x, [0.1, 0.5])
result = solver.Solve(prog)
print(f"Starting from x=[0.1, 0.5], the solver result is {result.get_solution_result()}")
print(f"The found solution is x={result.GetSolution(x)}")# A sample (quadratic) program
prog = MathematicalProgram()
x = prog.NewContinuousVariables(3, "x")
prog.AddQuadraticCost(x[0] * x[0] + 2 * x[0] + 3)
prog.Add2NormSquaredCost(A = [[1, 3], [2, 4]], b=[1, 4], vars=[x[1], x[2]])
prog.AddLinearEqualityConstraint(x[0] + 2*x[1] == 5)
prog.AddLinearConstraint(x[0] + 4 *x[1] <= 10)
prog.AddBoundingBoxConstraint(-1, 10, x)

# Now print a summary:
print(prog)display(Markdown(prog.ToLatex()))prog = MathematicalProgram()
x = prog.NewContinuousVariables(2, "point")
print(x)prog = MathematicalProgram()
x = prog.NewContinuousVariables(2, "point")
constraint = prog.AddConstraint(lambda z: [np.sum(z**2)], [1.], [1.], x)
constraint.evaluator().set_description("unit-length constraint")
print(constraint)prog = MathematicalProgram()
x = prog.NewContinuousVariables(2, "point")
# Add the cost on the distance to (1, 2)
cost1 = prog.AddCost(lambda z: np.sqrt((z[0]-1)**2 + (z[1]-2)**2), x)
cost1.evaluator().set_description("distance to (1, 2)")
# Add the cost on the distance to (3, -1)
cost2 = prog.AddCost(lambda z: np.sqrt((z[0]-3)**2 + (z[1] + 1)**2), x)
cost2.evaluator().set_description("distance to (3, -1)")
print(f"cost1: {cost1}")
print(f"cost2: {cost2}")import numpy as np
from pydrake.autodiffutils import (
    InitializeAutoDiff,
    ExtractGradient,
)
prog = MathematicalProgram()
x = prog.NewContinuousVariables(2, "x")

# Add the constraint dist(x, 0) >= 1
constraint1 = prog.AddConstraint(lambda z: [z.dot(z)], [1], [np.inf], x)
constraint1.evaluator().set_description("outside unit circle")

# Add the constraint x[0]**2 + 4 * x[1]**2 <= 4
constraint2 = prog.AddConstraint(lambda z: [z[0]**2 + 4 * z[1]**2], [0], [4], x)
constraint2.evaluator().set_description("inside ellipsoid 1")

solver = IpoptSolver()
prog.SetInitialGuess(x, [0, 0])
result = solver.Solve(prog)
print("Start from initial guess x = [0, 0]")
print(f"optimization status: {result.get_solution_result()}")
infeasible_constraints = result.GetInfeasibleConstraints(prog)
for c in infeasible_constraints:
    print(f"infeasible constraint: {c}")
x_stuck = result.GetSolution(x)
print(f"x_stuck={x_stuck.T}")
# Now evaluate the gradient of the constraint at x_stuck (where the solver gets stuck)
print(f"Gradient of the infeasible constraint at x_stuck: {ExtractGradient(infeasible_constraints[0].evaluator().Eval(InitializeAutoDiff(x_stuck)))}")

# For a different initial state, the constraint that was infeasible now has non-zero gradient
x_new = np.array([0.1, 0.2])
print(f"\nStart from initial guess x_new = {x_new}")
print(f"Gradient of the infeasible constraint at x_new: {ExtractGradient(infeasible_constraints[0].evaluator().Eval(InitializeAutoDiff(x_new)))}")
prog.SetInitialGuess(x, x_new)
# With this new initial guess, the solver will be able to find the solution
result = solver.Solve(prog)
print(f"optimization status: {result.get_solution_result()}")from pydrake.solvers import CommonSolverOption, SolverOptions

# Create a simple program consisting of
#  constraint: 1 <= squared_norm(x) <= 2.
#  constraint: 2 <= x[0]**2 + 4 * x[0]*x[1] + 4*x[1]**2 + 2 * x[0] <= 5.
prog = MathematicalProgram()
x = prog.NewContinuousVariables(2)
prog.AddConstraint(lambda z: [z.dot(z)], [1], [2], x)
prog.AddConstraint(lambda z: [z[0] ** 2 + 4 * z[0] * z[1] + 4 * z[1]**2 + 2 * z[0]], [2], [5], x)
prog.SetInitialGuess(x, [0, 0])

# Solve with printing.
ipopt_solver = IpoptSolver()
filename = "/tmp/debug.txt"
solver_options = SolverOptions()
solver_options.SetOption(CommonSolverOption.kPrintFileName, filename)
result = ipopt_solver.Solve(prog, solver_options=solver_options)
with open(filename) as f:
    print(f.read())
# Enable a slightly more verbose level of printing this time.
# The options can be attached to the program, instead of passed to Solve().
solver_options.SetOption(IpoptSolver.id(), "print_level", 5)
prog.SetSolverOptions(solver_options)
result = ipopt_solver.Solve(prog)# Visualize the solver progress in each iteration through a callback
# Find the closest point on a curve to a desired point.
from pydrake.solvers import Solve

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')
ax.axis('equal')

def visualization_callback(x):
    ax.plot(x[0], x[1], 'x', color='gray')
    
prog = MathematicalProgram()
x = prog.NewContinuousVariables(2)
prog.AddConstraint(x[0] * x[1] == 9)
prog.AddCost(x[0]**2 + x[1]**2)
prog.AddVisualizationCallback(visualization_callback, x)
result = Solve(prog, x_init)## Demonstrate EvalBinding function
prog = MathematicalProgram()
p1 = prog.NewContinuousVariables(2, "p1")
p2 = prog.NewContinuousVariables(2, "p2")

# Add the constraint that p1 is in an ellipsoid (p1(0)-1)**2 + 4*(p1(1)-2)**2 <= 1
constraint1 = prog.AddConstraint(lambda z: [(z[0]-1)**2 + 4 * (z[1]-2)**2], [0], [1], p1)
# Add the constraint that p2 is in an ellipsoid (p2(0) + 2)**2 + 0.25*(p2(1)+ 1)**2) <= 1
constraint2 = prog.AddConstraint(lambda z: [(z[0]+2)**2 + 0.25*(z[1]+1)**2], [0], [1], p2)
# Add a cost to minimize the distance between p1 and p2
cost = prog.AddCost((p1-p2).dot(p1-p2))

# Evaluate the constraint and cost at a guess p1=[0, 1], p2 = [-1, -4]
p1_val = [0, 1]
p2_val = [-1, -4]
prog.SetInitialGuess(p1, p1_val)
prog.SetInitialGuess(p2, p2_val)
print(f"constraint 1 evaluated at p1={p1_val}, p2={p2_val} is {prog.EvalBinding(constraint1, prog.initial_guess())}")
print(f"constraint 2 evaluated at p1={p1_val}, p2={p2_val} is {prog.EvalBinding(constraint2, prog.initial_guess())}")
print(f"cost evaluated at p1={p1_val}, p2={p2_val} is {prog.EvalBinding(cost, prog.initial_guess())}")
# Relaxing the constraint
prog = MathematicalProgram()
x = prog.NewContinuousVariables(2)

# Add the constraint x^T * x <= 1
constraint1 = prog.AddConstraint(lambda z: [z.dot(z)], [0], [1], x)
constraint1.evaluator().set_description("inside unit circle")

# Add the constraint norm(x-[3, 0]) <= 1
constraint2 = prog.AddConstraint(lambda z: [np.sum((z - np.array([3, 0]))**2)], [0], [1], x)
constraint2.evaluator().set_description("distance to [3, 0] less than 1")

prog.SetInitialGuess(x, [1, 0])
solver = IpoptSolver()
result = solver.Solve(prog)
print(f"For the original problem, the solver status is {result.get_solution_result()}")
print(f"x is stuck at {result.GetSolution(x)}")
# Now get the infeasible constraint
infeasible_constraints = result.GetInfeasibleConstraints(prog)
for c in infeasible_constraints:
    print(f"infeasible constraint: {c}")

# Now update the upper bound of the first infeasible constraint
infeasible_constraints[0].evaluator().UpdateUpperBound([4])
# I also update the description of the constraint. Without updating the description, the
# problem still solves fine, but it would be confusing if you print out this constraint.
infeasible_constraints[0].evaluator().set_description("inside a circle with radius=2")
result = solver.Solve(prog)
print(f"For the relaxed problem, the solver status is {result.get_solution_result()}")
print(f"Solution is x = {result.GetSolution(x)}")