import torch
from transformers import T5ForConditionalGeneration, T5Tokenizer
from difflib import SequenceMatcher
import gradio as gr

# Sample Input
sample_input = '''BRILLAT SAVARIN - Physilogie du goût ou méditations de gastronomie tarnscendante : ouvrage thé0r1qu3, historique et à l’order du jour, déidé aux g45tr00m35 parisiens par un professeur. - 4e éd. - Paris, Just Tessire, 1834. - 2 vol., 384 p. ; 403 p. ; 22 cm. R35 XIX 260 y'''
# Load the model and tokenizer
model = T5ForConditionalGeneration.from_pretrained("./finetuned_t5_ocr_ppm_v2")
tokenizer = T5Tokenizer.from_pretrained("./finetuned_t5_ocr_ppm_v2")
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)

# Function to correct OCR text
def correct_text(text):
    input_text = "correct: " + text
    inputs = tokenizer(input_text, return_tensors="pt", truncation=True, padding=True).to(device)
    output_ids = model.generate(**inputs, max_length=128)
    return tokenizer.decode(output_ids[0], skip_special_tokens=True)

# HTML diff viewer
def highlight_diff(original, corrected):
    matcher = SequenceMatcher(None, original.split(), corrected.split())
    html = ""
    for tag, i1, i2, j1, j2 in matcher.get_opcodes():
        if tag == "equal":
            html += " " + " ".join(original.split()[i1:i2])
        elif tag == "replace":
            html += f' <span style="background-color:#fdd;">{" ".join(original.split()[i1:i2])}</span>'
            html += f' <span style="color:green;">→ {" ".join(corrected.split()[j1:j2])}</span>'
        elif tag == "delete":
            html += f' <span style="background-color:#faa;">{" ".join(original.split()[i1:i2])}</span>'
        elif tag == "insert":
            html += f' <span style="color:blue;">{" ".join(corrected.split()[j1:j2])}</span>'
    return html

# Gradio interface function
def process_text(ocr_input):
    corrected = correct_text(ocr_input)
    diff_html = highlight_diff(ocr_input, corrected)
    return corrected, diff_html

# Gradio UI
demo = gr.Interface(
    fn=process_text,
    inputs=gr.Textbox(lines=10, label="Paste Noisy OCR Text",value = sample_input),
    outputs=[
        gr.Textbox(label="Cleaned (Corrected) Text"),
        gr.HTML(label="Highlighted Corrections (Red = original, Green = correction)")
    ],
    title="OCR Post-Processing",
    description="Paste OCR text below and the model will output a cleaned version with differences shown in color-code."
)

if __name__ == "__main__":
    demo.launch()