from typing import Dict

import json

from pathlib import Path

from zensols.nlp import (

    FeatureToken, FeatureSentence, FeatureDocument, FeatureDocumentParser

)

from util import TestBase

class TestParse(TestBase):

    def test_feature_parse(self):

        DEBUG: bool = False

        keeps = set('cui_ pref_name_'.split())

        parser: FeatureDocumentParser = self._get_doc_parser()

        self.assertTrue(isinstance(parser, FeatureDocumentParser))

        med_toks: Dict[str, str] = []

        doc: FeatureDocument = parser(self.text_1)

        for tok in doc.token_iter():

            fd = tok.asdict()

            med_toks.append({k: fd[k] for k in fd.keys() & keeps})

        if DEBUG:

            print()

            for tok in doc.token_iter():

                print(tok, tok.cui_, tok.pref_name_)

            print()

            for tok in med_toks:

                print(tok)

        none = FeatureToken.NONE

        for i, mtok in enumerate(med_toks):

            if DEBUG:

                print(i, mtok)

            if i >= 4 and i <= 5:

                self.assertEqual(

                    {'cui_': 'C0035078', 'pref_name_': 'Kidney Failure'}, mtok)

            else:

                self.assertEqual({'cui_': none, 'pref_name_': none}, mtok)

    def test_doc_parse(self):

        parser: FeatureDocumentParser = self._get_doc_parser()

        self.assertTrue(isinstance(parser, FeatureDocumentParser))

        doc: FeatureDocument = parser.parse(self.text_1)

        self.assertTrue(isinstance(doc, FeatureDocument))

        sent: FeatureSentence = doc[0]

        self.assertTrue(isinstance(sent, FeatureSentence))

        self.assertEqual(10, len(sent))

        self.assertTrue(isinstance(sent[0], FeatureToken))

        self.assertTrue('C0011900', sent[2].cui_)

        self.assertTrue('Diagnosis', sent[2].pref_name_)

        self.assertTrue('C0035078', sent[4].cui_)

        self.assertTrue('Kidney Failure', sent[4].pref_name_)

    def test_multi_entity(self):

        WRITE: bool = False

        parser: FeatureDocumentParser = self._get_doc_parser()

        doc: FeatureDocument = parser.parse(self.text_2)

        path = Path('test-resources/doc-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']

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

        # remove all `context_simirity` entries

        if WRITE:

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

                json.dump(obj, f, indent=4)

        with open(path) as f:

            should = json.load(f)

        self.assertEqual(should, obj)