import torch
from transformers import pipeline
from transformers import AutoModelForSeq2SeqLM
from transformers import AutoTokenizer
from textblob import TextBlob

# Load trained model
model = AutoModelForSeq2SeqLM.from_pretrained("output/reframer")
tokenizer = AutoTokenizer.from_pretrained("output/reframer")
reframer = pipeline('summarization', model=model, tokenizer=tokenizer)


def reframe(text, strategy):
    text_with_strategy = text +  "Strategy: ['" + strategy + "']"
    #Input Control
    #The input text cannot be too short to ensure it has substantial content to be reframed. It also cannot be too long to ensure the text has a focused idea.
    if len(text) < 15:
      return "[Error]: Input is too short. Please try again by inputing text with moderate length."
    if len(text) > 150:
      return "[Error]: Input is too long. Please try again by inputing text with moderate length."
    if TextBlob(text).sentiment.polarity > 0.2:
      return "[Error]: Input is too positive. Please try again by inputing text that is negative."
    from hatesonar import Sonar
    sonar = Sonar()
    if sonar.ping(text)['classes'][1]['confidence'] > 0.8:
      return "[Error]: Input is offensive. Please try again by not inputing offensive language."
    #if TextBlob(text).sentiment.polarity > 0:
    #  return "Please try again by inputing text that is negative."
    return reframer(text_with_strategy)[0]['summary_text']


import gradio as gr
with gr.Blocks() as demo:
    text = gr.Textbox(label="Original Text")
    radio = gr.Radio(
        ["thankfulness", "neutralizing", "optimism", "growth", "impermanence", "self_affirmation"], label="Strategy to use?"
    )
    output = gr.Textbox(label="Reframed Output")
    greet_btn = gr.Button("Reframe")
    greet_btn.click(fn=reframe, inputs=[text, radio], outputs=output)
demo.launch()