import pytest
from spacy.tokens import Span

import edsnlp.data
from edsnlp.data.converters import (
    FILENAME,
    get_dict2doc_converter,
    get_doc2dict_converter,
)


@pytest.fixture(autouse=True, scope="module")
def set_extensions():
    if not Span.has_extension("negation"):
        Span.set_extension("negation", default=None)


def test_read_omop_dict(blank_nlp):
    json = {
        "note_id": 1234,
        "note_text": "This is a test.",
        "entities": [
            {
                "note_nlp_id": 0,
                "start_char": 0,
                "end_char": 4,
                "lexical_variant": "This",
                "note_nlp_source_value": "test",
                "negation": True,
            },
            {
                "note_nlp_id": 1,
                "start_char": 5,
                "end_char": 7,
                "lexical_variant": "is",
                "note_nlp_source_value": "test",
            },
        ],
    }
    doc = get_dict2doc_converter(
        "omop",
        dict(
            nlp=blank_nlp,
            span_attributes="negation",
            bool_attributes="negation",
        ),
    )[0](json)
    assert doc.text == "This is a test."
    assert doc._.note_id == 1234
    assert len(doc.ents) == 2
    assert doc.ents[0].text == "This"
    assert doc.ents[0]._.negation is True
    assert doc.ents[1]._.negation is False


def test_read_standoff_dict(blank_nlp):
    json = {
        "doc_id": 1234,
        "text": "This is a test.",
        "entities": [
            {
                "entity_id": 0,
                "fragments": [
                    {
                        "begin": 0,
                        "end": 4,
                    }
                ],
                "attributes": {
                    "negation": True,
                },
                "label": "test",
            },
            {
                "entity_id": 1,
                "fragments": [
                    {
                        "begin": 5,
                        "end": 7,
                    }
                ],
                "attributes": {},
                "label": "test",
            },
        ],
    }
    doc = get_dict2doc_converter(
        "standoff",
        dict(
            nlp=blank_nlp,
            span_attributes={"negation": "negation"},
            bool_attributes="negation",
        ),
    )[0](json)
    assert doc.text == "This is a test."
    assert doc._.note_id == 1234
    assert len(doc.ents) == 2
    assert doc.ents[0].text == "This"
    assert doc.ents[0]._.negation is True
    assert doc.ents[1]._.negation is False


def test_write_omop_dict(blank_nlp):
    doc = blank_nlp("This is a test.")
    doc._.note_id = 1234
    doc.ents = [Span(doc, 0, 1, label="test"), Span(doc, 1, 2, label="test")]
    doc.ents[0]._.negation = True
    doc.ents[1]._.negation = False
    json = {
        FILENAME: 1234,
        "note_id": 1234,
        "note_text": "This is a test.",
        "entities": [
            {
                "note_nlp_id": 0,
                "start_char": 0,
                "end_char": 4,
                "lexical_variant": "This",
                "note_nlp_source_value": "test",
                "sent.text": "This is a test.",
                "negation": True,
            },
            {
                "note_nlp_id": 1,
                "start_char": 5,
                "end_char": 7,
                "lexical_variant": "is",
                "note_nlp_source_value": "test",
                "sent.text": "This is a test.",
                "negation": False,
            },
        ],
    }
    assert (
        get_doc2dict_converter(
            "omop",
            dict(
                span_getter={"ents": True},
                span_attributes=["negation", "sent.text"],
            ),
        )[0](doc)
        == json
    )


def test_write_standoff_dict(blank_nlp):
    doc = blank_nlp("This is a test.")
    doc._.note_id = 1234
    doc.ents = [Span(doc, 0, 1, label="test"), Span(doc, 1, 2, label="test")]
    if not Span.has_extension("negation"):
        Span.set_extension("negation", default=None)
    doc.ents[0]._.negation = True
    doc.ents[1]._.negation = False
    json = {
        FILENAME: 1234,
        "doc_id": 1234,
        "text": "This is a test.",
        "entities": [
            {
                "entity_id": 0,
                "fragments": [
                    {
                        "begin": 0,
                        "end": 4,
                    }
                ],
                "attributes": {
                    "negation": True,
                },
                "label": "test",
            },
            {
                "entity_id": 1,
                "fragments": [
                    {
                        "begin": 5,
                        "end": 7,
                    }
                ],
                "attributes": {
                    "negation": False,
                },
                "label": "test",
            },
        ],
    }
    assert (
        get_doc2dict_converter(
            "standoff",
            dict(
                span_getter={"ents": True},
                span_attributes={"negation": "negation"},
            ),
        )[0](doc)
        == json
    )


def test_write_ents_dict(blank_nlp):
    doc = blank_nlp("This is a test.")
    doc._.note_id = 1234
    doc.ents = [Span(doc, 0, 1, label="test"), Span(doc, 1, 2, label="test")]
    doc.ents[0]._.negation = True
    doc.ents[1]._.negation = False
    jsons = [
        {
            "note_id": 1234,
            "start": 0,
            "end": 4,
            "lexical_variant": "This",
            "label": "test",
            "span_type": "ents",
            "sent.text": "This is a test.",
            "negation": True,
        },
        {
            "note_id": 1234,
            "start": 5,
            "end": 7,
            "lexical_variant": "is",
            "label": "test",
            "span_type": "ents",
            "sent.text": "This is a test.",
            "negation": False,
        },
    ]
    assert (
        get_doc2dict_converter(
            "ents",
            dict(
                span_getter={"ents": True},
                span_attributes=["negation", "sent.text"],
            ),
        )[0](doc)
        == jsons
    )


def test_unknown_converter():
    with pytest.raises(ValueError):
        get_dict2doc_converter("test", {})

    with pytest.raises(ValueError):
        get_doc2dict_converter("test", {})


def test_callable_converter():
    raw = lambda x: x  # noqa: E731
    assert get_dict2doc_converter(raw, {}) == (raw, {})
    assert get_doc2dict_converter(raw, {}) == (raw, {})


def test_method_converter(blank_nlp):
    data = ["Ceci", "est", "un", "test"]
    texts = list(
        edsnlp.data.from_iterable(data, converter=blank_nlp.make_doc).map(
            lambda x: x.text
        )
    )
    assert texts == data


def test_converter_types(blank_nlp):
    class Text:
        def __init__(self, text):
            self.text = text

    for converter in (blank_nlp.make_doc, Text, lambda x, k=2: Text(x)):
        data = ["Ceci", "est", "un", "test"]
        texts = list(
            edsnlp.data.from_iterable(data, converter=blank_nlp.make_doc).map(
                lambda x: x.text
            )
        )
        assert texts == data