"""
Example adapted from SOSTOOLS sosdemo2.

See page 34 of PDF below for details.
https://github.com/oxfordcontrol/SOSTOOLS/blob/SOSTOOLS400/docs/SOSTOOLS_400.pdf
"""
import polymatrix as poly
import sumofsquares as sos

# TODO: Expose Solver class in sumofsquares.__init__
from sumofsquares.abc import Solver

# state variables and dynamics RHS
x1, x2, x3 = poly.from_names("x_1, x_2, x_3")
x = poly.give_name(poly.v_stack((x1, x2, x3)), "x")

f = poly.give_name(poly.v_stack((
    (-x1**3 - x1*x3**2) * (x3**2 + 1),
    (-x2 - x1**2 * x2) * (x3**2 + 1),
    (-x3 + 3*x1**2*x3) * (x3**2 + 1) - 3*x3
)), "f(x) * ((x_3 ** 2) + 1)")

# Create Lyapunov function parametrized by w
monomials = x ** 2
w = sos.from_name("w", shape=monomials.shape)

V = poly.give_name(w.T @ monomials, "V(x)")
dV = V.diff(x)

# Solve SOS problem
constraints = (
    sos.make_sos_constraint(V - (x.T @ x)), # last term to impose definiteness
    sos.make_sos_constraint(- (dV @ f)),
)

prob, state, result = sos.solve_problem(cost=None, constraints=constraints, solver=Solver.CVXOPT)
print(prob)

# TODO: create SolverStatus enumeration
assert result.solver_info.status == "optimal", f"Status {result.solver_info.status}"

V_sol = V.eval(result.values)
print(poly.to_sympy(V_sol).read(poly.make_state()))