import difflib
import gradio as gr
from transformers import pipeline
from tokenizers.pre_tokenizers import Whitespace
from tokenizers.normalizers import BertNormalizer

audio_input = gr.inputs.Audio(source='microphone', label='Read the passage', type="filepath")
text_input = gr.inputs.Textbox(label='Sample passage')
text_output = gr.outputs.Textbox(label='Output')
highlighted_text_output = gr.outputs.HighlightedText(color_map={"+": "green", "-": "pink"})

speech_to_text = pipeline('automatic-speech-recognition')

sm = difflib.SequenceMatcher(None)

splitter = Whitespace()
normalizer = BertNormalizer()

def preprocess(s):
    return [i[0] for i in splitter.pre_tokenize_str(normalizer.normalize_str(s))]

def diff_texts(text1, text2):
    d = difflib.Differ()
    return [
        (token[2:], token[0] if token[0] != " " else None)
        for token in d.compare(preprocess(text1), preprocess(text2))
    ]

def func(audio, text):
    # print(audio)
    # print(text)
    results = speech_to_text(audio)['text'].lower()
    text = text.lower()
    sm.set_seqs(preprocess(results), preprocess(text))

    r = f"""
    Original passage:
    {text}

    What we heard:
    {results}

    Ratio:
    {sm.ratio()}
    """

    d = diff_texts(results, text)

    return r, d

title = "Reading Practice Application"
description = """
This application is a POC for reading practice.
It compares some input text against an audio recording.
The intention is to help individuals with reading challenges identify how to improve their reading.
"""

gr.Interface(
    func, 
    inputs=[audio_input, text_input], 
    outputs=[text_output, highlighted_text_output],
    title=title,
    description=description,
    examples=[["sample.wav", "the quick brown fox jumped over the lazy dog"]]
).launch(debug=True)