"""Ellipsoid x'Ax + b'x + c = 0 and the KKT trajectory of the projection onto it.

The KKT conditions of min ||x - x0|| subject to x on the ellipsoid give

    x(mu) = (I + mu A)^{-1} (x0 - mu b / 2),

and the projections are the points x(mu) with f(mu) = psi(x(mu)) = 0.

With d = -A^{-1} b / 2, gamma = c + b'd + d'Ad < 0 and A = V diag(eig) V', the standardized
coordinates u = V'(x - d) / sqrt(-gamma) turn the ellipsoid into sum_i eig_i u_i^2 = 1 and the
trajectory into u_i(mu) = u0_i / (1 + mu eig_i), which has a pole at mu = -1/eig_i when u0_i != 0.
"""

import numpy as np
from scipy.optimize import brentq

TOL = 1e-6


class Ellipsoid:
    def __init__(self, A, b, c):
        A = np.asarray(A, dtype=float)
        b = np.asarray(b, dtype=float)
        d = np.linalg.solve(A, -b / 2)
        gamma = c + b @ d + d @ A @ d
        if gamma > 0:  # same set, written with gamma < 0
            A, b, c, gamma = -A, -b, -c, -gamma

        eig, V = np.linalg.eigh(A)
        order = np.argsort(-eig)
        self.eig = eig[order]
        self.V = V[:, order]
        if not np.all(self.eig > 0):
            raise ValueError('the quadric is not a non-empty ellipsoid')

        self.A, self.b, self.c, self.d = A, b, c, d
        self.dim = len(b)
        self.scale = np.sqrt(-gamma)

    def psi(self, x):
        return x @ self.A @ x + self.b @ x + self.c

    def to_standardized(self, x):
        return self.V.T @ (x - self.d) / self.scale

    def to_non_standardized(self, u):
        return self.d + self.scale * (self.V @ u)

    def _solve(self, mu, y):
        """Return (I + mu A)^{-1} y."""
        return self.V @ ((self.V.T @ y) / (1 + mu * self.eig))

    def x(self, mu, x0):
        return self._solve(mu, x0 - mu * self.b / 2)

    def f(self, mu, x0):
        return self.psi(self.x(mu, x0))

    def df(self, mu, x0):
        x = self.x(mu, x0)
        dx = -self._solve(mu, self.A @ x + self.b / 2)
        return (2 * self.A @ x + self.b) @ dx

    def extreme_roots(self, x0):
        """Return (mu*, mu**), the roots of f right of the rightmost and left of the leftmost active pole.

        x(mu*) and x(mu**) are the nearest and farthest non-degenerate KKT points; a degenerate
        point (see degenerate_point) may be nearer than x(mu*). Return None if x0 is the centre.
        """
        u0 = self.to_standardized(x0)
        active = np.abs(u0) > TOL
        if not active.any():
            return None
        eig, u0 = self.eig[active], u0[active]

        def secular(mu):
            return np.sum(eig * (u0 / (1 + mu * eig))**2) - 1

        mu_star = brentq(secular, *_bracket(secular, -1 / eig.max(), +1))
        mu_star_star = brentq(secular, *_bracket(secular, -1 / eig.min(), -1))
        return mu_star, mu_star_star

    def _pole_point(self, u0, k):
        """Standardized limit of u(mu) as mu -> -1/eig_k, with the eig_k components set to 0."""
        on_pole = np.abs(self.eig - self.eig[k]) < TOL
        u = np.zeros(self.dim)
        np.divide(u0, 1 - self.eig / self.eig[k], out=u, where=~on_pole)
        return u, on_pole

    def pole_line(self, x0, k):
        """Return (point, direction) of the line {x : u_i = u0_i / (1 - eig_i/eig_k) for eig_i != eig_k}.

        It is the asymptote of x(mu) at mu = -1/eig_k when u0_k != 0, and the line containing
        the degenerate points otherwise.
        """
        u, _ = self._pole_point(self.to_standardized(x0), k)
        return self.to_non_standardized(u), self.V[:, k]

    def degenerate_point(self, x0, k):
        """Return the KKT point at mu = -1/eig_k when u0 vanishes on that eigenspace, or None."""
        u0 = self.to_standardized(x0)
        u, on_pole = self._pole_point(u0, k)
        if np.any(np.abs(u0[on_pole]) > TOL):
            return None
        r = (1 - self.eig @ u**2) / self.eig[k]
        if r <= 0:
            return None
        u[np.argmax(on_pole)] = np.sqrt(r)
        return self.to_non_standardized(u)

    def boundary(self, n=1000):
        """Return n points of a 2D ellipse as an (n, 2) array."""
        assert self.dim == 2, 'only 2D ellipses can be sampled'
        t = np.linspace(-np.pi, np.pi, n)
        u = np.column_stack((np.cos(t), np.sin(t))) / np.sqrt(self.eig)
        return self.d + self.scale * u @ self.V.T


def _bracket(f, pole, direction):
    """Return an interval on the `direction` side of `pole` where f changes sign.

    Assumes f -> +inf next to the pole and f -> -1 far from it.
    """
    near = far = 1.0
    while f(pole + direction * near) <= 0:
        near /= 2
    while f(pole + direction * far) >= 0:
        far *= 2
    return sorted((pole + direction * near, pole + direction * far))