import math

import numpy as np
import matplotlib.pyplot as plt

from pydrake.solvers import MathematicalProgram, Solve
import pydrake.symbolic as sym"""
Code for example 1.
"""

# Initialize an empty optimization program.
prog = MathematicalProgram()
# Declare "a" as decision variables.
a = prog.NewContinuousVariables(2)
# Declare "x" as indeterminates
x = prog.NewIndeterminates(1)

# Declare p1(x) = pow(x, 4) + a0 * x + 2 + a1. The second argument
# [x] indicates that the indeterminates in the polynomial is x.
p1 = sym.Polynomial(x[0]**4 + a[0] * x[0] + 2 + a[1], [x])
# Declare p2(x) = (a0 - a1 + 1) * pow(x, 2) + a1 * x + 1. The second argument
# [x] indicates that the indeterminates in the polynomial is x.
p2 = sym.Polynomial((a[0] - a[1] + 1) * x[0] * x[0] + a[1] * x[0] + 1, [x])
# Add the constraint that p1(x) and p2(x) are SOS
prog.AddSosConstraint(p1)
prog.AddSosConstraint(p2)

# Add the cost -a0 - a1
prog.AddCost(-a[0] - a[1])

# Solve the problem.
result = Solve(prog)

# Retrieve the solution.
print(result.get_solution_result())
print("a =", result.GetSolution(a))"""
Code for example 2.
"""

# Initialize an empty optimization program.
prog = MathematicalProgram()
# Declares indeterminates
x = prog.NewIndeterminates(1)

# Declares decision variable a, b, and c
a = prog.NewContinuousVariables(1)[0]
b = prog.NewContinuousVariables(1)[0]
c = prog.NewContinuousVariables(1)[0]

# Declare p(x), the second argument indicates this polynomial p(x) has x as the
# indeterminates.
poly_p = sym.Polynomial(x[0]**4 + a * x[0]**3 + b * x[0]** 2 + c * x[0] + 1, [x])
# Declares SOS polynomial t(x), such that s(x) = p(x) - x(1-x) * t(x) is SOS.
# The second return argument of NewSosPolynomial is the positive-semidefinite
# Gramian matrix of poly_s. We ignore this matrix here.
# t(x) should be a quadratic polynomial of x, hence the second input argument is
# 2, indicating the degree of the returned polynomial.
poly_t, _ = prog.NewSosPolynomial(sym.Variables(x), 2)
# Compute s(x) = p(x) - x(1-x)t(x) as a polynomial of indeterminate x.
poly_s = poly_p - sym.Polynomial(x[0] * (1 - x[0]), sym.Variables(x)) * poly_t
prog.AddSosConstraint(poly_s)

# Add the cost a + b + c
prog.AddCost(a + b + c)

result = Solve(prog)
print("Is success? ", result.is_success())
print(result.get_solution_result())
a_val = result.GetSolution(a)
b_val = result.GetSolution(b)
c_val = result.GetSolution(c)

# Now plot the polynomial p(x), that is always non-negative in the range [0, 1]
x_val = np.linspace(0, 1, 100)
y_val = np.power(x_val, 4) + a_val * np.power(x_val, 3) + b_val * np.power(x_val, 2) + c_val * x_val + 1
plt.xlabel("x")
plt.ylabel("p(x)")
plt.plot(x_val, y_val)