import numpy as np
import transformers
from onnxruntime.quantization import quantize_dynamic,QuantType
import transformers.convert_graph_to_onnx as onnx_convert
from pathlib import Path
import os

import torch
from transformers import AutoModelForSequenceClassification,AutoTokenizer

chkpt='distilbert-base-uncased-finetuned-sst-2-english'
model=AutoModelForSequenceClassification.from_pretrained(chkpt)
tokenizer=AutoTokenizer.from_pretrained(chkpt)
# tokenizer=AutoTokenizer.from_pretrained('sentiment_classifier/')

def classify_sentiment(texts,model,tokenizer):
    """
        user will pass texts separated by comma
    """
    try:
        texts=texts.split(',')
    except:
        pass

    input = tokenizer(texts, padding=True, truncation=True,
                      return_tensors="pt")
    logits = model(**input)['logits'].softmax(dim=1)
    logits = torch.argmax(logits, dim=1)
    output = ['Positive' if i == 1 else 'Negative' for i in logits]
    return output


def create_onnx_model_sentiment(_model, _tokenizer):
    """

    Args:
        _model: model checkpoint with AutoModelForSequenceClassification
        _tokenizer: model checkpoint with AutoTokenizer

    Returns:
        Creates a simple ONNX model & int8 Quantized Model in the directory "sent_clf_onnx/" if directory not present

    """
    if not os.path.exists('sent_clf_onnx_dir'):
        try:
            os.mkdir('sent_clf_onnx_dir')
        except:
            pass
        pipeline=transformers.pipeline("text-classification", model=_model, tokenizer=_tokenizer)

        onnx_convert.convert_pytorch(pipeline,
                                     opset=11,
                                     output=Path("sent_clf_onnx_dir/sentiment_classifier_onnx.onnx"),
                                     use_external_format=False
                                     )

        quantize_dynamic("sent_clf_onnx_dir/sentiment_classifier_onnx.onnx","sent_clf_onnx_dir/sentiment_classifier_onnx_quant.onnx",
                         weight_type=QuantType.QUInt8)
    else:
        pass


def classify_sentiment_onnx(texts, _session, _tokenizer):
    """

    Args:
        texts: input texts from user
        _session: pass ONNX runtime session
        _tokenizer: Relevant Tokenizer e.g. AutoTokenizer.from_pretrained("same checkpoint as the model")

    Returns:
        list of Positve and Negative texts

    """
    try:
        texts=texts.split(',')
    except:
        pass

    _inputs = _tokenizer(texts, padding=True, truncation=True,
                      return_tensors="np")

    input_feed={
        "input_ids":np.array(_inputs['input_ids']),
        "attention_mask":np.array((_inputs['attention_mask']))
    }

    output = _session.run(input_feed=input_feed, output_names=['output_0'])[0]

    output=np.argmax(output,axis=1)
    output = ['Positive' if i == 1 else 'Negative' for i in output]
    return output