import streamlit as st
from happytransformer import HappyTextToText, TTSettings
from annotated_text import annotated_text
import difflib
from bokeh.models.widgets import Button
from bokeh.models import CustomJS
from streamlit_bokeh_events import streamlit_bokeh_events

checkpoint = "team-writing-assistant/t5-base-c4jfleg"

def diff_strings(a, b):
  result = []
  diff = difflib.Differ().compare(a.split(), b.split())
  replacement = ""
  for line in diff:
    if line.startswith("  "):
      if len(replacement) == 0:
        result.append("  ")
        result.append(line[2:])
      else: 
        result.append("  ")
        result.append(("", replacement, "#ffd"))
        replacement = ""
        result.append(line[2:])
    elif line.startswith("- "):
      if len(replacement) == 0:
        replacement = line[2:]
      else:
        result.append("  ")
        result.append(("", replacement, "#fdd"))
        replacement = ""
    elif line.startswith("+ "):
      if len(replacement) == 0:
        result.append((line[2:], "", "#dfd"))
      else:
        result.append("  ")
        result.append((line[2:], replacement, "#ddf"))
        replacement = ""
  return result
  
@st.cache(suppress_st_warning=True, allow_output_mutation=True)
def get_happy_text(model_name):
  return HappyTextToText("T5", model_name)
         
happy_tt = get_happy_text(checkpoint)
args = TTSettings(num_beams=5, min_length=1)

st.title("Grammar Corrector Two")
st.markdown("Paste or type text. Submit. The machine will attempt to correct your text's grammar and highlight its corrections.")

st.subheader("Example text: ")
col1, col2, col3 = st.columns([1, 2, 1])
with col1:
  example_1 = st.button("Intrailly, the costumers was mad about why they will not but Fast Fashion again as they")
with col2:
  example_2 = st.button("Firstly,why i think this policy should be changed is because sometime the customer may buy wrong size,if our company’s no-exchange policy,customers have threatened no never buy from Fast Fashion again.")
with col3:
  example_3 = st.button("I try my best but still nervous. I hope I can get a good result.")

input_text = st.text_area('Paste or type text')
button = st.button('Submit')

def output(text):
  with st.spinner('Correcting'):
    text = "grammar: " + text
    result = happy_tt.generate_text(text, args=args)
    diff = diff_strings(text[9:], result.text)
    annotated_text(*diff)

    copy_button = Button(label="Copy the Result")
    copy_button.js_on_event("button_click", CustomJS(args=dict(result=result.text), code="""
      navigator.clipboard.writeText(result);
      """))
    streamlit_bokeh_events(
      copy_button,
      events="GET_TEXT",
      key="get_text",
      refresh_on_update=True,
      override_height=75,
      debounce_time=0)

if example_1:
  output("Intrailly, the costumers was mad about why they will not but Fast Fashion again as they")
elif example_2:
  output("Firstly,why i think this policy should be changed is because sometime the customer may buy wrong size,if our company’s no-exchange policy,customers have threatened no never buy from Fast Fashion again.")
elif example_3:
  output("I try my best but still nervous. I hope I can get a good result.")
elif input_text:
  output(input_text)