from pathlib import Path
from collections import Counter
import gradio as gr
import pandas as pd
import numpy
from scipy.special import softmax

import torch
from transformers import pipeline, TextClassificationPipeline
from transformers import (
    AutoTokenizer,
    AutoConfig,
    AutoModelForSequenceClassification
)

# model_path = "EmotiScan/amazon-comments-bert" # To use your model, on line 55, you'll have to change i to i+1
# model_path = "nlptown/bert-base-multilingual-uncased-sentiment"
model_path = "Mofe/emotiscan_model_2"

tokenizer = AutoTokenizer.from_pretrained(model_path)
model = AutoModelForSequenceClassification.from_pretrained(model_path)


description = f"""# <h2 style="text-align: center;"> Opinion orbit Amazon Products Review</h2>
<p style="text-align: center;font-size:15px;"> Input a single review to predict sentiment or multiple from a text file.</p>
"""

sample_reviews = [
    "A comfortable pair of boots, but a little difficult to \
    remove without unlacing. So be prepared to unlace before \
    removing. Takes a few seconds to get them on and off, but \
        otherwise a very comfortable, sturdy product.",
    "Very nice boot, it's narrow whc I wasn't expecting, more \
     for a not so wide of a foot. My foot is not that wide, but \
    with socks on it will feel a little tight in the middle and back \
     of the heel, I would go a half or whole size up.",
    "If you like to continue having dry skin AND to smell like \
    the inside of a Shoppers Drug Mart then this product is for you!! \
     Sadly, I don't care for either of those things. I will continue \
     to hunt for a moisturizer/shower oil",
    "I really dislike this product.",
    "This is the best product ever!."
]

def upload_file(file):
    return file.name

def get_sentiments(text):
    encoded_input = tokenizer(text, return_tensors='pt')
    with torch.no_grad():
        logits = model(**encoded_input).logits

    logits_list = logits[0].tolist()
    logits_labels = [model.config.id2label[i]
                     for i in range(0, len(logits_list))]

    probs = softmax(logits_list)
    scores = {l:float(s) for (l,s) in zip(logits_labels, probs)}

    return scores

def get_multiple_sentiments(input_file):
    with open(input_file) as fn:
        reviews = fn.readlines()
        reviews = [review.strip() for review in reviews]

    classifier = pipeline('text-classification',
                          model=model_path)
    scores = classifier(reviews)
    scores = [pred['label'] for pred in scores]

    preds = dict(Counter(scores))
    labels = preds.keys()
    counts = preds.values()
    review_counts = pd.DataFrame(
    {
        "labels": labels,
        "counts": counts,
    }
    )
    total = len(scores)

    return gr.BarPlot(
            review_counts,
            x="labels",
            y="counts",
            title=f"# of Reviews: {total}",
            tooltip=["labels", "counts"],
            y_lim=[0, total + 10],
            min_width=500
        )

with gr.Blocks() as demo:
    with gr.Row():
        gr.Markdown(value=description)

    with gr.Row():
        with gr.Column():
            with gr.Tab("Single Input"):
                single = True
                input_text = gr.Textbox(label="Input Text",
                                placeholder="Input the product review...")
                predict_btn = gr.Button("Get Sentiment")
                with gr.Accordion("Here are some sample reviews!"):
                    examples = gr.Examples(examples=sample_reviews,
                                           inputs=[input_text])
            with gr.Tab("File Upload"):
                multiple = True
                file_output = gr.File()
                upload_button = gr.UploadButton("Upload a .txt file with each review on a line.",
                                                file_types=["text"])
                upload_button.upload(upload_file, upload_button, file_output)
                file_predict_btn = gr.Button("Get Sentiments")

        with gr.Column():
            output = gr.Label(label="Single Prediction")
            plots = gr.BarPlot(label="Plot Multiple Reviews")

    predict_btn.click(fn=get_sentiments,
                    inputs=input_text,
                    outputs=output,)
                    # api_name="product_review")
    file_predict_btn.click(fn=get_multiple_sentiments,
                    inputs=file_output,
                    outputs=plots)

demo.launch(share=True)