#!/usr/bin/env python
# coding: utf-8

# # Quadratic Program (QP) Tutorial
# 
# For instructions on how to run these tutorial notebooks, please see the [README](https://github.com/RobotLocomotion/drake/blob/master/tutorials/README.md).

# ## Important Note
# Please refer to [mathematical program tutorial](./mathematical_program.ipynb) for constructing and solving a general optimization program in Drake.

# A (convex) quadratic program (QP) is a special type of convex optimization. Its cost function is a convex quadratic function. Its constraints are linear, same as the constraints in linear program. A (convex) quadratic program has the following form
# \begin{align}
# \min_x 0.5 x^TQx + b^Tx + c\\
# \text{s.t } Ex \leq f
# \end{align}
# where `Q` is a positive semidefinite matrix.
# 
# A quadratic program can be solved by many different solvers. Drake supports some solvers including OSQP, SCS, Gurobi, Mosek, etc. Please see our [Doxygen page]( https://drake.mit.edu/doxygen_cxx/group__solvers.html) for a complete list of supported solvers. Note that some commercial solvers (such as Gurobi and Mosek) are not included in the pre-compiled Drake binaries, and therefore not on Binder/Colab. 
#     
# Drake's API supports multiple functions to add quadratic cost and linear constraints. We briefly go through some of the functions in this tutorial. For a complete list of functions, please check our [Doxygen](https://drake.mit.edu/doxygen_cxx/classdrake_1_1solvers_1_1_mathematical_program.html).
# 
# There are many applications of quadratic programs in robotics, for example, we can solve differential inverse kinematics problem as a QP, see [DifferentialInverseKinematics](https://drake.mit.edu/doxygen_cxx/namespacedrake_1_1manipulation_1_1planner.html#ab53fd2e1578db60ceb43b754671ae539) for more details. For more examples, check out [Underactuated Robotics code repo](https://github.com/RussTedrake/underactuated)
# 
# ## Add quadratic cost
# 
# ### AddQuadraticCost function
# 
# The easist way to add a quadratic cost is to call the `AddQuadraticCost` function. In the following code snippet, we first construct a program with 2 decision variables, and then show how to call the `AddQuadraticCost` function.

# In[ ]:


from pydrake.solvers.mathematicalprogram 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")


# We can call `AddQuadraticCost(expression)` to add the quadratic cost, where `expression` is a symbolic expression representing a quadratic cost.

# In[ ]:


# 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])


# If we call `AddQuadraticCost` again, the total cost of `prog` is the summation of all the added costs. You can see that `prog.quadratic_costs()` has two entries. And the total cost of `prog` is `cost1 + cost2`.

# In[ ]:


# 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())}")


# If you know the coefficients of the quadratic cost `Q, b, c`, you could also add the cost without using the symbolic expression, as shown in the following code snippet

# In[ ]:


# 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}")


# ### AddQuadraticErrorCost
# You could also add the quadratic cost
# \begin{align}
# (x - x_{desired})^TQ(x-x_{desired})
# \end{align}
# to the program by calling `AddQuadraticErrorCost`. Here is the code example

# In[ ]:


# 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}")


# ### AddL2NormCost
# You could also add the quadratic cost
# \begin{align}
# |Ax-b|^2
# \end{align}
# which is the squared L2 norm of the vector `Ax-b` to the program by calling `AddL2NormCost`. Here is the code example

# In[ ]:


# Adds the squared norm of (x[0]+2*x[1]-2, x[1] - 3) to the program cost.
cost5 = prog.AddL2NormCost(A=[[1, 2], [0, 1]], b=[2, 3], vars=x)
print(f"cost5 is {cost5}")


# ## Add linear cost
# 
# You could also add linear costs to a quadratic program. For an introduction on the different APIs on adding a linear cost, please refer to our [linear programming tutorial](./linear_program.ipynb). Here is an example

# In[ ]:


# 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())}")


# ## Add linear constraints
# 
# To add linear constraints into quadratic program, please refer to the section `Add linear constraints` in our [linear programming tutorial](./linear_program.ipynb). Here is a brief example

# In[ ]:


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)


# ## A complete code example
# Here we show a complete example to construct and solve a QP

# In[ ]:


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.AddL2NormCost(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()}")


# For more details on quadratic programming, you could refer to
# 
# [Quadratic Programming wiki](https://en.wikipedia.org/wiki/Quadratic_programming)
# 
# [Numerical Optimization by J. Nocedal and S.Wright](http://www.apmath.spbu.ru/cnsa/pdf/monograf/Numerical_Optimization2006.pdf)