from typing import Union, Tuple

import numpy
import pandas

from lexnlp.extract.ml.detector.detecting_settings import DetectingSettings
from lexnlp.extract.ml.detector.sample_processor import process_sample, get_target_start_end_from_corgetes
from lexnlp.extract.ml.detector.artifact_detector import ArtifactDetector

__author__ = "ContraxSuite, LLC; LexPredict, LLC"
__copyright__ = "Copyright 2015-2019, ContraxSuite, LLC"
__license__ = "https://github.com/LexPredict/lexpredict-lexnlp/blob/master/LICENSE"
__version__ = "1.4.0"
__maintainer__ = "LexPredict, LLC"
__email__ = "support@contraxsuite.com"


class DefinitionPhraseDetector(ArtifactDetector):
    """
    Search for the phrase surrounding the term being defined

    Let the prase be <agrees to serve the Company in such capacity during the
    term of employment (the "Employment Period”).

    ... model_definition will find <term of employment (the "Employment Period”)>
    """

    def process_sample(self,
                       sample_df: pandas.DataFrame,
                       build_target_data: bool = False) -> Union[numpy.ndarray, Tuple[numpy.ndarray, numpy.ndarray]]:
        return process_sample(sample_df, self.model,
                              get_target_start_end=get_target_start_end_from_corgetes,
                              column_name_formatted='labels',
                              feature_mask_column='feature_mask')

    def train_and_save(self,
                       settings: DetectingSettings,
                       train_file: str,
                       train_size: int = -1,
                       save_path: str = '',
                       compress: bool = False) -> None:
        # load sample
        train_sample_df = self.read_sample_df(train_file, train_size)

        self.train_and_save_on_dataframe(settings,
                                         train_sample_df,
                                         save_path,
                                         compress)

    def train_and_save_on_dataframe(
            self,
            settings: DetectingSettings,
            train_sample_df: pandas.DataFrame,
            save_path: str = '',
            compress: bool = False) -> None:
        # setup token match strings in locale
        def_tokens = ['shall', 'includes', 'including',
                      'mean', 'meaning', 'referred', 'known',
                      'refers', 'used', 'hereby', 'interpreted',
                      'defined', 'interpreted', 'collectively',
                      'include', 'individually', 'together',
                      'being', 'foregoing', 'purpose', 'purposed']

        # create model class
        self.train_and_save_on_tokens(def_tokens, save_path,
                                      settings, train_sample_df,
                                      punc_set=":\"()",
                                      compress=compress)
