import numpy as np
import pytest

from quadric import Ellipsoid

AXIS_ELLIPSE = Ellipsoid([[2, 0], [0, 1]], [0, 0], -1)
ROTATED_ELLIPSE = Ellipsoid([[3, 1], [1, 2]], [1, -1], -4)


def boundary_distances(Q, x0):
    return np.linalg.norm(Q.boundary(200_000) - x0, axis=1)


@pytest.mark.parametrize('x0', [[0.3, 0.2], [2.0, -1.5], [-0.1, 0.4]])
def test_extreme_roots_give_nearest_and_farthest_points(x0):
    Q, x0 = ROTATED_ELLIPSE, np.array(x0)
    mu_star, mu_star_star = Q.extreme_roots(x0)
    x_star, x_star_star = Q.x(mu_star, x0), Q.x(mu_star_star, x0)

    assert abs(Q.psi(x_star)) < 1e-8
    assert abs(Q.psi(x_star_star)) < 1e-8
    distances = boundary_distances(Q, x0)
    assert np.linalg.norm(x_star - x0) == pytest.approx(distances.min(), abs=1e-6)
    assert np.linalg.norm(x_star_star - x0) == pytest.approx(distances.max(), abs=1e-6)


def test_trajectory_satisfies_kkt_conditions():
    Q, x0, mu = ROTATED_ELLIPSE, np.array([0.3, 0.2]), 0.7
    x = Q.x(mu, x0)
    np.testing.assert_allclose(x - x0 + mu * (Q.A @ x + Q.b / 2), 0, atol=1e-12)


def test_df_matches_finite_difference():
    Q, x0, mu, h = ROTATED_ELLIPSE, np.array([0.3, 0.2]), 0.7, 1e-6
    finite_difference = (Q.f(mu + h, x0) - Q.f(mu - h, x0)) / (2 * h)
    assert Q.df(mu, x0) == pytest.approx(finite_difference, rel=1e-6)


def test_degenerate_point_is_the_projection():
    Q, x0 = AXIS_ELLIPSE, np.array([0.0, -0.1])
    xd = Q.degenerate_point(x0, 0)
    assert abs(Q.psi(xd)) < 1e-12
    assert np.linalg.norm(xd - x0) == pytest.approx(boundary_distances(Q, x0).min(), abs=1e-6)
    assert Q.degenerate_point(x0, 1) is None


def test_pole_line_is_the_asymptote():
    Q, x0 = AXIS_ELLIPSE, np.array([0.3, 0.2])
    point, direction = Q.pole_line(x0, 1)
    x = Q.x(-1 / Q.eig[1] + 1e-9, x0)
    residual = (x - point) - ((x - point) @ direction) * direction
    np.testing.assert_allclose(residual, 0, atol=1e-6)


def test_sign_convention_does_not_matter():
    flipped = Ellipsoid(-ROTATED_ELLIPSE.A, -ROTATED_ELLIPSE.b, -ROTATED_ELLIPSE.c)
    x0 = np.array([0.3, 0.2])
    np.testing.assert_allclose(flipped.extreme_roots(x0), ROTATED_ELLIPSE.extreme_roots(x0))