import inspect
import json
from pathlib import Path

import numpy as np
import pytest

from logic import DIVERGED, OPTIMISERS, SINGULAR, Objective, Run, parse_function
from manager import MAX_ITERATIONS, Manager, UserError

SPEC = json.loads((Path(__file__).parent.parent / "src" / "algorithms.json").read_text())["algorithms"]
DEFAULTS = {spec["id"]: {p["name"]: p["default"] for p in spec["params"]} for spec in SPEC}


def run(expression: str, algorithm: str, start: list[float], **params) -> Run:
    f = Objective(parse_function(expression, len(start)), len(start))
    return Run(f, OPTIMISERS[algorithm](f, **(DEFAULTS[algorithm] | params)), np.array(start, dtype=float))


def test_spec_matches_optimisers():
    assert [spec["id"] for spec in SPEC] == list(OPTIMISERS)
    for spec in SPEC:
        names = list(inspect.signature(OPTIMISERS[spec["id"]]).parameters)[1:]
        assert names == [p["name"] for p in spec["params"]], spec["id"]


@pytest.mark.parametrize("expression, n, expected", [
    ("x^2 - 0.1cos(20x)", 1, 0.81 - 0.1 * np.cos(18)),
    ("2x y + e", 2, 2 * 0.9 * -0.5 + np.e),
    ("sqrt(abs(x)) + ln(exp(y))", 2, np.sqrt(0.9) - 0.5),
])
def test_parse_function(expression, n, expected):
    f = Objective(parse_function(expression, n), n)
    point = np.array([0.9, -0.5][:n])
    assert f.value(point) == pytest.approx(expected)


@pytest.mark.parametrize("expression, message", [
    ("", "Enter a function"),
    ("x^2 + (", "can read"),
    ("x + z", "Unknown name z"),
])
def test_parse_function_errors(expression, message):
    with pytest.raises(ValueError, match=message):
        parse_function(expression, 2)


def test_gradient_and_hessian():
    f = Objective(parse_function("x^2 y + y^3", 2), 2)
    p = np.array([1.0, 2.0])
    assert f.gradient(p).tolist() == [4.0, 13.0]
    assert f.hessian(p).tolist() == [[4.0, 2.0], [2.0, 12.0]]


def test_gradient_descent_on_a_quadratic():
    # x_k = (1 - 2η)^k x_0
    r = run("x^2", "gradient_descent", [1.0], learning_rate=0.1)
    assert r.point(10)["position"][0] == pytest.approx(0.8**10)


def test_momentum_overshoots_and_nesterov_damps_it():
    plain = run("x^2", "gradient_descent", [1.0], learning_rate=0.1, momentum=0.9)
    nesterov = run("x^2", "nesterov", [1.0], learning_rate=0.1, momentum=0.9)
    assert min(plain.point(k)["position"][0] for k in range(20)) < -0.5
    assert abs(nesterov.point(20)["position"][0]) < abs(plain.point(20)["position"][0])


def test_newton_solves_a_quadratic_in_one_step():
    r = run("(x - 1)^2 + 3(y + 2)^2 + x y", "newton", [5.0, 5.0])
    assert r.point(1)["gradient"] == pytest.approx([0, 0], abs=1e-12)


@pytest.mark.parametrize("algorithm", [spec["id"] for spec in SPEC])
def test_every_optimiser_goes_downhill_on_a_bowl(algorithm):
    r = run("x^2 + 3y^2", algorithm, [0.8, -0.6])
    assert r.point(50)["value"] < r.point(0)["value"]


def test_divergence_stops_the_run():
    r = run("x^2", "gradient_descent", [1.0], learning_rate=1.5)
    # |x| doubles each step, so it passes 1e6 after 20 steps
    assert r.point(20) is None and r.point(19) is not None
    assert r.stopped == DIVERGED


def test_leaving_the_domain_stops_the_run():
    r = run("sqrt(x)", "gradient_descent", [0.5], learning_rate=1.0)
    assert r.point(3) is None
    assert r.stopped == DIVERGED


def test_singular_hessian_stops_newton():
    r = run("x^3", "newton", [0.0])
    assert r.point(1) is None
    assert r.stopped == SINGULAR


def test_singular_hessian_stops_newton_in_pyodide(monkeypatch):
    # Pyodide's solve returns inf rather than raising LinAlgError
    monkeypatch.setattr(np.linalg, "solve", lambda a, b: np.full_like(b, np.inf))
    r = run("x^2 + 0y", "newton", [1.0, 1.0])
    assert r.point(1) is None
    assert r.stopped == SINGULAR


def test_points_are_cached_and_asked_for_in_any_order():
    r = run("x^2 + y^2", "adam", [1.0, 1.0])
    late = r.point(30)
    assert r.point(3) == r.points[3]
    assert r.point(30) == late


def manager_with(expression: str, *runs: dict) -> Manager:
    m = Manager()
    m.handle_set_function({"variables": len(runs[0]["start"]), "expression": expression})
    m.handle_set_runs({"runs": list(runs)})
    return m


def test_state_past_the_last_iteration():
    m = manager_with("x^2", {"algorithm": "gradient_descent", "start": [1.0], "params": {}})
    assert m.handle_state({"iteration": MAX_ITERATIONS}) is not None
    assert m.handle_state({"iteration": MAX_ITERATIONS + 1}) is None


def test_state_ends_when_every_run_has_stopped():
    m = manager_with(
        "x^3 + x^2",
        {"algorithm": "gradient_descent", "start": [1.0], "params": {"learning_rate": 1.5}},
        {"algorithm": "newton", "start": [-1/3], "params": {}},  # f'' = 0 here
    )
    states = [m.handle_state({"iteration": k}) for k in range(MAX_ITERATIONS)]
    end = states.index(None)
    assert 1 < end and all(s is None for s in states[end:])
    # Newton stops first and says why while gradient descent carries on
    assert states[1]["runs"][1] == {"point": None, "stopped": SINGULAR}
    assert states[1]["runs"][0]["point"] is not None
    # Each run says why at its last point
    assert states[0]["runs"][1]["stopped"] == SINGULAR
    assert states[end - 2]["runs"][0]["stopped"] is None
    assert states[end - 1]["runs"][0]["stopped"] == DIVERGED


def test_start_point_outside_the_domain():
    m = Manager()
    m.handle_set_function({"variables": 1, "expression": "log(x)"})
    with pytest.raises(UserError, match="no finite value at the initial point"):
        m.handle_set_runs({"runs": [{"algorithm": "adam", "start": [-1.0], "params": {}}]})


def test_failed_runs_replace_the_old_ones():
    m = manager_with("log(x)", {"algorithm": "adam", "start": [1.0], "params": {}})
    with pytest.raises(UserError):
        m.handle_set_runs({"runs": [{"algorithm": "adam", "start": [-1.0], "params": {}}]})
    assert m.handle_state({"iteration": 0}) is None


def test_function_values_one_variable():
    m = Manager()
    m.handle_set_function({"variables": 1, "expression": "log(x)"})
    values = m.handle_function_values({"x_range": [-1, 1], "y_range": [0, 1]})
    assert len(values["x"]) == len(values["f"])
    assert values["f"][0] is None and values["f"][-1] == 0.0


def test_function_values_two_variables():
    m = Manager()
    m.handle_set_function({"variables": 2, "expression": "x + 2y"})
    values = m.handle_function_values({"x_range": [0, 1], "y_range": [0, 2]})
    assert values["f"][-1][0] == 4.0  # rows follow y
    assert set(values["arrows"]["u"]) == {1.0} and set(values["arrows"]["v"]) == {2.0}


def test_no_function_values_without_a_function():
    assert Manager().handle_function_values({"x_range": [0, 1], "y_range": [0, 1]}) is None