import numpy as np
import pytest
from quadproj.fun import Fun
from quadproj.project import project
from quadproj.quadrics import Quadric

from plot_mu_trajectories import degenerate_point, ellipse_boundary, kkt_roots, pole_line

ROTATED_ELLIPSE = Quadric(np.array([[3.0, 1.0], [1.0, 2.0]]), np.array([1.0, -1.0]), -4.0)


def trajectory(Q, x0):
    return Fun(Q, Q.to_normalized(x0)).x_not_std


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


def test_boundary_is_on_the_ellipse():
    assert np.all([abs(ROTATED_ELLIPSE.evaluate_point(x)) < 1e-10 for x in ellipse_boundary(ROTATED_ELLIPSE)])


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

    distances = boundary_distances(Q, x0)
    assert abs(Q.evaluate_point(x(mu_star_star))) < 1e-8
    assert np.linalg.norm(x(mu_star) - x0) == pytest.approx(distances.min(), abs=1e-6)
    assert np.linalg.norm(x(mu_star_star) - x0) == pytest.approx(distances.max(), abs=1e-6)


@pytest.mark.parametrize('eig, x0', [([2.0, 1.0], [0.0, -0.1]), ([2.0, 0.75], [0.0, -0.6])])
def test_degenerate_point_is_the_projection(eig, x0):
    Q, x0 = Quadric(np.diag(eig), np.zeros(2), -1.0), np.array(x0)
    xd = degenerate_point(Q, x0, 0)
    np.testing.assert_allclose(xd, project(Q, x0))
    assert np.linalg.norm(xd - x0) == pytest.approx(boundary_distances(Q, x0).min(), abs=1e-6)
    assert degenerate_point(Q, x0, 1) is None


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