from flair.data import Sentence

from deidentify.base import Annotation, Document
from deidentify.dataset.corpus_loader import DUMMY_CORPUS, CorpusLoader
from deidentify.methods.bilstmcrf import flair_utils
from deidentify.tokenizer import TokenizerFactory


def test_standoff_to_flair_sents():
    corpus = CorpusLoader().load_corpus(DUMMY_CORPUS)
    tokenizer = TokenizerFactory().tokenizer('ons')
    docs = corpus.train
    sents, parsed_docs = flair_utils.standoff_to_flair_sents(docs, tokenizer)

    assert len(sents) == 14
    assert len(parsed_docs) == 14

    bio_tags = [token.get_tag('ner').value for token in sents[0]]
    token_texts = [token.text for token in sents[0]]
    assert token_texts == [
        'Linders',
        ',',
        'Xandro',
        '<'
    ]
    assert bio_tags == [
        'B-Name',
        'I-Name',
        'I-Name',
        'O'
    ]

    bio_tags = [token.get_tag('ner').value for token in sents[1]]
    token_texts = [token.text for token in sents[1]]
    assert token_texts == [
        't.njg.nmmeso@rcrmb.nl'
    ]
    assert bio_tags == [
        'B-Email'
    ]

    bio_tags = [token.get_tag('ner').value for token in sents[2]]
    token_texts = [token.text for token in sents[2]]
    assert token_texts == [
        '>',
        '<SPACE>',
        '07',
        'apr',
        '.',
        '<SPACE>'
    ]
    assert bio_tags == [
        'O',
        'O',
        'B-Date',
        'I-Date',
        'O',
        'O',
    ]


def test_flair_sents_to_standoff():
    corpus = CorpusLoader().load_corpus(DUMMY_CORPUS)
    tokenizer = TokenizerFactory().tokenizer('ons')
    docs_expected = corpus.train

    sents, parsed_docs = flair_utils.standoff_to_flair_sents(docs_expected, tokenizer)
    docs_actual = flair_utils.flair_sents_to_standoff(sents, parsed_docs)

    assert len(docs_actual) == 1
    assert len(docs_expected) == 1

    assert len(docs_actual[0].annotations) == 16
    assert len(docs_expected[0].annotations) == 16

    for ann_expected, ann_actual in zip(docs_expected[0].annotations, docs_actual[0].annotations):
        assert ann_expected.text == ann_actual.text
        assert ann_expected.tag == ann_actual.tag


def test_filtered_corpus():
    def ignore_sentence(sent):
        return sent[0].text.startswith('===')

    filtered_corpus = flair_utils.FilteredCorpus(
        train=[Sentence('=== Answer: 123 ==='), Sentence('this is should be included')],
        dev=[Sentence('this is should be included'), Sentence('=== Answer: 456 ===')],
        test=[Sentence('this is should be included'), Sentence('and this as well')],
        ignore_sentence=ignore_sentence
    )

    assert len(filtered_corpus.train) == 1
    assert filtered_corpus.train[0].to_plain_string() == 'this is should be included'
    assert len(filtered_corpus.dev) == 1
    assert filtered_corpus.dev[0].to_plain_string() == 'this is should be included'
    assert len(filtered_corpus.test) == 2
    assert filtered_corpus.test[0].to_plain_string() == 'this is should be included'
    assert filtered_corpus.test[1].to_plain_string() == 'and this as well'

    assert len(filtered_corpus.train_ignored) == 1
    assert filtered_corpus.train_ignored[0].to_plain_string() == '=== Answer: 123 ==='
    assert len(filtered_corpus.dev_ignored) == 1
    assert filtered_corpus.dev_ignored[0].to_plain_string() == '=== Answer: 456 ==='
    assert len(filtered_corpus.test_ignored) == 0


def test_flair_sentence_with_whitespace_tokens():
    text = 'Mw geniet zichtbaar.  Maarten is de afgelopen periode veelal afwezig.'
    annotation = Annotation(
        text='Maarten',
        start=text.index('Maarten'),
        end=text.index('Maarten') + len('Maarten'),
        tag='PERSON'
    )
    doc = Document(name='', text=text, annotations=[annotation])

    tokenizer = TokenizerFactory().tokenizer('ons')
    flair_sents, docs = flair_utils.standoff_to_flair_sents([doc], tokenizer)

    # spaCy adds consecutive whitespace tokens as a single whitespace. These should be retained
    # in the Flair sentence, otherwise it's not possible to reconstruct the original document from
    # the tokenized representation.
    assert [token.text for token in flair_sents[0]] == ['Mw', 'geniet', 'zichtbaar', '.', '<SPACE>']

    spacy_doc = docs[0].spacy_doc
    spacy_sents = list(spacy_doc.sents)
    assert len(flair_sents) == 2
    assert len(spacy_sents) == 2

    assert len(flair_sents[0]) == 5
    assert len(spacy_sents[0]) == 5
    assert len(flair_sents[1]) == 8
    assert len(spacy_sents[1]) == 8