from os.path import dirname, join

from deidentify.base import Annotation

from deidentify.dataset.brat import load_brat_document

from deidentify.surrogates.dataset_deidentifier import DatasetDeidentifier, Document

def _load_documents():

    doc_names = ['example-1', 'example-2']

    docs = []

    for doc in doc_names:

        annotations, text = load_brat_document(join(dirname(__file__), 'data/original'), doc)

        docs.append(Document(annotations, text))

    return docs

def test_dataset_deidentifier():

    docs = _load_documents()

    dataset_deidentifier = DatasetDeidentifier()

    docs = dataset_deidentifier.generate_surrogates(docs)

    print()

    for doc in docs:

        print('===================')

        for annotation, surrogate in zip(*doc.annotation_surrogate_pairs()):

            print('{:<30} => {:<20} ({})'.format(annotation.text, surrogate, annotation.tag))

            if annotation.tag == 'Other' or annotation.tag == 'Age':

                # other will be replaced manually, and there is no age above 89 in the test corpus

                assert annotation.text == surrogate

            else:

                assert annotation.text != surrogate

def test_generate_surrogates_without_choices():

    text = 'Patient is being treated at UMCU.'

    annotations = [Annotation('UMCU', text.index('UMCU'), text.index('UMCU') + 4, 'Hospital')]

    doc = Document(annotations, text)

    surrogate_doc = DatasetDeidentifier().generate_surrogates([doc])[0]

    original_annotations, surrogates = surrogate_doc.annotation_surrogate_pairs()

    assert len(original_annotations) == 1

    assert len(surrogates) == 1

    assert original_annotations[0].text == 'UMCU'

    assert surrogates[0] == 'UMCU'

def test_generate_surrogates_shuffle_choices():

    text = 'Patient is being treated at UMCU.'

    annotations = [Annotation('UMCU', text.index('UMCU'), text.index('UMCU') + 4, 'Hospital')]

    doc_1 = Document(annotations, text)

    text = 'Patient is being treated at MST.'

    annotations = [Annotation('MST', text.index('MST'), text.index('MST') + 3, 'Hospital')]

    doc_2 = Document(annotations, text)

    surrogate_docs = DatasetDeidentifier().generate_surrogates([doc_1, doc_2])

    original_annotations, surrogates = surrogate_docs[0].annotation_surrogate_pairs()

    assert len(original_annotations) == 1 and len(surrogates) == 1

    assert original_annotations[0].text == 'UMCU'

    assert surrogates[0] == 'MST'

    original_annotations, surrogates = surrogate_docs[1].annotation_surrogate_pairs()

    assert len(original_annotations) == 1 and len(surrogates) == 1

    assert original_annotations[0].text == 'MST'

    assert surrogates[0] == 'UMCU'