
from nlu.pipe.extractors.extractor_base_data_classes import SparkNLPExtractor,SparkNLPExtractorConfig
from nlu.pipe.extractors.extractor_methods.helper_extractor_methods import *
from nlu.pipe.extractors.extractor_methods.base_extractor_methods import *

"""
This file contains methods to get pre-defined configurations for every annotator.
Extractor_resolver.py should be used to resolve SparkNLP Annotator classes to methods 
in this file, which return the corrosponding configs that need to be passed to 
the master_extractor() call.

This file is where all the in extractor_base_data_classes.py Dataclasses are combined with the 
extractors defined in extractor_methods.py.

"""

def default_get_nothing(output_col_prefix):
    return SparkNLPExtractorConfig(
        output_col_prefix   = output_col_prefix,
        name                = 'nothing_extractor',
        description         = 'Extracts nothing. Useful for annotators with irrelevant data'
    )
def default_only_result_config(output_col_prefix):
    return SparkNLPExtractorConfig(
        output_col_prefix   = output_col_prefix,
        get_result          = True,
        name                = 'Default result extractor',
        description         = 'Just gets the result field'
    )

def default_full_config(output_col_prefix='DEFAULT'):
    return SparkNLPExtractorConfig(
        output_col_prefix   = output_col_prefix,
        get_positions       = True,
        get_begin           = True,
        get_end             = True,
        get_embeds          = True,
        get_result          = True,
        get_meta            = True,
        get_full_meta       = True,
        get_annotator_type  = True,
        name                = 'default_full',
        description         = 'Default full configuration, keeps all data and gets all metadata fields',

    )




def default_NER_converter_licensed_config(output_col_prefix='entities'):
    """Extracts NER tokens withouth positions, just the converted IOB tags,confidences and classified tokens """
    return SparkNLPExtractorConfig(
        output_col_prefix   = output_col_prefix,
        get_result          = True,
        get_meta            = True,
        meta_white_list     = ['entity','confidence'], #sentence, chunk
        name                = 'default_ner',
        description         = 'Converts IOB-NER representation into entity representation and generates confidences for the entire entity chunk',
    )

def default_chunk_resolution_config(output_col_prefix='resolved_entities'):
    """Extracts NER tokens withouth positions, just the IOB tags,confidences and classified tokens """
    return SparkNLPExtractorConfig(
        output_col_prefix   = output_col_prefix,
        get_result          = True,
        get_meta            = True,
        meta_white_list     = ['confidence','resolved_text'], #sentence, chunk
        name                = 'default_ner',
        description         = 'Converts IOB-NER representation into entity representation and generates confidences for the entire entity chunk',
    )

def default_relation_extraction_positional_config(output_col_prefix='extracted_relations'):
    """Extracts NER tokens withouth positions, just the IOB tags,confidences and classified tokens """
    return SparkNLPExtractorConfig(
        output_col_prefix   = output_col_prefix,
        get_result          = True,
        meta_white_list     = [],
        get_meta            = True,
        meta_black_list     = ['entity1_begin','entity2_begin','entity1_end','entity2_end','entity1','entity2','chunk1','chunk2',],
        name                = 'default_relation_extraction',
        description         = 'Get relation extraction result and all metadata, positions of entities excluded',
    )


def default_relation_extraction_config(output_col_prefix='extracted_relations'):
    """Extracts NER tokens withouth positions, just the IOB tags,confidences and classified tokens """
    return SparkNLPExtractorConfig(
        output_col_prefix   = output_col_prefix,
        get_result          = True,
        meta_white_list     = [],
        get_meta            = True,
        meta_black_list     = ['entity1','entity2','chunk1','chunk2',],
        name                = 'default_relation_extraction',
        description         = 'Get relation extraction result and all metadata, positions of entities excluded',
    )




def default_de_identification_config(output_col_prefix='de_identified'):
    """Extracts NER tokens withouth positions, just the IOB tags,confidences and classified tokens """
    return SparkNLPExtractorConfig(
        output_col_prefix   = output_col_prefix,
        get_result          = True,
        name                = 'positional_relation_extraction',
        description         = 'Get relation extraction result and all metadata, which will include positions of entities chunks',
    )


def default_assertion_config(output_col_prefix='assertion'):
    """Extracts NER tokens withouth positions, just the IOB tags,confidences and classified tokens """
    return SparkNLPExtractorConfig(
        output_col_prefix   = output_col_prefix,
        get_result          = True,
        name                = 'default_assertion_extraction',
        get_meta            = True,
        meta_white_list     = ['confidence'],
        description         = 'Gets the assertion result and confidence',
    )

def default_ner_config(output_col_prefix='med_ner'):
    return default_only_result_config(output_col_prefix)

def default_ner_config(output_col_prefix='med_ner'):
    return default_get_nothing(output_col_prefix)


def default_feature_assembler_config(output_col_prefix='feature_assembler'):
    return SparkNLPExtractorConfig(
        output_col_prefix   = output_col_prefix,
        get_result          = False,
        name                = 'features_assembled',
        get_meta            = False,
        description         = 'Gets nothing',
    )

def default_generic_classifier_config(output_col_prefix='generic_classifier'):
    return SparkNLPExtractorConfig(
        output_col_prefix   = output_col_prefix,
        get_result          = True,
        name                = 'generic_classifier',
        get_meta            = True,
        meta_white_list     = ['confidence'],
        description         = 'Gets the  result and confidence',
    )