import gradio as gr
import cv2
import numpy as np
import pytesseract
import re
import google.generativeai as genai
from rapidfuzz.distance import Levenshtein
import os

os.system('apt-get update && apt-get install -y tesseract-ocr')
# Configure Generative AI
OPENAI_API_KEY = os.getenv("API_KEY")
genai.configure(api_key=OPENAI_API_KEY)
model = genai.GenerativeModel("gemini-1.5-flash")

# Image processing functions
def threshold_image(img, threshold_value=None):
    if threshold_value is None:  # Adaptive thresholding
        thresholded_image = cv2.adaptiveThreshold(img, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C,
                                                  cv2.THRESH_BINARY, 11, 2)
    else:  # Manual thresholding
        _, thresholded_image = cv2.threshold(img, threshold_value, 255, cv2.THRESH_BINARY)
    return thresholded_image

def bm3d_denoising(img, sigma_psd=55):
    return cv2.fastNlMeansDenoising(img, None, sigma_psd)

def remove_noise(img, kernel_size=3):
    kernel = np.ones((kernel_size, kernel_size), np.float32) / (kernel_size**2)
    denoised = cv2.filter2D(img, -1, kernel)
    return cv2.medianBlur(denoised, 3)

def sharpen_image(img):
    kernel = np.array([[-1, -1, -1], [-1, 9, -1], [-1, -1, -1]])
    return cv2.filter2D(img, -1, kernel)

def remove_extra_spaces_and_lines(text):
    text = re.sub(r'\s+', ' ', text).strip()
    text = re.sub(r'\n\s*\n', '\n\n', text)
    return text

def calculate_accuracy(text1, text2):
    # matcher = difflib.SequenceMatcher(None, generated_text, transcribed_text)
    # return matcher.ratio()
    distance = Levenshtein.distance(text1, text2)
    max_length = max(len(text1), len(text2))
    accuracy = (1 - (distance / max_length))
    return accuracy

# Gradio app
def process_image(image, threshold_value=None, correct_transcription=None):
    img = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)

    # Process the image
    thresholded = threshold_image(img, threshold_value)
    bm3d_denoised_image = bm3d_denoising(thresholded)
    denoised = remove_noise(thresholded)
    sharpened_image = sharpen_image(bm3d_denoised_image)

    # OCR
    original_text = pytesseract.image_to_string(img)
    thresholded_text = pytesseract.image_to_string(thresholded)
    bm3d_denoised_text = pytesseract.image_to_string(bm3d_denoised_image)
    denoised_text = pytesseract.image_to_string(denoised)
    sharpened_text = pytesseract.image_to_string(sharpened_image)
    
    # Clean up text
    original_text = remove_extra_spaces_and_lines(original_text)
    thresholded_text = remove_extra_spaces_and_lines(thresholded_text)
    bm3d_denoised_text = remove_extra_spaces_and_lines(bm3d_denoised_text)
    denoised_text = remove_extra_spaces_and_lines(denoised_text)
    sharpened_text = remove_extra_spaces_and_lines(sharpened_text)

    # Generative AI model response
    user_prompt = user_prompt = f"""
    below are the output texts of OCR on multiple image processing techniques of a faded image with text written in English, can you use all the texts to predict the original text, provide only the text.
    Pre-Processing Image Text:
    {original_text}
    Sharpened Image Text:
    {sharpened_text}
    Thresholded Image Text:
    {thresholded_text}
    BM3D Denoised Image Text:
    {bm3d_denoised_text}
    Denoised Image Text:
    {denoised_text}
    """  
    response = model.generate_content(user_prompt)
    model_text = response.text

    if not correct_transcription:
        correct_transcription = model_text
    # Accuracy metrics
    if correct_transcription:
        original_accuracy = calculate_accuracy(original_text, correct_transcription)
        thresholded_accuracy = calculate_accuracy(thresholded_text, correct_transcription)
        bm3d_denoised_accuracy = calculate_accuracy(bm3d_denoised_text, correct_transcription)
        denoised_accuracy = calculate_accuracy(denoised_text, correct_transcription)
        sharpened_accuracy = calculate_accuracy(sharpened_text, correct_transcription)
        model_accuracy = calculate_accuracy(model_text, correct_transcription)
        accuracy_metrics = f"""
        Original Image Accuracy: {original_accuracy:.2%}
        Thresholded Image Accuracy: {thresholded_accuracy:.2%}
        BM3D Denoised Image Accuracy: {bm3d_denoised_accuracy:.2%}
        Denoised Image Accuracy: {denoised_accuracy:.2%}
        Sharpened Image Accuracy: {sharpened_accuracy:.2%}
        Model Response Accuracy: {model_accuracy:.2%}
        """
    else:
        accuracy_metrics = "No correct transcription provided."

    # Return results
    return (
        image, thresholded, bm3d_denoised_image, denoised, sharpened_image,
        original_text, thresholded_text, bm3d_denoised_text, denoised_text, sharpened_text,
        model_text, accuracy_metrics
    )

# Interface
with gr.Blocks() as demo:
    with gr.Row():
        gr.Markdown("## Faded text restoration")
    with gr.Row():
        gr.Markdown("""
        ### Legend
        - **Model Response**: Text generated by the Generative AI model.
        - **Accuracy Metrics**: Comparison of OCR results with the provided correct transcription if provided, otherwise with the model response.
        """)
    with gr.Row():
        with gr.Column():
            image_input = gr.Image(label="Upload Image", type="numpy")
            threshold_slider = gr.Slider(label="Threshold Value", minimum=0, maximum=255, step=1, value=242)
            adaptive_checkbox = gr.Checkbox(label="Use Adaptive Thresholding", value=False)
            transcription_input = gr.Textbox(label="Correct Transcription (Optional)")
            process_button = gr.Button("Process Image")

        with gr.Column():
            tabs = gr.Tabs()
            with tabs:
                with gr.TabItem("Original"):
                    original_image_display = gr.Image(label="Original Image")
                    original_text_display = gr.Textbox(label="Original Image Text", lines=5)
                with gr.TabItem("Thresholded"):
                    thresholded_image_display = gr.Image(label="Thresholded Image")
                    thresholded_text_display = gr.Textbox(label="Thresholded Image Text", lines=5)
                with gr.TabItem("BM3D Denoised"):
                    bm3d_denoised_image_display = gr.Image(label="BM3D Denoised Image")
                    bm3d_denoised_text_display = gr.Textbox(label="BM3D Denoised Image Text", lines=5)
                with gr.TabItem("Denoised"):
                    denoised_image_display = gr.Image(label="Denoised Image")
                    denoised_text_display = gr.Textbox(label="Denoised Image Text", lines=5)
                with gr.TabItem("Sharpened"):
                    sharpened_image_display = gr.Image(label="Sharpened Image")
                    sharpened_text_display = gr.Textbox(label="Sharpened Image Text", lines=5)
            accuracy_output = gr.Textbox(label="Accuracy Metrics")
            model_text_display = gr.Textbox(label="Model Response Text")

    # Link button to processing function
    def update_process(image, threshold_value, use_adaptive, correct_transcription):
        threshold_value = None if use_adaptive else threshold_value
        return process_image(image, threshold_value, correct_transcription)

    process_button.click(
        update_process,
        inputs=[image_input, threshold_slider, adaptive_checkbox, transcription_input],
        outputs=[
            original_image_display, thresholded_image_display,
            bm3d_denoised_image_display, denoised_image_display, 
            sharpened_image_display, original_text_display,
            thresholded_text_display, bm3d_denoised_text_display,
            denoised_text_display, sharpened_text_display,
            model_text_display, accuracy_output
            ],
    )

# Launch app
demo.launch()