"""Wrappers for different approaches to text classification, including scikit-learn text classification, text classification with Hugging Face Transformers, and few-shot classification (via SetFit)."""

# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/04_pipelines.classifier.ipynb.

# %% auto 0
__all__ = ['DATASET_TEXT', 'DATASET_LABEL', 'DEFAULT_SETFIT_MODEL', 'SMALL_SETFIT_MODEL', 'ClassifierBase', 'SKClassifier',
           'HFClassifier', 'FewShotClassifier']

# %% ../../nbs/04_pipelines.classifier.ipynb 3
from typing import List, Union

import warnings

with warnings.catch_warnings():
    warnings.filterwarnings("ignore",category=DeprecationWarning)
    from setfit import SetFitModel, TrainingArguments, Trainer, sample_dataset

import numpy as np

# %% ../../nbs/04_pipelines.classifier.ipynb 4
from abc import ABC, abstractmethod

DATASET_TEXT = "text"
DATASET_LABEL = "label"

class ClassifierBase(ABC):

  
    @abstractmethod
    def train(self,
              X:List[str],
              y:Union[List[int], List[str]],
              max_steps:int=50,
              num_epochs:int=10,
              batch_size:int=32,
              metric='accuracy',
              callbacks=None,
              **kwargs,
             ):
        """
        Trains the classifier on a list of texts (`X`) and a list of labels (`y`).
        Additional keyword arguments are passed directly to `SetFit.TrainingArguments`.

        **Args:**

        - *X*: List of texts
        - *y*: List of integers representing labels
        - *max_steps*: If set to a positive number, the total number of training steps to perform. Overrides num_epochs. 
        - *num_epochs*: Number of epochs to train
        - *batch_size*: Batch size
        - *metric*: metric to use
        - *callbacks*:  A list of callbacks to customize the training loop.

        **Returns:**

        - None
        """
        pass

    @abstractmethod
    def save(self, save_path:str):
        """
        Save model to specified folder path, `save_path`
        """
        pass

    
    def get_explain_predictor(self):
        """
        Get predictor and tokenizer used for shap predictions.
        """
        return None, None

    
    def get_labels(self):
        return self.model.labels if hasattr(self.model, 'labels') else self.labels

        
    def sample_examples(self, X:list, y:list, num_samples:int=8,
                        text_key:str=DATASET_TEXT, label_key:str=DATASET_LABEL):
        """
        Sample a dataset with `num_samples` per class
        """
        full_dataset = self.arrays2dataset(X, y, text_key=text_key, label_key=label_key)
                
        sample = sample_dataset(full_dataset, label_column=label_key, num_samples=num_samples)
        return sample.to_dict()[text_key], sample.to_dict()[label_key]
        

    
    def arrays2dataset(self, X:List[str], y:Union[List[int], List[str]], 
                       text_key:str=DATASET_TEXT, label_key:str=DATASET_LABEL):
        """
        Convert train or test examples to HF dataset
        """
        from datasets import Dataset
        return Dataset.from_dict({text_key:X, label_key:y})


    def dataset2arrays(self, dataset, text_key:str=DATASET_TEXT, label_key:str=DATASET_LABEL):
        """
        Convert a Hugging Face dataset to X, y arrays
        """
        return dataset.to_dict()['text'], dataset.to_dict()['label']
        
    def get_trainer(self):
        """
        Retrieves last trainer
        """
        if not self.trainer:
            raise ValueError('A trainer has not been created yet. You must first train a model on some labeled examples ' +\
                             'using the FewShotClassifier.train method.')
        return self.trainer

    
    def evaluate(self, X_eval:list, y_eval:list, print_report:bool=True, labels:list=[], **kwargs):
        """
        Evaluates labeled data using the trained model. 
        If `print_report` is True, prints classification report and returns nothing.
        Otherwise, returns and prints a dictionary of the results.
        Extra kwargs fed to `self.predict`.

        
        """
        labels = labels if labels else self.get_labels()
        labels = labels if labels else None
        
        from sklearn.metrics import classification_report
        y_pred= self.predict(X_eval, **kwargs)
        if isinstance(y_pred[0], str) and isinstance(y_eval[0], (np.integer, int)) and labels:
            y_eval = [labels[y] for y in y_eval]
        if isinstance(y_eval[0], str) and isinstance(y_pred[0], (np.integer, int)) and labels:
            y_pred = [labels[y] for y in y_pred]

        result = classification_report(y_eval, y_pred, 
                                       output_dict=not print_report,
                                       target_names = labels)
        if print_report:
            print(result)
            return
        else:
            import yaml
            #print(yaml.dump(result, allow_unicode=True, default_flow_style=False))
            return result


    def explain(self, X:list, labels:list=[]):
        """
        Explain the predictions on given examples in `X`. (Requires `shap` and `matplotlib` to be installed.)
        """
        X = [X] if isinstance(X, str) else X
        output_names = labels if labels else self.get_labels()
        output_names = output_names if output_names else None
        try:
            import shap
        except ImportError:
            raise ImportError('Please install the shap library: pip install shap')

        try:
            import matplotlib
        except ImportError:
            raise ImportError('Please install the matplotlib library: pip install matplotlib')

        f, tokenizer = self._get_explain_predictor()
        if f is None:
            raise NotImplementedError('Exlpanations are not currently supported for this model.')
        explainer = shap.Explainer(f, tokenizer, output_names=output_names)
        shap_values = explainer(X)
        shap.plots.text(shap_values)

# %% ../../nbs/04_pipelines.classifier.ipynb 10
class SKClassifier(ClassifierBase):
    def __init__(
        self,
        model_path=None,
        labels = [],
        **kwargs,
    ):
        """
        `SKClassifier` is a wrapper to scikit-learn text classifiation models.
        Extra kwargs are fed directly to `onprem.sk.clf.Classifier.create_model`.
        If no arguments are supplied, then a default Logistic Regression model is used.

        **Args:**

        - *model_path*: path to an already saved model file to be reloaded
        - *labels*: 
        - *labels*: list of strings intended to map label indices to string labels
        """

        from onprem.sk.clf import Classifier
        self.model = Classifier()
        self.labels = labels

        if model_path:
            self.model.load(model_path)
        else:
            # set defaults if necessary        
            if 'ctype' not in kwargs: 
                kwargs['ctype'] = 'sgdclassifier'
                kwargs['clf__loss'] = 'hinge'
                kwargs['clf__penalty'] = 'l2'
                kwargs['clf__alpha'] = 1e-3
                kwargs['clf__max_iter'] = 5
                kwargs['clf__tol'] = None
                kwargs['clf__random_state'] = 42
                kwargs['vec__token_pattern'] = '(?u)\\b\\w\\w+\\b'
                kwargs['use_tfidf'] = False

                # replace TF-IDF with this for slight accuracy boost
                # in relevant examples
                kwargs['vec__ngram_range'] = (1,3)
                kwargs['vec__binary'] = True
                kwargs['vec__max_features'] = 100000

            self.model.create_model(**kwargs) 

        
    def predict(self, X, **kwargs):
        """
        predict labels
        """
        labels = self.get_labels()
        preds = self.model.predict(X, **kwargs)
        preds = [preds] if not isinstance(preds, (list, np.ndarray)) else preds
        preds =  [labels[p] for p in preds] if labels else preds
        return preds[0] if len(preds) == 1 else preds

    
    def predict_proba(self, X, **kwargs):
        """
        predict label probabilities
        """
        return self.predict_proba(X, **kwargs)
    
    def train(self,
              X:List[str],
              y:Union[List[int], List[str]],
              **kwargs,
             ):
        """
        Trains the classifier on a list of texts (`X`) and a list of labels (`y`).
        Additional keyword arguments are passed directly to `self.model.fit`.

        **Args:**

        - *X*: List of texts
        - *y*: List representing labels

        **Returns:**

        - None
        """

        self.model.fit(X, y)                                                                                                                   
  

    def save(self, filename:str):
        """
        Save model to specified `filename` (e.g., `/tmp/mymodel.gz`).
        Model saved as pickle file.
        To reload the model, supply `model_path` when instantiating`SKClassifier`.
        """
        self.model.save(filename)    

        

# %% ../../nbs/04_pipelines.classifier.ipynb 21
from ..hf import HFTrainer
from transformers import pipeline
import numpy as np
import os.path


class HFClassifier(ClassifierBase):
    def __init__(
        self,
        model_id_or_path:str='google/bert_uncased_L-2_H-128_A-2',
        device=None,
        labels=[],
        **kwargs,
    ):
        """
        `HFClassifier` can be used to train and run text Hugging Face transformer text classifiers.
                Additional keyword arguments are fed directly to `from_pretrained`.


        **Args:**

        - *model_id_or_path*: The Hugging Face model_id or path to model folder (e.g, path previously trained and saved model).
        - *device*: 'cuda' or 'cpu'
        - *labels*: list of strings intended to map label indices to string labels

        """
        self.model_id_or_path = model_id_or_path
        self.device=device
        self.model = None
        self.tokenizer = None
        self.labels = labels
        if os.path.isdir(self.model_id_or_path):
            from transformers import AutoTokenizer
            from transformers import AutoModelForSequenceClassification 
            self.model = AutoModelForSequenceClassification.from_pretrained(self.model_id_or_path)
            self.tokenizer = AutoTokenizer.from_pretrained(self.model_id_or_path)

    def _tempsave(self):
        """
        Temporarily save model
        """
        import tempfile
        temp_dir = tempfile.TemporaryDirectory()
        if self.model:
            self.model.save_pretrained(temp_dir.name)
            self.tokenizer.save_pretrained(temp_dir.name)
        return temp_dir
    
    def train(self,
              X:List[str],
              y:Union[List[int], List[str]],
              **kwargs,
             ):
        """
        Trains the classifier on a list of texts (`X`) and a list of labels (`y`).
        Extra kwargs are treated as arguments to `transformers.TrainingArguments`.

        **Args:**

        - *X*: List of texts
        - *y*: List of integers representing labels
        - *num_epochs*: Number of epochs to train
        - *batch_size*: Batch size
        - *metric*: metric to use
        - *callbacks*:  A list of callbacks to customize the training loop.

        **Returns:**

        - None
        """

        temp_dir = self._tempsave()
              
        # convert to HFTrainer format
        data = [{'text': x, 'label':y[i]} for i, x in enumerate(X)]

        # create a trainer
        trainer = HFTrainer()

        # train
        training_path = self.model_id_or_path if not self.model else temp_dir.name
        self.model, self.tokenizer = trainer(training_path,  data, **kwargs)

        # cleanup
        temp_dir.cleanup()
  

    def save(self, save_path:str):
        """
        Save model to specified folder path, `save_path`.
        To reload the model, supply path in `model_id_or_path` argument when
        instantiating`FewShotClassifier`.

        """
        self.model.save_pretrained(save_path)
        self.tokenizer.save_pretrained(save_path)


    def _get_pipeline(self):
        """
        Create transformers pipeline using current model
        """
        if not self.model:
            raise ValueError('There is not trained model yet. Please invoke the HFClassifier.train(...)')
        save_path = None
        if not os.path.isdir(self.model_id_or_path):
            temp_dir = self._tempsave()
            modelpath = temp_dir.name
        else:
            modelpath = self.model_id_or_path
        clf = pipeline('text-classification', model=modelpath, device=self.device)
        return clf
        
    def _predict(self, X:list, return_proba=True, max_length=512, **kwargs):
        """
        Predicts labels using `transformers.pipeline`
        """
        from operator import itemgetter
        X = [X] if isinstance(X, str) else X
        clf = self._get_pipeline(**kwargs)
        preds = []
        labels = self.get_labels()
        for example in clf(X, top_k=None, truncation=True, max_length=max_length):
            # unlike deprecated return_all_scores, top_k re-sorts
            example = sorted(example, key=itemgetter('label'))
            pred = [d['score'] for d in example]
            pred = np.argmax(pred) if not return_proba else pred
            pred = labels[pred] if labels else pred
            preds.append(pred)
        if len(preds) == 1:
            preds = preds[0]
        return preds

    def predict(self, X:list, max_length=512, **kwargs):
        """
        Predict labels. 
        Extra kwargs fed to Hugging Face transformers text-classification pipeline.
        """
        return self._predict(X, max_length=max_length, return_proba=False)
        
    def predict_proba(self, X:list, max_length=512, **wargs):
        """
        Predict labels.
        Extra kwargs fed to Hugging Face transformers text-classification pipeline.
        """
        return self._predict(X, max_length=max_length, return_proba=True)

    
    def _get_explain_predictor(self, device=None):
        return self._get_pipeline(), None
               


# %% ../../nbs/04_pipelines.classifier.ipynb 27
from sklearn.datasets import fetch_20newsgroups
from .classifier import HFClassifier

# %% ../../nbs/04_pipelines.classifier.ipynb 32
DEFAULT_SETFIT_MODEL = "sentence-transformers/paraphrase-mpnet-base-v2"
SMALL_SETFIT_MODEL = "sentence-transformers/all-MiniLM-L6-v2"

class FewShotClassifier(ClassifierBase):
    def __init__(
        self,
        model_id_or_path:str=DEFAULT_SETFIT_MODEL,
        use_smaller:bool=False,
        **kwargs,
    ):
        """
        `FewShotClassifier` can be used to train and run text classifiers. Currently based on SetFit.
                Additional keyword arguments are fed directly to `from_pretrained`.


        **Args:**

        - *model_id_or_path*: The Hugging Face model_id or path to model folder (e.g, path previously trained and saved model).
        - *use_smaller*:  If True, will use a smaller but performant model.

        """
        self.model_id_or_path = model_id_or_path
        if use_smaller and model_id_or_path != DEFAULT_SETFIT_MODEL:
            warnings.warn(f'Over-writing supplied model ({model_id_or_path}) with {SMALL_MODEL} because use_smaller=True.')
        self.model_id_or_path = SMALL_SETFIT_MODEL if use_smaller else self.model_id_or_path
        self.model = SetFitModel.from_pretrained(self.model_id_or_path, **kwargs)
        self.predict = self.model.predict
        self.predict_proba = self.model.predict_proba
        self.trainer = None # set in `FewShotClassifier.train`
        self.labels = [] # set in `FewShotClassifier.train`
            
    def train(self,
              X:List[str],
              y:Union[List[int], List[str]],
              num_epochs:int=10,
              batch_size:int=32,
              metric='accuracy',
              callbacks=None,
              **kwargs,
             ):
        """
        Trains the classifier on a list of texts (`X`) and a list of labels (`y`).
        Additional keyword arguments are passed directly to `SetFit.TrainingArguments`

        **Args:**

        - *X*: List of texts
        - *y*: List of integers representing labels
        - *num_epochs*: Number of epochs to train
        - *batch_size*: Batch size
        - *metric*: metric to use
        - *callbacks*:  A list of callbacks to customize the training loop.

        **Returns:**

        - None
        """

        # convert to HF dataset
        train_dataset = self.arrays2dataset(X, y, text_key='text', label_key='label')

        args = TrainingArguments(
                batch_size=batch_size,
                num_epochs=num_epochs,
                **kwargs
        )      

        trainer = Trainer(
                    model=self.model,
                    args=args,
                    metric=metric,
                    callbacks=callbacks,
                    train_dataset=train_dataset,
                    column_mapping={"text": "text", "label": "label"}
        )
        trainer.train()
      
        self.trainer = trainer
  

    def predict(self, X, **kwargs):
        """
        predict labels
        """
        return self.predict(X, **kwargs)


    def predict_proba(self, X, **kwargs):
        """
        predict label probabilities
        """
        return self.predict_proba(X, **kwargs)


    def save(self, save_path:str):
        """
        Save model to specified folder path, `save_path`.
        To reload the model, supply path in `model_id_or_path` argument when
        instantiating`FewShotClassifier`.

        """
        self.model.save_pretrained(save_path)        

    def _get_explain_predictor(self, device=None):
        def f(x):
            return self.predict_proba(x)

        from transformers import AutoTokenizer
        tokenizer = AutoTokenizer.from_pretrained(self.model_id_or_path)
        return f, tokenizer
