from typing import List, Dict, Set, Any

from collections import OrderedDict

import json

from zensols.nlp import FeatureToken, FeatureDocument, FeatureDocumentParser

from util import TestBase

class TestCombinedParsers(TestBase):

    _DEFAULT_ATTRS = 'i i_sent idx norm ent_ ent_iob ent_iob_'.split()

    _TRACE = False

    def _compare_sents(self, parser_name: str, idx: int, write: bool,

                       sent: str, attrs: List[str], missing: Set[str]):

        def map_tok_features(t: FeatureToken) -> Dict[str, Any]:

            # sort keys to make diffing easier

            dct = t.asdict()

            return OrderedDict(sorted(dct.items(), key=lambda t: t[0]))

        actual_file: str = f'test-resources/should/{parser_name}-{idx}.json'

        p: FeatureDocumentParser = self._get_doc_parser('combined', parser_name)

        doc: FeatureDocument = p(sent)

        actuals = tuple(map(map_tok_features, doc.token_iter()))

        if self._TRACE:

            from pprint import pprint

            pprint(actuals)

            return

        if write:

            print(sent)

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

                json.dump(actuals, f, indent=4, sort_keys=False)

            for attr in attrs:

                vals: str = ', '.join(map(lambda d: str(d[attr]), actuals))

                print(f'  {attr}: {vals}')

            print('_' * 79)

        with open(actual_file) as f:

            shoulds: List[Dict[str, Any]] = json.load(f)

        attr: str

        for attr in attrs:

            actual = tuple(map(lambda d: d[attr], actuals))

            should = tuple(map(lambda d: d[attr], shoulds))

            self.assertEqual(should, actual, f'for attribute: <{attr}>')

        if missing is not None:

            for attr in missing:

                for tok in actuals:

                    self.assertFalse(hasattr(tok, attr),

                                     f'expected missing {attr} in {tok}')

    def _compare(self, parser_name: str, write: bool = False,

                 attrs: List[str] = None, missing: Set[str] = None):

        attrs = TestCombinedParsers._DEFAULT_ATTRS if attrs is None else attrs

        for i in range(2):

            sent: str = getattr(self, f'text_{i + 1}')

            self._compare_sents(parser_name, i, write, sent, attrs, missing)

    def test_default(self):

        self._compare('doc_parser', missing='cui_'.split())

    def test_biomed_ner(self):

        self._compare('mednlp_biomed_doc_parser', missing='cui_'.split())

    def test_medcat(self):

        self._compare('mednlp_medcat_doc_parser',

                      attrs=self._DEFAULT_ATTRS + 'cui_ tuis_'.split())

    def test_biomed_combined(self):

        self._compare('mednlp_combine_biomed_doc_parser',

                      missing='cui_'.split())

    def test_medcat_combined(self):

        self._compare('mednlp_combine_medcat_doc_parser',

                      attrs=self._DEFAULT_ATTRS + 'cui_ tuis_'.split())

    def test_medcat_biomded_combined(self):

        self._compare('mednlp_combine_biomed_medcat_doc_parser',

                      attrs=self._DEFAULT_ATTRS + 'cui_ tuis_'.split())