import abc

import pytest
import sympy

from comb_spec_searcher.utils import taylor_expand
from tilings import GriddedPerm, Tiling
from tilings.algorithms import LocalEnumeration, MonotoneTreeEnumeration
from tilings.exception import InvalidOperationError


class CommonTest(abc.ABC):
    @abc.abstractmethod
    @pytest.fixture
    def enum_verified(self):
        raise NotImplementedError

    @abc.abstractmethod
    @pytest.fixture
    def enum_not_verified(self):
        raise NotImplementedError

    def test_verified(self, enum_verified, enum_not_verified):
        assert enum_verified.verified()
        assert not enum_not_verified.verified()

    @abc.abstractmethod
    def test_get_genf(self, enum_verified):
        raise NotImplementedError

    def test_get_genf_not_verified(self, enum_not_verified):
        with pytest.raises(InvalidOperationError):
            enum_not_verified.get_genf()


class TestLocalEnumeration(CommonTest):
    @pytest.fixture
    def enum_verified(self):
        t = Tiling(
            obstructions=[
                GriddedPerm((0, 1, 2), ((0, 0),) * 3),
                GriddedPerm((0, 2, 1), ((1, 0),) * 3),
                GriddedPerm((0, 1, 2), ((1, 0),) * 3),
                GriddedPerm((0, 1), ((1, 1),) * 2),
            ],
            requirements=[
                [
                    GriddedPerm((0, 1), ((0, 0),) * 2),
                    GriddedPerm((1, 0), ((0, 0),) * 2),
                ],
                [GriddedPerm((0,), ((1, 0),))],
            ],
        )
        return LocalEnumeration(t)

    @pytest.fixture
    def enum_not_verified(self):
        t = Tiling(
            obstructions=[
                GriddedPerm((0, 1, 2), ((0, 0),) * 3),
                GriddedPerm((0, 2, 1), ((1, 0),) * 3),
                GriddedPerm((0, 1, 2), ((1, 0),) * 3),
                GriddedPerm((0, 1), ((1, 1),) * 2),
                GriddedPerm((0, 1), ((0, 0), (1, 1))),
            ],
            requirements=[
                [
                    GriddedPerm((0, 1), ((0, 0),) * 2),
                    GriddedPerm((0, 1), ((1, 0),) * 2),
                ]
            ],
        )
        return LocalEnumeration(t)

    @pytest.fixture
    def onebyone_enum(self):
        return LocalEnumeration(Tiling.from_string("123"))

    @pytest.fixture
    def enum_no_req(self):
        t = Tiling(
            obstructions=[
                GriddedPerm((0, 1, 2), ((0, 0),) * 3),
                GriddedPerm((0, 2, 1), ((1, 0),) * 3),
                GriddedPerm((0, 1, 2), ((1, 0),) * 3),
                GriddedPerm((0, 1), ((1, 1),) * 2),
            ],
            requirements=[
                [
                    GriddedPerm((0, 1), ((0, 0),) * 2),
                    GriddedPerm((0, 1), ((1, 0),) * 2),
                ]
            ],
        )
        return LocalEnumeration(t, no_req=True)

    def test_req_is_single_cell(self):
        assert LocalEnumeration._req_is_single_cell([GriddedPerm((0,), ((0, 1),))])
        assert LocalEnumeration._req_is_single_cell(
            [GriddedPerm((0, 1), ((0, 1), (0, 1)))]
        )
        assert not LocalEnumeration._req_is_single_cell(
            [GriddedPerm((0, 1), ((0, 0), (0, 1)))]
        )
        assert not LocalEnumeration._req_is_single_cell(
            [GriddedPerm((0,), ((0, 1),)), GriddedPerm((0,), ((1, 0),))]
        )
        assert LocalEnumeration._req_is_single_cell(
            [
                GriddedPerm((0, 1), ((0, 0),) * 2),
                GriddedPerm((1, 0), ((0, 0),) * 2),
            ]
        )
        assert not LocalEnumeration._req_is_single_cell(
            [
                GriddedPerm((0, 1), ((1, 0),) * 2),
                GriddedPerm((1, 0), ((0, 0),) * 2),
            ]
        )

    def test_verified(self, enum_verified, enum_not_verified):
        assert enum_verified.verified()
        assert not enum_not_verified.verified()

    def test_crossing_req_list(self):
        """
        This tiling is not local verified because of the requirement list in multiple
        cells.
        """
        t = Tiling(
            obstructions=[
                GriddedPerm((0, 2, 1), ((0, 1),) * 3),
                GriddedPerm((0, 2, 1), ((1, 0),) * 3),
            ],
            requirements=[
                [
                    GriddedPerm((0,), ((0, 1),)),
                    GriddedPerm((2, 0, 1), ((1, 0),) * 3),
                ],
                [GriddedPerm((0,), ((1, 0),))],
            ],
        )
        assert not LocalEnumeration(t).verified()

    def test_get_genf(self, enum_verified):
        with pytest.raises(NotImplementedError):
            enum_verified.get_genf()

    def test_get_genf_not_verified(self, enum_not_verified):
        with pytest.raises(InvalidOperationError):
            enum_not_verified.get_genf()

    def test_1x1_verified(self, onebyone_enum):
        assert onebyone_enum.verified()

    def test_no_req_option(self, enum_no_req):
        assert not enum_no_req.verified()


class TestMonotoneTreeEnumeration(CommonTest):
    @pytest.fixture
    def enum_verified(self):
        t = Tiling(
            obstructions=[
                GriddedPerm((0, 1), ((0, 0),) * 2),
                GriddedPerm((0, 1), ((0, 1),) * 2),
                GriddedPerm((0, 1), ((0, 2),) * 2),
                GriddedPerm((0, 1), ((2, 0),) * 2),
                GriddedPerm((0, 1, 2), ((1, 1),) * 3),
            ]
        )
        return MonotoneTreeEnumeration(t)

    @pytest.fixture
    def enum_not_verified(self):
        t = Tiling(
            obstructions=[
                GriddedPerm((0, 1), ((0, 0),) * 2),
                GriddedPerm((0, 1), ((0, 1),) * 2),
                GriddedPerm((0, 1), ((0, 2),) * 2),
                GriddedPerm((0, 1), ((2, 0),) * 2),
                GriddedPerm((0, 1), ((2, 2),) * 2),
                GriddedPerm((0, 1, 2), ((1, 1),) * 3),
            ]
        )
        return MonotoneTreeEnumeration(t)

    @pytest.fixture
    def enum_with_list_req(self):
        t = Tiling(
            obstructions=[
                GriddedPerm((0, 1), ((0, 0), (0, 0))),
                GriddedPerm((0, 1), ((1, 0), (1, 0))),
            ],
            requirements=[
                [
                    GriddedPerm((0,), ((0, 0),)),
                    GriddedPerm((0,), ((1, 0),)),
                ]
            ],
        )
        return MonotoneTreeEnumeration(t)

    @pytest.fixture
    def onebyone_enum(self):
        return MonotoneTreeEnumeration(Tiling.from_string("123"))

    @pytest.fixture
    def enum_with_crossing(self):
        t = Tiling(
            obstructions=[
                GriddedPerm((0, 1), ((0, 0),) * 2),
                GriddedPerm((0, 1), ((0, 1),) * 2),
                GriddedPerm((0, 1), ((0, 2),) * 2),
                GriddedPerm((0, 1), ((2, 0),) * 2),
                GriddedPerm((0, 1), ((0, 0), (0, 1))),
                GriddedPerm((0, 1, 2), ((1, 1),) * 3),
            ]
        )
        return MonotoneTreeEnumeration(t)

    def test_visited_cells_aligned(self, enum_verified):
        visited = {(1, 1), (0, 1)}
        assert sorted(enum_verified._visted_cells_aligned((0, 2), visited)) == [(0, 1)]

    def test_cell_tree_traversal(self, enum_verified):
        order = list(enum_verified._cell_tree_traversal((1, 1)))
        assert len(order) == 4
        assert (1, 1) not in order
        assert order[0] == (0, 1)
        assert order[3] == (2, 0)
        assert set(order[1:3]) == {(0, 0), (0, 2)}

    def test_not_verified(self, enum_with_list_req, onebyone_enum, enum_with_crossing):
        assert not enum_with_crossing.verified()
        assert not enum_with_list_req.verified()
        assert not onebyone_enum.verified()
        forest_tiling = Tiling(
            obstructions=[
                GriddedPerm((0,), ((0, 0),)),
                GriddedPerm((0,), ((1, 1),)),
                GriddedPerm((0,), ((2, 1),)),
                GriddedPerm((0, 1), ((1, 0), (1, 0))),
                GriddedPerm((0, 1), ((2, 0), (2, 0))),
                GriddedPerm((0, 1, 2), ((0, 1), (0, 1), (0, 1))),
            ],
            requirements=[[GriddedPerm((0,), ((0, 1),))]],
        )
        assert not MonotoneTreeEnumeration(forest_tiling).verified()

    def test_get_genf(self, enum_verified):
        x = sympy.Symbol("x")
        expected_gf = -(
            sympy.sqrt(-(4 * x**3 - 14 * x**2 + 8 * x - 1) / (2 * x**2 - 4 * x + 1)) - 1
        ) / (2 * x * (x**2 - 3 * x + 1))
        assert sympy.simplify(enum_verified.get_genf() - expected_gf) == 0
        t = Tiling(
            obstructions=[
                GriddedPerm((0, 1), ((0, 0),) * 2),
                GriddedPerm((0, 1), ((1, 0),) * 2),
            ]
        )
        enum_no_start = MonotoneTreeEnumeration(t)
        expected_gf = -1 / ((x - 1) * (x / (x - 1) + 1))
        assert sympy.simplify(enum_no_start.get_genf() - expected_gf) == 0

    def test_get_genf_simple(self):
        t = Tiling(
            obstructions=[
                GriddedPerm((0, 1), ((0, 0),) * 2),
                GriddedPerm((1, 0), ((1, 0),) * 2),
            ]
        )
        enum = MonotoneTreeEnumeration(t)
        print(t)
        assert enum.verified()
        assert sympy.simplify(enum.get_genf() - sympy.sympify("1/(1-2*x)")) == 0

    def test_with_finite_monotone_cell(self):
        t = Tiling(
            obstructions=[
                GriddedPerm((0, 1), ((0, 0),) * 2),
                GriddedPerm((1, 0), ((0, 0),) * 2),
                GriddedPerm((0, 1), ((1, 0),) * 2),
                GriddedPerm((1, 0), ((1, 0),) * 2),
            ]
        )
        enum = MonotoneTreeEnumeration(t)
        print(t)
        assert enum.verified()
        assert enum.get_genf().expand() == sympy.sympify("1+2*x+2*x**2")

    def test_with_finite_monotone_cell2(self):
        t = Tiling(
            obstructions=[
                GriddedPerm((0, 1), ((0, 0),) * 2),
                GriddedPerm((1, 0), ((0, 1),) * 2),
                GriddedPerm((0, 1), ((0, 1),) * 2),
                GriddedPerm((1, 0), ((1, 1),) * 2),
            ]
        )
        enum = MonotoneTreeEnumeration(t)
        print(t)
        assert enum.verified()
        assert (
            sympy.sympify("x/(1-x)**4 + 1/(1-x)**2") - enum.get_genf()
        ).simplify() == 0

    def test_interleave_fixed_length(self, enum_verified):
        track_var = MonotoneTreeEnumeration._tracking_var
        cell_var = enum_verified._cell_variable((1, 0))
        dummy_var = enum_verified._cell_variable((0, 0))
        x = sympy.var("x")
        F = x**8 * track_var**3 * dummy_var**3
        assert (
            enum_verified._interleave_fixed_length(F, (1, 0), 1)
            == 4 * x**9 * dummy_var**3 * cell_var**1
        )
        assert (
            enum_verified._interleave_fixed_length(F, (1, 0), 3)
            == 20 * x**11 * dummy_var**3 * cell_var**3
        )
        assert (
            enum_verified._interleave_fixed_length(F, (1, 0), 0) == x**8 * dummy_var**3
        )

    def test_interleave_fixed_lengths(self, enum_verified):
        track_var = MonotoneTreeEnumeration._tracking_var
        cell_var = enum_verified._cell_variable((1, 0))
        dummy_var = enum_verified._cell_variable((0, 0))
        x = sympy.var("x")
        F = x**8 * track_var**3 * dummy_var**3
        assert (
            enum_verified._interleave_fixed_lengths(F, (1, 0), 1, 1)
            == 4 * x**9 * dummy_var**3 * cell_var**1
        )
        assert (
            enum_verified._interleave_fixed_lengths(F, (1, 0), 3, 3)
            == 20 * x**11 * dummy_var**3 * cell_var**3
        )
        assert (
            enum_verified._interleave_fixed_lengths(F, (1, 0), 0, 0)
            == x**8 * dummy_var**3
        )
        assert (
            enum_verified._interleave_fixed_lengths(F, (1, 0), 0, 2)
            == x**8 * dummy_var**3
            + 4 * x**9 * dummy_var**3 * cell_var**1
            + 10 * x**10 * dummy_var**3 * cell_var**2
        )
        assert (
            enum_verified._interleave_fixed_lengths(F, (1, 0), 1, 3)
            == 4 * x**9 * dummy_var**3 * cell_var**1
            + 10 * x**10 * dummy_var**3 * cell_var**2
            + 20 * x**11 * dummy_var**3 * cell_var**3
        )

    def test_genf_with_req(self):
        t = Tiling(
            obstructions=[
                GriddedPerm((0, 1), ((0, 0),) * 2),
                GriddedPerm((0, 1), ((1, 0),) * 2),
            ],
            requirements=[
                [GriddedPerm((1, 0), ((0, 0),) * 2)],
                [GriddedPerm((0,), ((1, 0),))],
            ],
        )
        enum = MonotoneTreeEnumeration(t)
        print(t)
        assert enum.verified()
        genf = enum.get_genf().expand()
        terms = [0, 0, 0, 3, 10, 25, 56, 119, 246, 501, 1012]
        assert taylor_expand(genf) == terms

    def test_genf_with_big_finite_cell(self):
        t = Tiling(
            obstructions=[
                GriddedPerm((0, 1), ((0, 0),) * 2),
                GriddedPerm((0, 1), ((1, 0),) * 2),
                GriddedPerm((3, 2, 1, 0), ((0, 0),) * 4),
                GriddedPerm((3, 2, 1, 0), ((1, 0),) * 4),
            ]
        )
        enum = MonotoneTreeEnumeration(t)
        print(t)
        assert enum.verified()
        genf = enum.get_genf().expand()
        x = sympy.var("x")
        assert (
            genf == 1 + 2 * x + 4 * x**2 + 8 * x**3 + 14 * x**4 + 20 * x**5 + 20 * x**6
        )

    def test_with_two_reqs(self):
        t = Tiling(
            obstructions=(
                GriddedPerm((0,), ((1, 1),)),
                GriddedPerm((0, 1), ((0, 0), (0, 0))),
                GriddedPerm((0, 1), ((0, 1), (0, 1))),
                GriddedPerm((0, 1), ((1, 0), (1, 0))),
                GriddedPerm((1, 0), ((0, 1), (0, 1))),
            ),
            requirements=(
                (GriddedPerm((0,), ((0, 0),)),),
                (GriddedPerm((0,), ((0, 1),)),),
            ),
        )
        enum = MonotoneTreeEnumeration(t)
        expected_enum = [0, 0, 2, 7, 19, 47, 111, 255, 575, 1279, 2815]
        assert enum.verified()
        assert taylor_expand(enum.get_genf()) == expected_enum

    def test_corner(self):
        t = Tiling(
            obstructions=(
                GriddedPerm((0,), ((1, 1),)),
                GriddedPerm((0, 1), ((0, 0), (0, 0))),
                GriddedPerm((0, 1), ((0, 1), (0, 1))),
                GriddedPerm((0, 1), ((1, 0), (1, 0))),
            ),
            requirements=((GriddedPerm((0,), ((0, 0),)),),),
        )
        enum = MonotoneTreeEnumeration(t)
        expected_enum = [0, 1, 5, 17, 50, 138, 370, 979, 2575, 6755, 17700]
        assert enum.verified()
        assert taylor_expand(enum.get_genf()) == expected_enum
