import json

import re

from pathlib import Path

from zensols.nlp import FeatureDocument, FeatureDocumentParser, FeatureToken

from util import TestBase

class TestMultiEntity(TestBase):

    def test_multi_entity(self):

        def filter_json_line(s: str) -> bool:

            return re.match(r'^\s*"context_similarity":\s*-?1,?\s*$', s) is None

        DEBUG: bool = False

        WRITE: bool = False

        sent = 'John was diagnosed with kidney failure. He has lung cancer too.'

        parser: FeatureDocumentParser = \

            self._get_doc_parser('mednlp-add-linker')

        doc: FeatureDocument = parser.parse(sent)

        if DEBUG:

            tok: FeatureToken

            for tok in doc.token_iter():

                print(tok, tok.definition_)

                tok.write()

        path = Path('test-resources/entlink-features.json')

        json_str = doc.asjson(indent=4)

        obj = json.loads(json_str)

        for s in obj['sentences']:

            for t in s['tokens']:

                del t['context_similarity']

        if WRITE:

            with open(path, 'w') as f:

                line: str

                # enable to re-write `should` test data for API changes; but

                # have to remove all `context_simirity` entries

                for line in filter(filter_json_line, json_str.split('\n')):

                    f.write(line + '\n')

        with open(path) as f:

            should = json.load(f)

        self.assertEqual(should, obj)