import re

import pytest

import asdf
from asdf.tests import helpers


def make_complex_asdf(string):
    yaml = """
a: !core/complex-1.0.0
  {}
    """.format(
        string
    )

    return helpers.yaml_to_asdf(yaml)


@pytest.mark.parametrize(
    "invalid",
    [
        "3 + 4i",
        "3+-4i",
        "3-+4i",
        "3i+4i",
        "X3+4iX",
        "3+X4i",
        "3+4",
        "3i+4" "3+4z",
        "3.+4i",
        "3+4.i",
        "3e-4.0+4i",
        "3+4e4.0i",
        "",
    ],
)
def test_invalid_complex(invalid):

    with pytest.raises(asdf.ValidationError):
        with asdf.open(make_complex_asdf(invalid)):
            pass


@pytest.mark.parametrize(
    "valid",
    [
        "3+4j",
        "(3+4j)",
        ".3+4j",
        "3+.4j",
        "3e10+4j",
        "3e-10+4j",
        "3+4e10j",
        "3.0+4j",
        "3+4.0j",
        "3.0+4.0j",
        "3+4e-10j",
        "3+4J",
        "3+4i",
        "3+4I",
        "inf",
        "inf+infj",
        "inf+infi",
        "infj",
        "infi",
        "INFi",
        "INFI",
        "3+infj",
        "inf+4j",
    ],
)
def test_valid_complex(valid):

    with asdf.open(make_complex_asdf(valid)) as af:
        assert af.tree["a"] == complex(re.sub(r"[iI]$", r"j", valid))


@pytest.mark.parametrize("valid", ["nan", "nan+nanj", "nan+nani", "nanj", "nani", "NANi", "NANI", "3+nanj", "nan+4j"])
def test_valid_nan_complex(valid):
    with asdf.open(make_complex_asdf(valid)):
        pass


def test_roundtrip(tmpdir):
    tree = {"a": 0 + 0j, "b": 1 + 1j, "c": -1 + 1j, "d": -1 - 1j}

    helpers.assert_roundtrip_tree(tree, tmpdir)
