import json
import os

import spacy

from medacy.pipeline_components.feature_extractors.discrete_feature_extractor import FeatureExtractor
from medacy.pipeline_components.feature_extractors.text_extractor import TextExtractor
from medacy.pipeline_components.feature_overlayers.metamap.metamap import MetaMap
from medacy.pipeline_components.feature_overlayers.metamap.metamap_all_types_component import MetaMapAllTypesOverlayer
from medacy.pipeline_components.feature_overlayers.metamap.metamap_component import MetaMapOverlayer
from medacy.pipeline_components.learners.bert_learner import BertLearner
from medacy.pipeline_components.learners.bilstm_crf_learner import BiLstmCrfLearner
from medacy.pipeline_components.learners.crf_learner import get_crf
from medacy.pipeline_components.tokenizers.character_tokenizer import CharacterTokenizer
from medacy.pipeline_components.tokenizers.clinical_tokenizer import ClinicalTokenizer
from medacy.pipeline_components.tokenizers.systematic_review_tokenizer import SystematicReviewTokenizer
from medacy.pipelines.base.base_pipeline import BasePipeline


required_keys = [
    'learner',
    'spacy_pipeline',
]


def json_to_pipeline(json_path):
    """
    Constructs a custom pipeline from a json file

    The json must have the following keys:

    'learner': 'CRF', 'BiLSTM', or 'BERT'
    'spacy_pipeline': the spaCy model to use

    The following keys are optional:
    'spacy_features': a list of features that exist as spaCy token annotations
    'window_size': the number of words +/- the target word whose features should be used along with the target word; defaults to 0
    'tokenizer': 'clinical', 'systematic_review', or 'character'; defaults to the spaCy model's tokenizer
    'metamap': the path to the MetaMap binary; MetaMap will only be used if this key is present
        if 'metamap' is a key, 'semantic_types' must also be a key, with value 'all', 'none', or
        a list of semantic type strings

    :param json_path: the path to the json file, or a dict of what that json would be
    :return: a custom pipeline class
    """

    if isinstance(json_path, (str, os.PathLike)):
        with open(json_path, 'rb') as f:
            input_json = json.load(f)
    elif isinstance(json_path, dict):
        input_json = json_path

    missing_keys = [key for key in required_keys if key not in input_json.keys()]
    if missing_keys:
        raise ValueError(f"Required key(s) '{missing_keys}' was/were not found in the json file.")

    class CustomPipeline(BasePipeline):
        """A custom pipeline configured from a JSON file"""

        def __init__(self, entities, **kwargs):
            super().__init__(entities, spacy_pipeline=spacy.load(input_json['spacy_pipeline']))

            if 'metamap' in input_json.keys():
                if 'semantic_types' not in input_json.keys():
                    raise ValueError("'semantic_types' must be a key when 'metamap' is a key.")

                metamap = MetaMap(input_json['metamap'])

                if input_json['semantic_types'] == 'all':
                    self.add_component(MetaMapAllTypesOverlayer, metamap)
                elif input_json['semantic_types'] == 'none':
                    self.add_component(MetaMapOverlayer, metamap, semantic_type_labels=[])
                elif isinstance(input_json['semantic_types'], list):
                    self.add_component(MetaMapOverlayer, metamap, semantic_type_labels=input_json['semantic_types'])
                else:
                    raise ValueError("'semantic_types' must be 'all', 'none', or a list of strings")

            # BERT values
            self.cuda_device = kwargs['cuda_device'] if 'cuda_device' in kwargs else -1
            self.batch_size = kwargs['batch_size'] if 'batch_size' in kwargs else 8
            self.learning_rate = kwargs['learning_rate'] if 'learning_rate' in kwargs else 1e-5
            self.epochs = kwargs['epochs'] if 'epochs' in kwargs else 3
            self.pretrained_model = kwargs['pretrained_model'] if 'pretrained_model' in kwargs else 'bert-large-cased'
            self.using_crf = kwargs['using_crf'] if 'using_crf' in kwargs else False

            # BiLSTM value
            if input_json['learner'] == 'BiLSTM':
                if 'word_embeddings' not in kwargs:
                    raise ValueError("BiLSTM learner requires word embeddings; use the parameter '--word_embeddings' "
                                     "to specify an embedding path")
            self.word_embeddings = kwargs['word_embeddings']

        def get_tokenizer(self):
            if 'tokenizer' not in input_json.keys():
                return None

            selection = input_json['tokenizer']
            options = {
                'clinical': ClinicalTokenizer,
                'systematic_review': SystematicReviewTokenizer,
                'character': CharacterTokenizer
            }

            if selection not in options:
                raise ValueError(f"Tokenizer selection '{selection}' not an option")

            Tokenizer = options[selection]
            return Tokenizer(self.spacy_pipeline)

        def get_learner(self):
            learner_selection = input_json['learner']

            if learner_selection == 'CRF':
                return "CRF_l2sgd", get_crf()
            if learner_selection == 'BiLSTM':
                return 'BiLSTM+CRF', BiLstmCrfLearner(self.word_embeddings, self.cuda_device)
            if learner_selection == 'BERT':
                learner = BertLearner(
                    self.cuda_device,
                    pretrained_model=self.pretrained_model,
                    batch_size=self.batch_size,
                    learning_rate=self.learning_rate,
                    epochs=self.epochs,
                    using_crf=self.using_crf
                )
                return 'BERT', learner
            else:
                raise ValueError(f"'learner' must be 'CRF', 'BiLSTM', or 'BERT', but is {learner_selection}")

        def get_feature_extractor(self):
            if input_json['learner'] == 'BERT':
                return TextExtractor()

            return FeatureExtractor(
                window_size=input_json['window_size'] if 'window_size' in input_json else 0,
                spacy_features=input_json['spacy_features'] if 'spacy_features' in input_json else ['text']
            )

        def get_report(self):
            report = super().get_report() + '\n'
            report += f"Pipeline configured from a JSON: {json.dumps(input_json)}\nJSON path: {json_path}"
            return report

    return CustomPipeline