import contourpy
import matplotlib
import matplotlib.pyplot as plt
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, kkt_roots, off_pole, plot_trajectories, pole_line

matplotlib.use('Agg')

GRID = np.linspace(-4, 4, 4001)

# (A, b, c, x0): generic points, and points with a zero normalized component (degenerate)
QUADRICS = {
    'ellipse': ([[3, 1], [1, 2]], [1, -1], -4, [0.3, 0.2]),
    'ellipse_outside': ([[3, 1], [1, 2]], [1, -1], -4, [2.0, -1.5]),
    'ellipse_degenerate': ([[2, 0], [0, 1]], [0, 0], -1, [0.0, -0.1]),
    'ellipse_degenerate_2': ([[2, 0], [0, 1]], [0, 0], -1, [0.0, 0.5]),
    'hyperbola': ([[1, 2], [2, -1]], [1, 0], -1, [0.3, 0.2]),
    'hyperbola_degenerate': ([[1, 0], [0, -1]], [0, 0], -1, [0.0, 0.5]),
    'parabola': ([[1, 1], [1, 1]], [1, -1], 0, [0.3, 0.2]),
    'parabola_degenerate': ([[1, 0], [0, 0]], [0, -1], 0, [0.0, 0.5]),
    'crossing_lines': ([[1, 0], [0, -1]], [0, 0], 0, [0.3, 0.2]),
    'parallel_lines': ([[1, 0], [0, 0]], [0, 0], -1, [0.3, 0.2]),
}
# Quadrics without finite roots of f: the projection is the limit of x(mu) as |mu| -> inf.
LIMIT_QUADRICS = {
    'single_line': ([[1, 0], [0, 0]], [0, 0], 0, [0.3, 0.2]),
    'point': ([[1, 0], [0, 1]], [0, 0], 0, [0.3, 0.2]),
}


def make_quadric(A, b, c):
    return Quadric(np.array(A, dtype=float), np.array(b, dtype=float), float(c))


ROTATED_ELLIPSE = make_quadric(*QUADRICS['ellipse'][:3])


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


def ellipse_boundary(Q, n=1000):
    """Return n points of the ellipse Q as an (n, 2) array."""
    t = np.linspace(-np.pi, np.pi, n)
    u = np.column_stack((np.cos(t), np.sin(t))) / np.sqrt(Q.eig)
    return Q.d + np.sqrt(abs(Q.gamma)) * u @ Q.V.T


def distances_to_quadric(Q, x0):
    """Distances from x0 to the points of the contour psi = 0 on a fine grid."""
    X, Y = np.meshgrid(GRID, GRID)
    P = np.stack((X, Y), axis=-1)
    psi = np.einsum('...i,ij,...j->...', P, Q.A, P) + P @ Q.b + Q.c
    return np.linalg.norm(np.vstack(contourpy.contour_generator(X, Y, psi).lines(0)) - x0, axis=1)


@pytest.mark.parametrize('name', QUADRICS)
def test_kkt_points_contain_the_projection(name):
    *quadric, x0 = QUADRICS[name]
    Q, x0 = make_quadric(*quadric), np.array(x0)
    mu_star, other_roots = kkt_roots(Q, x0)
    x = trajectory(Q, x0)
    points = [x(mu) for mu in other_roots + ([] if mu_star is None else [mu_star])]
    points += [xd for k in range(Q.dim) if (xd := degenerate_point(Q, x0, k)) is not None]
    if Q.is_conical:
        points.append(Q.d)

    assert all(abs(Q.evaluate_point(point)) < 1e-8 for point in points)
    nearest = min(np.linalg.norm(point - x0) for point in points)
    assert nearest == pytest.approx(distances_to_quadric(Q, x0).min(), abs=1e-4)


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

    distances = np.linalg.norm(ellipse_boundary(Q, 200_000) - x0, axis=1)
    assert np.linalg.norm(x(mu_star) - x0) == pytest.approx(distances.min(), abs=1e-6)
    assert max(np.linalg.norm(x(mu) - x0) for mu in other_roots) == pytest.approx(distances.max(), abs=1e-6)


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('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 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)
    direction = direction / np.linalg.norm(direction)
    residual = (x - point) - ((x - point) @ direction) * direction
    np.testing.assert_allclose(residual, 0, atol=1e-6)


@pytest.mark.parametrize('name', [*QUADRICS, *LIMIT_QUADRICS])
def test_every_quadric_type_can_be_plotted(name):
    *quadric, x0 = {**QUADRICS, **LIMIT_QUADRICS}[name]
    plt.close(plot_trajectories(make_quadric(*quadric), np.array(x0, dtype=float)))


def test_empty_quadric_is_rejected():
    with pytest.raises(ValueError):
        plot_trajectories(make_quadric([[1, 0], [0, 1]], [0, 0], 1), np.zeros(2))