import re
import warnings
from typing import Any, Type

import multipledispatch
import pytest
import tensorflow as tf
from _pytest.capture import CaptureFixture
from packaging.version import Version

import gpflow
from gpflow.utilities import Dispatcher


class A1:
    pass


class A2(A1):
    pass


class B1:
    pass


class B2(B1):
    pass


def get_test_fn() -> Dispatcher:
    test_fn = Dispatcher("test_fn")

    @test_fn.register(A1, B1)
    def test_a1_b1(x: A1, y: B1) -> str:
        return "a1-b1"

    @test_fn.register(A2, B1)
    def test_a2_b1(x: A2, y: B1) -> str:
        return "a2-b1"

    @test_fn.register(A1, B2)
    def test_a1_b2(x: A1, y: B2) -> str:
        return "a1-b2"

    return test_fn


def test_our_multipledispatch() -> None:
    test_fn = get_test_fn()

    assert test_fn(A1(), B1()) == "a1-b1"
    assert test_fn(A2(), B1()) == "a2-b1"
    assert test_fn(A1(), B2()) == "a1-b2"

    # test the ambiguous case:

    with warnings.catch_warnings(record=True) as w:
        warnings.simplefilter("always")

        assert test_fn(A2(), B2()) == "a1-b2"  # the last definition wins

        assert len(w) == 1
        assert issubclass(w[0].category, multipledispatch.conflict.AmbiguityWarning)

    # test that adding the child-child definition removes ambiguity warning:

    @test_fn.register(A2, B2)
    def test_a2_b2(x: A2, y: B2) -> str:
        return "a2-b2"

    with warnings.catch_warnings(record=True) as w:
        warnings.simplefilter("always")

        assert test_fn(A2(), B2()) == "a2-b2"

        assert len(w) == 0


def test_dispatcher__no_match() -> None:
    test_fn = get_test_fn()

    with pytest.raises(NotImplementedError):
        test_fn(3, "foo")

    assert None is test_fn.dispatch(int, str)

    with pytest.raises(NotImplementedError):
        test_fn.dispatch_or_raise(int, str)


@pytest.mark.parametrize(
    "Dispatcher, expect_autograph_warning",
    [
        (
            multipledispatch.Dispatcher,
            True,
        ),  # checks test is written correctly and captures warnings generated by base Dispatcher
        (gpflow.utilities.Dispatcher, False),  # no warnings with our custom Dispatcher
    ],
)
def test_dispatcher_autograph_warnings(
    capsys: CaptureFixture[str], Dispatcher: Type[Any], expect_autograph_warning: bool
) -> None:
    if Dispatcher is multipledispatch.Dispatcher and (Version(tf.__version__) >= Version("2.3.0")):
        pytest.skip("In TensorFlow >= 2.3, multipledispatch.Dispatcher no longer works at all")

    tf.autograph.set_verbosity(0, alsologtostdout=True)  # to be able to capture it using capsys

    test_fn = Dispatcher("test_fn")

    # generator would only be invoked when defining for base class...
    @test_fn.register(gpflow.inducing_variables.InducingVariables)
    def test_iv(x: gpflow.inducing_variables.InducingVariables) -> tf.Tensor:
        return tf.reduce_sum(x.Z)

    test_fn_compiled = tf.function(test_fn)  # with autograph=True by default

    # ...but calling using subclass
    result = test_fn_compiled(gpflow.inducing_variables.InducingPoints([[1.0, 2.0]]))
    assert result.numpy() == 3.0  # expect computation to work either way

    captured = capsys.readouterr()

    tf_warning = "WARNING:.*Entity .* appears to be a generator function. It will not be converted by AutoGraph."
    assert bool(re.match(tf_warning, captured.out)) == expect_autograph_warning