import pytest

from einops import EinopsError
from einops.parsing import ParsedExpression, AnonymousAxis, _ellipsis

__author__ = "Alex Rogozhnikov"


class AnonymousAxisPlaceholder:
    def __init__(self, value: int):
        self.value = value
        assert isinstance(self.value, int)

    def __eq__(self, other):
        return isinstance(other, AnonymousAxis) and self.value == other.value


def test_anonymous_axes():
    a, b = AnonymousAxis("2"), AnonymousAxis("2")
    assert a != b
    c, d = AnonymousAxisPlaceholder(2), AnonymousAxisPlaceholder(3)
    assert a == c and b == c
    assert a != d and b != d
    assert [a, 2, b] == [c, 2, c]


def test_elementary_axis_name():
    for name in [
        "a",
        "b",
        "h",
        "dx",
        "h1",
        "zz",
        "i9123",
        "somelongname",
        "Alex",
        "camelCase",
        "u_n_d_e_r_score",
        "unreasonablyLongAxisName",
    ]:
        assert ParsedExpression.check_axis_name(name)

    for name in ["", "2b", "12", "_startWithUnderscore", "endWithUnderscore_", "_", "...", _ellipsis]:
        assert not ParsedExpression.check_axis_name(name)


def test_invalid_expressions():
    # double ellipsis should raise an error
    ParsedExpression("... a b c d")
    with pytest.raises(EinopsError):
        ParsedExpression("... a b c d ...")
    with pytest.raises(EinopsError):
        ParsedExpression("... a b c (d ...)")
    with pytest.raises(EinopsError):
        ParsedExpression("(... a) b c (d ...)")

    # double/missing/enclosed parenthesis
    ParsedExpression("(a) b c (d ...)")
    with pytest.raises(EinopsError):
        ParsedExpression("(a)) b c (d ...)")
    with pytest.raises(EinopsError):
        ParsedExpression("(a b c (d ...)")
    with pytest.raises(EinopsError):
        ParsedExpression("(a) (()) b c (d ...)")
    with pytest.raises(EinopsError):
        ParsedExpression("(a) ((b c) (d ...))")

    # invalid identifiers
    ParsedExpression("camelCase under_scored cApiTaLs ß ...")
    with pytest.raises(EinopsError):
        ParsedExpression("1a")
    with pytest.raises(EinopsError):
        ParsedExpression("_pre")
    with pytest.raises(EinopsError):
        ParsedExpression("...pre")
    with pytest.raises(EinopsError):
        ParsedExpression("pre...")


def test_parse_expression():
    parsed = ParsedExpression("a1  b1   c1    d1")
    assert parsed.identifiers == {"a1", "b1", "c1", "d1"}
    assert parsed.composition == [["a1"], ["b1"], ["c1"], ["d1"]]
    assert not parsed.has_non_unitary_anonymous_axes
    assert not parsed.has_ellipsis

    parsed = ParsedExpression("() () () ()")
    assert parsed.identifiers == set()
    assert parsed.composition == [[], [], [], []]
    assert not parsed.has_non_unitary_anonymous_axes
    assert not parsed.has_ellipsis

    parsed = ParsedExpression("1 1 1 ()")
    assert parsed.identifiers == set()
    assert parsed.composition == [[], [], [], []]
    assert not parsed.has_non_unitary_anonymous_axes
    assert not parsed.has_ellipsis

    aap = AnonymousAxisPlaceholder

    parsed = ParsedExpression("5 (3 4)")
    assert len(parsed.identifiers) == 3 and {i.value for i in parsed.identifiers} == {3, 4, 5}
    assert parsed.composition == [[aap(5)], [aap(3), aap(4)]]
    assert parsed.has_non_unitary_anonymous_axes
    assert not parsed.has_ellipsis

    parsed = ParsedExpression("5 1 (1 4) 1")
    assert len(parsed.identifiers) == 2 and {i.value for i in parsed.identifiers} == {4, 5}
    assert parsed.composition == [[aap(5)], [], [aap(4)], []]

    parsed = ParsedExpression("name1 ... a1 12 (name2 14)")
    assert len(parsed.identifiers) == 6
    assert parsed.identifiers.difference({"name1", _ellipsis, "a1", "name2"}).__len__() == 2
    assert parsed.composition == [["name1"], _ellipsis, ["a1"], [aap(12)], ["name2", aap(14)]]
    assert parsed.has_non_unitary_anonymous_axes
    assert parsed.has_ellipsis
    assert not parsed.has_ellipsis_parenthesized

    parsed = ParsedExpression("(name1 ... a1 12) name2 14")
    assert len(parsed.identifiers) == 6
    assert parsed.identifiers.difference({"name1", _ellipsis, "a1", "name2"}).__len__() == 2
    assert parsed.composition == [["name1", _ellipsis, "a1", aap(12)], ["name2"], [aap(14)]]
    assert parsed.has_non_unitary_anonymous_axes
    assert parsed.has_ellipsis
    assert parsed.has_ellipsis_parenthesized