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

processor = AutoProcessor.from_pretrained("facebook/wav2vec2-base-960h")
model = AutoModelForCTC.from_pretrained("facebook/wav2vec2-base-960h")

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
).launch(inline=True, debug=True)