"""
Provides some general visualization tools using matplotlib.  This is not
related to `rendering`.
"""

import matplotlib as mpl
import numpy as np

from pydrake.solvers import MathematicalProgram, Solve
from pydrake.symbolic import Evaluate, Jacobian, Polynomial


def plot_sublevelset_quadratic(ax, A, b=[0, 0], c=0, vertices=51, **kwargs):
    """
    Plots the 2D ellipse representing x'Ax + b'x + c <= 1, e.g.
    the one sub-level set of a quadratic form.

    Args:
        ax:     the matplotlib axis to receive the plot
        A:      an 2x2 PSD numpy array.
        b:      broadcastable to a 2x1 array
        c:      scalar
        vertices: number of sample points along the boundary
        kwargs: are passed to the matplotlib fill method

    Returns:
        the return values from matplotlib's fill command.
    """
    # TODO(russt): Add project_quadratic_form and slice_quadratic_form
    #  methods to drake::math and recommend users call them first to plot
    #  higher-dimensional ellipsoids.
    assert isinstance(ax, mpl.axes.Axes)
    A = np.asarray(A)
    assert A.shape == (2, 2), "A must be 2x2"

    # Note: this is closely related to
    # drake::math::DecomposePositiveQuadraticForm, but does not require the
    # function to be positive for all x.

    # f = x'Ax + b'x + c
    # dfdx = x'(A+A') + b'
    # H = .5*(A+A')   # aka, the symmetric part of A (note: x'Hx = x'Ax)
    # dfdx = 0 => xmin = -inv(A+A')*b = -.5 inv(H)*b
    H = 0.5 * (A + A.T)
    xmin = np.linalg.solve(-2 * H, np.reshape(b, (2, 1)))
    fmin = -xmin.T.dot(H).dot(xmin) + c  # since b = -2*H*xmin
    assert fmin <= 1, (
        "The minimum value is > 1; there is no sub-level set to plot"
    )

    # To plot the contour at f = (x-xmin)'H(x-xmin) + fmin = 1,
    # we make a circle of values y, such that: y'y = 1-fmin,
    th = np.linspace(0, 2 * np.pi, vertices)
    Y = np.sqrt(1 - fmin) * np.vstack([np.sin(th), np.cos(th)])
    # then choose L'*(x - xmin) = y, where H = LL'.
    L = np.linalg.cholesky(H)
    X = np.tile(xmin, vertices) + np.linalg.inv(np.transpose(L)).dot(Y)

    return ax.fill(X[0, :], X[1, :], **kwargs)


def plot_sublevelset_expression(ax, e, vertices=51, **kwargs):
    """
    Plots the 2D sub-level set e(x) <= 1, which must contain the origin.

    Args:
        ax:       the matplotlib axis to receive the plot
        e:        a symbolic expression in two variables
        vertices: number of sample points along the boundary
        kwargs:   are passed to the matplotlib fill method

    Returns:
        the return values from matplotlib's fill command.
    """

    x = list(e.GetVariables())
    assert len(x) == 2, "e must be an expression in two variables"

    # Handle the special case where e is a degree 2 polynomial.
    if e.is_polynomial():
        p = Polynomial(e)
        if p.TotalDegree() == 2:
            env = {a: 0 for a in x}
            c = e.Evaluate(env)
            e1 = e.Jacobian(x)
            b = Evaluate(e1, env)
            e2 = Jacobian(e1, x)
            A = 0.5 * Evaluate(e2, env)
            return plot_sublevelset_quadratic(ax, A, b, c, vertices, **kwargs)

    # Find the level-set in polar coordinates, by sampling theta and
    # root-finding (on the scalar expression) to find a rplus and rminus.

    Xplus = np.empty((2, vertices))
    Xminus = np.empty((2, vertices))
    i = 0
    for theta in np.linspace(0, np.pi, vertices):
        prog = MathematicalProgram()
        r = prog.NewContinuousVariables(1, "r")[0]
        env = {x[0]: r * np.cos(theta), x[1]: r * np.sin(theta)}
        scalar = e.Substitute(env)
        b = prog.AddBoundingBoxConstraint(0, np.inf, r)
        prog.AddConstraint(scalar == 1)
        prog.AddQuadraticCost([1], [0], [r])
        prog.SetInitialGuess(r, 0.1)  # or anything non-zero.
        result = Solve(prog)
        assert result.is_success(), "Failed to find the level set"
        rplus = result.GetSolution(r)
        Xplus[0, i] = rplus * np.cos(theta)
        Xplus[1, i] = rplus * np.sin(theta)
        b.evaluator().UpdateLowerBound([-np.inf])
        b.evaluator().UpdateUpperBound([0])
        prog.SetInitialGuess(r, -0.1)  # or anything non-zero.
        result = Solve(prog)
        assert result.is_success(), "Failed to find the level set"
        rminus = result.GetSolution(r)
        Xminus[0, i] = rminus * np.cos(theta)
        Xminus[1, i] = rminus * np.sin(theta)
        i = i + 1

    return ax.fill(
        np.hstack((Xplus[0, :], Xminus[0, :])),
        np.hstack((Xplus[1, :], Xminus[1, :])),
        **kwargs,
    )
